IREE项目中RDNA3架构下SDXL模型编译时linalg.transpose排列问题分析

问题背景

在IREE编译器项目中，当尝试在RDNA3架构上编译FP16精度的SDXL模型时，遇到了一个关于linalg.transpose操作的验证失败问题。具体表现为编译器在某个阶段生成了一个非法的转置操作，导致编译过程中断。

问题现象

错误信息显示："'linalg.transpose' op dim(result, 3) = 4 doesn't match dim(input, permutation[3]) = 2"，这表明转置操作的结果张量维度与输入张量维度不匹配。这个问题出现在PropagateReshapesByExpansionPass阶段之后。

技术分析

问题根源

通过深入分析，发现问题源于上游LLVM项目中的一个变更。具体来说，是LLVM的FoldWithProducerReshapeOpByExpansion模式发生了变化，该模式原本会插入linalg.generic操作，后来改为插入linalg.copy或linalg.transpose操作。

关键发现

在调试过程中，发现了一个可疑的张量展开操作：

%41 = "tensor.expand_shape"(%39) <{reassociation = [[0, 1], [2], [3], [4, 5]], static_output_shape = array<i64: 4, 16, 1, 2, 8, 16>}> : (tensor<4x16x1x2x8x16xf16>) -> tensor<4x16x1x2x8x16xf16>

这个操作特别之处在于输入和输出的秩都是6，实际上没有进行任何有效的形状扩展，这显然是一个无效操作。

上游变更影响

经过bisect定位，确认问题与上游LLVM的两个关键变更有关：

1. 6e59282235b2ba7b5bbae968cafb15bab9656cff：引入了expand shape bubbling模式
2. 70b95d16645dfe1e8d76bdf94e791d74ad36e780：修改了reshape操作的处理方式

这些变更导致IREE的FoldCollapseShapeIntoInterfaceTensorStore模式在后续处理中产生了无效的转置操作。

解决方案

上游开发者已经提交了修复补丁，主要解决了reshape作为消费者时的处理逻辑问题。该修复确保了在形状转换操作中正确处理张量维度和排列顺序。

技术启示

这个问题揭示了几个重要的技术要点：

1. 编译器pass的顺序敏感性：不同优化pass之间的交互可能导致意外的结果，特别是在处理张量形状转换时。

2. 张量操作验证的重要性：MLIR中的操作验证机制能够有效捕获维度不匹配问题，防止生成无效的IR。

3. 上游依赖的影响：IREE作为基于LLVM的项目，其行为会受到上游变更的显著影响，需要密切关注上游改动。

4. 形状操作的特殊情况处理：对于看似无操作但实际上可能改变语义的形状转换操作，需要特别小心处理。

总结

这个问题展示了深度学习编译器开发中的典型挑战——当多个优化pass和上游变更交互时，可能会产生微妙的错误。通过系统的bisect和深入分析，团队成功定位了问题根源并获得了上游修复。这也提醒开发者需要全面理解编译器各阶段的转换逻辑，特别是在处理张量形状操作时，要特别注意维度一致性和操作语义的正确性。