IREE项目中GPU卷积运算数值精度问题的分析与解决

问题背景

在IREE深度学习编译器项目中，开发人员发现当将反向卷积运算分解为正向卷积运算时，某些特定的卷积布局会导致显著的数值误差。这个问题在GPU后端（特别是ROCM/HIP）上表现尤为明显，而在CPU后端上计算结果则完全正确。

问题现象

以一个简单的1D卷积为例，输入张量形状为1x3x2，权重张量形状为2x2x2。当使用特定布局时，GPU计算结果与预期不符。例如：

• 输入数据：[[[0,1],[0,0],[0,0]]]
• 权重数据：[[[1,2],[3,4]],[[5,6],[7,8]]]
• 预期输出：[[5 6][0 0]]
• 实际GPU输出：[[3 4][0 0]]

技术分析

卷积运算的索引映射

问题根源在于卷积运算的索引映射关系。在原始IR中，定义了三个关键的affine映射：

1. 输入映射：将NHC格式的输入张量映射到计算空间
2. 权重映射：将CHF格式的权重张量映射到计算空间
3. 输出映射：将NHF格式的输出张量映射到计算空间

关键问题出现在权重张量的映射上，其reduction维度与输入张量的reduction维度顺序不一致，导致计算时维度对应关系错误。

im2col转换的问题

在GPU后端，卷积运算会被转换为im2col操作后接矩阵乘法。当前的im2col实现有一个隐含假设：权重张量的维度顺序必须与输入张量的对应维度顺序一致。当这个假设不成立时，就会导致数值计算错误。

解决方案探讨

开发团队提出了几种可能的解决方案：

1. 添加转置操作：在im2col转换前，显式添加转置操作调整权重张量的维度顺序
2. 扩展im2col元数据：为im2col操作添加额外的元数据，明确指定权重张量的维度布局
3. 调整im2col输出：扩展im2col输出的K维度，并添加后续转置操作

经过验证，第一种方案（添加转置操作）能够正确解决问题，但可能会引入性能开销。第二种方案更为优雅，但需要修改im2col操作的定义和实现。

实现细节

最终采用的解决方案是修改im2col的分解逻辑，使其能够正确处理权重张量与输入张量维度顺序不一致的情况。具体实现包括：

1. 分析输入和权重张量的维度对应关系
2. 在必要时插入隐式转置操作
3. 确保矩阵乘法的维度正确对齐

经验总结

这个问题的解决过程提供了几个有价值的经验：

1. 编译器中的卷积运算实现需要考虑各种可能的输入布局
2. 维度顺序的一致性检查在张量运算中至关重要
3. GPU后端与CPU后端的数值一致性验证是必要的质量保证手段
4. 设计通用运算（如im2col）时，需要明确其前提假设和约束条件

这个问题也提醒我们，在将高级运算分解为低级运算时，必须仔细处理维度语义的保持，特别是在涉及不同硬件后端时。