TVM项目中RemoveUnusedOutputs优化器的问题分析与修复

在深度学习编译器TVM项目中，最近发现了一个关于RemoveUnusedOutputs优化器的有趣问题。这个问题不仅揭示了优化器本身的实现缺陷，还暴露了TVM核心基础设施中关于NaN值比较的深层次问题。

问题背景

RemoveUnusedOutputs是TVM中的一个Relax优化器，其主要功能是识别并移除计算图中未被使用的输出。这个优化器会分析函数调用中哪些输出值被实际使用，然后相应地修改被调用函数的输出签名，从而减少不必要的计算。

在测试过程中发现，当遇到特定模式的代码时，这个优化器会产生两个异常行为：

1. 错误地将有效输出替换为NaN值
2. 导致assert_structural_equal对相同模块的结构比较失败

问题现象分析

优化器错误替换输出

在原始测试案例中，模块包含一个返回三元组的主函数，这个三元组由三个张量组成。经过RemoveUnusedOutputs处理后，优化器错误地将所有三个张量输出都替换为了NaN值，完全移除了原本有效的计算逻辑。

结构比较失败

更令人困惑的是，当尝试使用assert_structural_equal比较优化前后的模块时，即使比较的是完全相同的模块，也会抛出结构不等的异常。深入分析发现，这个问题源于TVM核心对NaN值的特殊处理。

根本原因

NaN比较问题

TVM的结构比较(StructuralEqual)对于浮点数的处理采用了相对误差比较方法，即检查abs(lhs-rhs) < 1e9。然而，NaN值与任何值（包括它自己）的比较都会返回false，这导致了结构比较的失败。这是一个基础架构层面的问题，影响了所有涉及NaN值的比较操作。

优化器逻辑缺陷

RemoveUnusedOutputs优化器的实现存在一个关键假设：它认为所有对函数输出的使用都会通过TupleGetItem操作显式地访问特定元素。当遇到直接返回整个元组的情况时，优化器无法正确识别这些使用，错误地认为所有输出都未被使用，从而将它们替换为NaN占位符。

解决方案

NaN比较修复

针对NaN比较问题，解决方案是在StructuralEqual和StructuralHash中增加对NaN值的特殊处理。具体来说，当检测到比较双方都是NaN值时，应该认为它们在结构上是相等的。

优化器使用分析修复

对于RemoveUnusedOutputs优化器，需要改进其使用分析逻辑，使其能够识别元组的完整使用情况，而不仅仅是TupleGetItem操作。这包括：

1. 跟踪元组作为整体的使用情况
2. 正确处理函数返回和传递场景
3. 确保不错误地将有效使用识别为未使用

技术影响

这个问题的修复对TVM项目有重要意义：

1. 提高了优化器的正确性，确保不会错误移除有效计算
2. 增强了基础设施的鲁棒性，特别是对特殊浮点值的处理
3. 为后续开发提供了更可靠的NaN值处理参考

总结

TVM作为深度学习编译器，其优化器的正确性直接影响生成代码的质量。这次发现的问题展示了编译器开发中常见的挑战：优化器需要在复杂的程序变换中保持语义不变性。通过深入分析NaN比较和元组使用分析这两个看似独立的问题，我们不仅修复了具体bug，还提升了TVM核心架构的健壮性。这对保证TVM在各种边缘情况下的正确行为具有重要意义。