IREE项目中ROCM后端处理带步长的矩阵乘法优化问题分析

问题背景

在IREE编译器项目中，当使用ROCM后端针对AMD GPU进行代码生成时，发现一个性能优化问题：带有步长访问的矩阵乘法运算无法正确降级到硬件原语指令（mfma）。这种情况会显著影响深度学习模型中卷积等运算在AMD GPU上的执行效率。

问题现象

原始代码展示了一个典型的卷积运算实现，其中包含一个带有步长访问的矩阵乘法运算。关键特征包括：

1. 输入张量形状为2x118x182x448（bf16类型）
2. 权重张量形状为896x1x1x448（bf16类型）
3. 输出张量形状为2x59x91x896（f32类型）
4. 左侧输入张量在高度和宽度维度上采用步长2的访问模式

技术分析

原始IR结构分析

原始IR使用linalg.generic操作实现矩阵乘法，其索引映射关系显示：

• 输入张量访问模式为(d0, d1*2, d2*2, d4)
• 权重张量访问模式为(d3, d4)
• 输出张量访问模式为(d0, d1, d2, d3)

这种带有步长的访问模式阻碍了编译器将其识别为标准的矩阵乘法运算，从而无法利用AMD GPU的矩阵融合乘加（MFMA）指令。

优化方案实现

通过实现一个专门的转换通道（ConvertStridedContractionToContractionPass），可以将带步长的收缩运算转换为标准收缩运算。该优化包含两个关键步骤：

1. 张量形状重构：首先将权重张量从896x1x1x448重构为896x448，消除不必要的维度

2. 输入切片提取：使用tensor.extract_slice操作从原始输入中提取出实际参与计算的子张量，消除步长访问模式

优化后IR分析

优化后的IR结构变为：

• 提取出2x59x91x448的有效输入子张量
• 使用标准的5维并行+归约迭代器实现矩阵乘法
• 索引映射简化为直接的维度访问模式

这种转换使得编译器能够识别出标准的矩阵乘法模式，从而可以正确降级到MFMA指令。

性能影响

这种优化对深度学习模型性能有显著影响：

1. 使能硬件加速指令：MFMA指令可以提供更高的计算吞吐量
2. 减少内存访问：通过消除步长访问模式，提高缓存利用率
3. 提高寄存器使用效率：标准矩阵乘法模式允许更好的寄存器分配

实现细节

优化转换的核心在于识别和重构张量访问模式：

1. 识别带有步长的收缩运算模式
2. 计算实际参与计算的子张量范围
3. 插入适当的张量提取操作
4. 重构索引映射关系

该优化已通过测试验证，能够正确生成MFMA指令，为AMD GPU上的深度学习运算提供显著的性能提升。

结论

IREE编译器通过这种针对性的优化转换，解决了ROCM后端在处理带步长矩阵乘法时的指令降级问题。这种优化不仅适用于当前案例，也为处理类似模式的张量运算提供了通用解决方案，对提升深度学习模型在AMD GPU上的执行效率具有重要意义。