TVM项目中RemoveUnusedOutputs优化器的异常行为分析

在深度学习编译器TVM的使用过程中，我们发现了一个关于RemoveUnusedOutputs优化器的异常行为，以及一个与结构相等性检查相关的bug。本文将详细分析这两个问题的技术细节及其解决方案。

RemoveUnusedOutputs优化器的问题

RemoveUnusedOutputs是TVM中的一个Relax优化器，其主要功能是移除未被使用的函数输出。但在实际使用中，该优化器在某些情况下会产生不符合预期的结果。

以一个具体案例为例，原始IR模块包含多个TIR原语函数和一个返回三个张量的Relax函数。经过RemoveUnusedOutputs优化后，原本应该保留的张量输出被替换成了NaN值，这显然不是预期的优化行为。

深入分析发现，问题的根源在于优化器在确定哪些输出被使用时，仅收集了通过TupleGetItem显式访问元组元素的情况。如果整个元组被直接使用（如作为函数返回值），优化器会错误地认为所有元素都未被使用，从而用NaN占位符替换它们。

结构相等性检查的NaN问题

另一个相关问题是TVM的结构相等性检查（assert_structural_equal）在处理NaN值时的不一致性。测试发现，即使两个完全相同的IR模块，如果包含NaN值，结构相等性检查也会错误地判定它们不相等。

这是由于TVM内部对浮点数的比较采用了相对误差方法（abs(lhs-rhs) < 1e9），而NaN值与任何值（包括它自己）的比较结果都是false。这种实现方式导致包含NaN的IR模块无法正确通过结构相等性验证。

解决方案

针对这两个问题，TVM社区已经提出了相应的修复方案：

1. 对于RemoveUnusedOutputs优化器，修复方案改进了使用情况收集逻辑，确保能够正确识别所有元组元素的使用情况，避免错误地用NaN替换实际使用的输出。

2. 对于结构相等性检查，修复方案增加了对NaN值的特殊处理，确保两个NaN值能够被正确识别为相等。

技术启示

这个案例给我们带来几个重要的技术启示：

1. 编译器优化器的实现需要考虑所有可能的使用场景，特别是边界情况。

2. 浮点数的特殊值（如NaN）在编译器内部处理时需要格外小心，常规的比较逻辑可能不适用。

3. 测试用例应当覆盖各种边界情况，包括特殊值和非常规使用模式。

TVM作为深度学习编译器，其正确性对下游应用至关重要。这类问题的发现和修复有助于提高编译器的稳定性和可靠性，为深度学习模型的部署提供更坚实的基础设施支持。