IREE项目中RDNA3架构下linalg.transpose排列问题的分析与解决

问题背景

在IREE编译器项目中，当使用RDNA3架构运行FP16精度的SDXL模型时，出现了卷积调度编译失败的问题。核心错误信息显示，在编译过程中生成了一个非法的linalg.transpose操作，具体表现为结果张量的维度与输入张量的排列不匹配。

问题现象

错误信息明确指出：

configured_compiled_unet_main$async_dispatch_1146.mlir:48:15: error: 'linalg.transpose' op dim(result, 3) = 4 doesn't match dim(input, permutation[3]) = 2

这个错误出现在PropagateReshapesByExpansionPass阶段，但不确定问题是在此阶段产生的还是仅仅在此阶段首次被发现。

技术分析

问题根源

通过深入分析，发现问题源于LLVM上游的一个变更。具体来说，上游的FoldWithProducerReshapeOpByExpansion模式被修改，从使用linalg.generic操作改为插入linalg.copy或linalg.transpose操作。

关键差异在于：

• 旧版本（正常工作的版本）使用linalg.generic操作
• 新版本（有问题的版本）使用linalg.copy或linalg.transpose操作

具体表现

在IR层面，可以看到以下变化：

正常版本：

%29 = linalg.generic {
    indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>,
                    affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>],
    iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]}
    ins(%transposed_8 : tensor<4x16x1x2x8x16xf16>)
    outs(%expanded_10 : tensor<4x16x1x2x8x16xf16>) {
    ^bb0(%in: f16, %out: f16):
        linalg.yield %in : f16
} -> tensor<4x16x1x2x8x16xf16>

问题版本：

%29 = linalg.copy 
    ins(%transposed_8 : tensor<4x16x1x2x8x16xf16>)
    outs(%expanded_10 : tensor<4x16x1x2x8x16xf16>)
    -> tensor<4x16x1x2x8x16xf16>

问题定位

通过二分法定位，确认问题源于LLVM上游的特定提交。特别是以下两个变更可能是问题根源：

1. 6e59282235b2ba7b5bbae968cafb15bab9656cff - 上游了一个expand shape冒泡模式
2. 70b95d16645dfe1e8d76bdf94e791d74ad36e780 - 也在问题出现的时间范围内

解决方案

上游开发者确认了问题并提供了修复方案。修复的核心思路是正确处理reshape操作作为消费者的情况。具体修复包括：

1. 确保在折叠操作时正确处理张量形状
2. 修复reshape操作作为消费者时的处理逻辑
3. 添加相应的测试用例

修复已经提交到上游并被合并，现在主分支已经可以正常工作。

技术启示

这个问题展示了编译器优化过程中几个重要的技术点：

1. IR变换的安全性：任何IR变换都必须保证结果的合法性，特别是在改变操作类型时（如从generic到copy/transpose）

2. 形状处理的重要性：在处理张量形状变换时，必须确保维度匹配，特别是在有复杂reshape和transpose组合的情况下

3. 上游依赖的影响：编译器项目依赖上游组件时，上游变更可能引入难以预料的问题，需要建立有效的监控和回滚机制

这个问题也提醒我们，在编译器优化过程中，形状推导和验证是确保正确性的关键环节，任何优化都必须通过严格的验证检查。