IREE项目中ROCM后端Strided Matmul优化问题分析

问题背景

在IREE项目的ROCM后端实现中，开发者发现了一个关于矩阵乘法（matmul）性能优化的问题。当矩阵乘法操作中的左侧矩阵（LHS）以步长2进行访问时，该操作无法正确降级（lower）到AMD GPU的矩阵融合乘法累加（MFMA）指令集。

技术细节

该问题涉及到一个特定的计算模式：一个2D卷积操作的前向传播实现，其核心计算部分是一个带有步长的矩阵乘法。原始实现中，左侧输入张量（形状为2x118x182x448）以步长2在高度和宽度维度上进行访问，右侧权重张量形状为896x1x1x448。

问题的关键在于，这种带有非单位步长的矩阵乘法操作无法被IREE编译器正确识别并优化为针对AMD GPU（特别是gfx942架构）的MFMA指令。MFMA指令是AMD GPU中专门用于加速矩阵乘法运算的硬件特性，能够显著提升深度学习工作负载的性能。

解决方案

开发者提出了一个优化方案，通过引入一个新的编译器传递（pass）来转换这种带有步长的矩阵乘法操作。该方案的核心思想是：

1. 首先对右侧权重张量进行形状折叠（collapse_shape），将其从896x1x1x448转换为896x448的2D张量
2. 然后对左侧输入张量使用extract_slice操作，显式地提取出步长为2的子张量（2x59x91x448）
3. 最后将原始的带有步长访问的矩阵乘法转换为标准的矩阵乘法操作

这种转换使得编译器能够更容易识别出矩阵乘法模式，从而成功降级到MFMA指令。转换后的IR显示了更清晰的计算模式，其中显式地表示了数据的提取和标准的矩阵乘法操作。

技术意义

这个优化对于IREE项目在AMD GPU上的性能至关重要，因为：

1. 许多深度学习模型中的卷积操作都可以表示为带有步长的矩阵乘法
2. MFMA指令能够提供比通用计算高得多的计算吞吐量
3. 这种转换保持了计算的语义等价性，同时为后端优化提供了更好的机会

该问题的解决展示了编译器优化中一个常见但重要的模式：将复杂的、非常规的内存访问模式转换为标准的、易于优化的计算模式。这种转换在保持计算正确性的同时，为硬件特定的优化打开了大门。

未来展望

虽然当前解决方案已经能够解决问题，但开发者提到还需要进一步收集反馈并完善实现。这可能包括：

1. 更通用的步长处理模式识别
2. 对其他数据类型（如FP16、INT8等）的支持
3. 对不同AMD GPU架构的适配
4. 性能基准测试和验证

这种优化不仅限于当前的特定案例，还可以扩展到其他类似的非单位步长内存访问模式，为IREE在AMD GPU上的性能优化开辟了新的可能性。