IREE项目中卷积操作转换为IGEMM时的维度处理问题分析

问题背景

在IREE编译器项目中，当将卷积操作转换为IGEMM（整数矩阵乘法）形式时，发现了一个关于维度处理的bug。该问题主要出现在处理具有特殊维度特征的卷积操作时，特别是在将命名卷积操作（named convolution op）转换为通用形式（generic convolution op）后。

问题现象

当卷积操作具有以下特征时会出现问题：

1. 输入张量格式为NHWC（批量大小×高度×宽度×通道数）
2. 卷积核格式为FHWC（输出通道×高度×宽度×输入通道）
3. 卷积核的某些空间维度（如宽度维度）为1

在这种情况下，如果直接将命名卷积操作转换为通用形式，编译器会在ConvolutionToIGEMM转换阶段失败。具体表现为im2col操作错误地转置了空间维度，导致后续的GEMM通用操作验证失败。

技术细节分析

维度处理逻辑

问题的核心在于维度处理逻辑。当卷积核的某些空间维度为1时，系统会将这些维度视为"单位维度"（unit dim）并进行折叠处理。在命名卷积操作中，只有当所有滤波器维度都是单位维度时才会进行通用化转换。但在手动通用化后，系统会折叠单个单位维度，导致维度推断出现问题。

im2col操作的特殊性

im2col操作在处理卷积时会：

1. 将批量维度移到最外层
2. 将输出窗口维度重新排列
3. 将滤波器窗口和输入通道维度合并为内部维度

在问题案例中，由于卷积核的第二维度（宽度）为1，系统将其视为1D卷积而非2D卷积处理，导致输出/输入宽度被当作另一个批量维度处理。

解决方案

经过深入分析，确定了以下解决方案：

1. 调整结果映射：修改im2col操作的结果映射逻辑，使其正确考虑批量维度的重新排列。具体来说，需要确保结果映射能够反映批量维度被移到最外层的变化。

2. 维度顺序一致性：确保在维度折叠后，输出索引映射与输入索引映射保持一致。在问题案例中，需要将第二个批量维度与m维度交换，使输出索引映射变为(d0, d2, d1, d3)。

3. 分块大小调整：由于维度顺序变化，需要相应调整工作组分块大小，确保分块大小能够整除输出张量的各维度大小。

实现影响

该问题的修复不仅解决了当前案例中的编译失败问题，还为以下方面奠定了基础：

1. 为分组卷积（grouped convolution）的支持铺平了道路
2. 增强了系统处理非标准卷积维度配置的能力
3. 提高了编译器对特殊维度卷积操作的鲁棒性

经验总结

通过这个问题的解决，我们获得了以下经验：

1. 维度折叠需要谨慎处理，特别是在部分维度为单位维度的情况下
2. im2col操作的结果映射必须准确反映维度重排的变化
3. 工作组分块策略需要与实际的维度顺序和大小相匹配

这个问题展示了IREE编译器在处理复杂张量操作时的挑战，也体现了编译器开发中维度分析和映射处理的重要性。