IREE项目中Deepseek模型编译失败的根源分析与解决方案

问题背景

在IREE编译器项目中，用户报告了一个关于Deepseek模型编译失败的问题。当尝试使用IREE编译器将Deepseek模型的MLIR中间表示转换为目标设备可执行文件时，遇到了验证错误。这个问题特别出现在使用HIP后端针对AMD GPU(gfx942)进行编译的场景中。

问题现象

编译过程中出现的核心错误信息表明，在Linalg批处理矩阵乘法(batch_matmul)操作中，操作数的秩(rank)与索引映射(indexing_map)的结果秩不匹配。具体错误提示为："'linalg.batch_matmul' op expected operand rank (4) to match the result rank of indexing_map #1 (3)"。

技术分析

最小复现案例

通过分析，可以将其简化为以下最小复现代码：

func.func @main(%arg0: memref<512x1x192xf32>,
               %arg1: memref<4x128x192x64xf32>,
               %arg2: memref<512x1x64xf32>) {
    linalg.batch_matmul indexing_maps = [
        affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>,
        affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>,
        affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
    ]
    ins(%arg0, %arg1: memref<512x1x192xf32>, memref<4x128x192x64xf32>) 
    outs(%arg2: memref<512x1x64xf32>)
    return
}

在这个例子中，第二个索引映射affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>产生了一个秩为3的结果，而对应的输入操作数memref<4x128x192x64xf32>的秩却是4，这违反了Linalg操作的验证规则。

根本原因

问题根源在于IREE编译器中的形状下沉传递(SinkReshapesPass)产生了非法的操作。具体来说，当编译器尝试通过折叠形状操作(如tensor.collapse_shape)来优化元素级操作融合时，新生成的折叠操作不能保证与原始折叠形状操作具有相同的类型。

在更深的层次上，这是MLIR上游Linalg方言中元素级操作融合实现的一个缺陷。当处理包含形状折叠的操作时，现有的融合逻辑没有正确处理维度折叠后的类型一致性。

解决方案

临时解决方案

IREE团队迅速响应，提供了一个临时修复方案，通过修改shouldBubbleCollapseShapeOp函数使其返回false，从而禁用可能导致问题的形状折叠操作冒泡优化。这个方案可以立即解决问题，但并非长期最优解。

长期解决方案

真正的解决方案需要在上游LLVM项目中修复。核心问题在于元素级操作融合过程中对形状折叠操作的处理逻辑。修复需要确保：

1. 新生成的折叠操作必须与原始折叠形状操作保持类型一致
2. 在融合过程中正确处理维度折叠后的类型转换
3. 完善相关的验证逻辑，防止生成非法IR

技术影响与启示

这个案例揭示了深度学习编译器开发中的几个重要方面：

1. 中间表示验证的重要性：MLIR的强大类型系统能够及早发现这类维度不匹配问题，防止更隐蔽的错误

2. 优化传递的顺序敏感性：形状操作处理需要在正确的阶段进行，过早或不当的优化可能导致语义错误

3. 上游依赖的管理：作为基于MLIR构建的编译器，IREE需要平衡快速修复与上游贡献的关系

4. 模型兼容性挑战：不同深度学习框架导出的模型可能包含特殊的操作模式，需要编译器具备鲁棒的处理能力

结论

Deepseek模型编译失败问题展示了深度学习编译器开发中可能遇到的复杂场景。通过深入分析，IREE团队不仅提供了快速解决方案，还推动了上游MLIR/LLVM项目的改进。这类问题的解决过程体现了现代编译器架构的模块化设计价值，使得问题能够被隔离和针对性修复，而不影响整个编译管道的其他部分。