TVM项目中VMBuiltinLower对ShapeExpr类型检查过严的问题分析

问题背景

在TVM深度学习编译器项目中，Relax中间表示层的VMBuiltinLower转换阶段出现了一个类型检查过于严格的问题。该问题导致在某些情况下，当代码中使用了非直接ShapeExpr类型的形状表达式时，编译器会报错终止。

问题现象

开发者在使用TVM编译一个包含reshape操作的Relax模块时遇到了编译错误。错误信息显示VMBuiltinLower转换阶段期望绑定的值必须是一个ShapeExpr类型，但实际上传入的是一个通过tensor_to_shape转换得到的形状值。

技术分析

错误发生的上下文

在示例代码中，开发者首先通过call_tir调用两个TIR函数(add和multiply)计算出一个目标形状张量，然后使用R.tensor_to_shape将这个张量转换为形状表达式，最后将这个形状表达式用于reshape操作。

问题本质

VMBuiltinLower转换阶段的实现中存在一个过于严格的类型检查：

1. 它要求所有形状相关的参数必须是直接的ShapeExpr类型
2. 但实际上，Relax的类型系统中，只要表达式具有ShapeStructInfo类型信息，就应该被认为是合法的形状表达式
3. 这种严格的检查阻止了通过函数调用等方式间接获得的形状表达式的使用

影响范围

这个问题会影响所有使用以下模式的代码：

• 动态计算形状并通过tensor_to_shape转换
• 通过函数调用返回形状表达式
• 任何非直接R.shape(...)方式创建的形状表达式

解决方案

修复方案的核心思想是放宽类型检查条件：

1. 不再要求必须是直接的ShapeExpr类型
2. 改为检查表达式是否具有ShapeStructInfo类型信息
3. 这样既保证了类型安全，又支持了更灵活的形状表达式使用方式

技术意义

这个修复体现了TVM类型系统的设计原则：

1. 应该基于结构类型(structural typing)而非名义类型(nominal typing)进行检查
2. 表达式的合法性应该由其类型信息决定，而非创建方式
3. 支持更灵活的表达式组合方式，增强语言表达能力

对开发者的建议

在使用TVM的Relax前端时，开发者可以：

1. 放心使用各种方式构造形状表达式
2. 当需要进行形状计算时，可以灵活组合各种操作
3. 遇到类似类型错误时，检查是否因类型检查过于严格导致

该问题的修复使得TVM的形状处理更加灵活和强大，为动态形状计算提供了更好的支持。