IREE项目中GPU编译失败的卷积优化问题分析

背景介绍

在深度学习框架IREE的GPU后端编译过程中，开发人员遇到了一个关于分组卷积反向传播操作的编译问题。该问题表现为两种逻辑上等价的IR（中间表示）实现方式，其中一种能够成功编译，而另一种则因共享内存超出限制而失败。

问题现象

开发人员提供了两种实现分组卷积反向传播的IR代码：

1. 成功编译的版本：在卷积操作前后分别进行维度扩展和折叠操作，卷积操作内部包含滤波器空间维度的翻转和输入数据的填充。

2. 编译失败的版本：在函数边界处进行维度扩展和折叠，将滤波器翻转和输入填充操作放在扩展之后执行。这个版本在编译时报错，提示共享内存使用超出限制（592944字节超过了65536字节的限制）。

技术分析

失败原因

编译失败的根本原因在于IR优化过程中未能有效地融合reshape操作。具体表现为：

1. 在失败的版本中，存在一个generic(collapse_shape(generic(...)))的操作链
2. 现有的FoldWithProducerReshapeOpByExpansion模式无法匹配这种情况，因为isFusableWithReshapeByDimExpansion返回false
3. 虽然FoldReshapeWithGenericOpByCollapsing模式理论上可以处理这种情况，但它并未在BubbleExpandShapes过程中被调用

解决方案探索

通过实验性修改，开发人员发现：

1. 如果强制应用FoldReshapeWithGenericOpByCollapsing模式，可以成功消除tensor.collapse_shape操作，使IR能够编译通过
2. 但直接添加这个模式可能会引入其他不需要的模式，需要谨慎处理

更深层次的解决方向

项目团队指出：

1. 这个问题实际上是由于分组卷积尚未启用im2col优化路径导致的
2. 一旦为分组卷积启用im2col转换路径，系统会自动应用正确的reshape融合模式
3. 当前dispatch创建过程中的reshape下沉操作受到控制函数的限制，未能处理这种情况

技术影响

这个问题揭示了IREE在以下方面的优化空间：

1. 分组卷积的特殊处理路径需要完善
2. reshape操作融合的触发条件需要更精细的控制
3. 共享内存使用估算和优化需要更智能的策略

结论与展望

虽然目前可以通过临时修改绕过这个问题，但更根本的解决方案是完善分组卷积的im2col转换路径。IREE团队计划在未来版本中：

1. 为分组卷积启用im2col优化
2. 优化reshape操作融合的触发逻辑
3. 改进共享内存使用估算机制

这将从根本上解决此类编译问题，同时提高分组卷积在GPU上的执行效率。对于开发者而言，目前建议采用第一种IR实现方式作为临时解决方案。