IREE项目中动态维度导致的scatter_和one_hot操作数值问题分析

问题概述

在IREE项目中发现了一个关于动态维度的数值计算问题，主要影响scatter_和one_hot两种操作。当输入张量中部分维度被标记为动态而其他维度保持静态时，会导致计算结果不正确。这个问题在GPU后端（特别是VectorDistribution阶段）表现得尤为明显，而在LLVM CPU后端则能正常工作。

问题表现

该问题表现为以下两种操作在特定条件下的错误：

1. scatter_操作：当执行稀疏更新操作时，如果输入张量的某些维度是动态的而其他维度是静态的，结果会出现数值错误。

2. one_hot操作：在创建one-hot编码时，同样在部分动态维度的条件下，生成的编码结果不正确。

问题复现条件

问题仅在以下特定条件下出现：

• 部分维度标记为动态（如num_tokens维度）
• 部分维度保持静态（如num_top_experts维度）
• 使用GPU后端（HIP目标）

当所有维度都是动态或都是静态时，计算结果正确。LLVM CPU后端在所有情况下都能产生正确结果。

根本原因分析

经过深入调查，发现问题出在IREE的优化阶段，具体是FoldCollapseShapeIntoInterfaceTensorStoreFullSlice这个优化转换存在缺陷。该转换在处理张量形状折叠时，错误地将：

%collapsed = tensor.collapse_shape %src [[0, 1]] : tensor<?x2xi32> into tensor<?xi32>
store %collapsed, %dest, offsets = [0, 0], sizes = [%dim, 1], strides = [1, 1]

转换成了：

store %src, %dest, offsets = [0, 0], sizes = [%dim, 2], strides = [1, 1]

这种转换在静态维度情况下是正确的，但在混合动态/静态维度情况下会导致尺寸计算错误，从而产生不正确的存储操作。

技术影响

这个问题影响了：

1. 张量形状折叠优化的正确性
2. 动态维度张量的存储操作
3. 依赖这些基础操作的更高级算子（如scatter_和one_hot）

解决方案

该问题已在提交中被修复。修复方案包括：

1. 增强FoldCollapseShapeIntoInterfaceTensorStoreFullSlice转换的维度检查
2. 确保在混合动态/静态维度情况下正确处理形状折叠
3. 添加相应的测试用例防止回归

开发者建议

对于遇到类似问题的开发者，建议：

1. 检查是否在混合动态/静态维度情况下使用高级操作
2. 验证计算结果与预期是否一致
3. 考虑使用最新版本的IREE，其中已包含此修复
4. 对于关键计算路径，可以在CPU和GPU后端都进行验证

这个问题展示了动态张量处理中的一些微妙之处，特别是在优化转换过程中需要特别注意维度信息的正确传播。