IREE项目中LinalgExt::SortOp在LLVMGPU管道中的内存分配问题分析

问题背景

在IREE编译器项目中，当尝试将LinalgExt::SortOp操作通过LLVMGPUTileAndFuse管道进行降级时，遇到了一个编译失败的问题。这个问题在LLVMGPUDistribute管道中不会出现，值得深入分析其根本原因。

问题现象

在编译一个ONNX非极大值抑制(NonMaxSuppression)操作时，编译器报告了一个关于未解析外部函数'malloc'的错误。错误信息显示在最终生成的bitcode中发现了对malloc函数的调用，这在GPU内核中是不应该出现的。

技术分析

内存分配差异

通过对比两种不同管道（LLVMGPUTileAndFuse和LLVMGPUDistribute）在IREEComprehensiveBufferizePass后的中间表示，发现了关键差异：

1. LLVMGPUTileAndFuse管道：

   • 生成了全局内存(fat_raw_buffer)的分配
   • 使用memref.alloc()创建了memref<6xf32, #amdgpu.address_space<fat_raw_buffer>>

2. LLVMGPUDistribute管道：

   • 正确地生成了工作组内存(workgroup)的分配
   • 使用memref.alloc()创建了memref<6xf32, #gpu.address_space<workgroup>>
   • 包含了必要的gpu.barrier同步操作

问题本质

在GPU编程中，内核内部动态分配全局内存是一个严重的设计问题。正确的做法应该是：

1. 对于工作组共享的数据，使用工作组内存(workgroup memory)
2. 对于线程私有的数据，使用私有内存(private memory)
3. 全局内存分配应该在内核启动前完成，通过参数传递

LLVMGPUTileAndFuse管道错误地生成了全局内存分配，这会导致：

1. 性能问题：全局内存访问延迟高
2. 功能问题：可能无法在某些GPU架构上正常工作
3. 编译问题：需要运行时支持动态内存分配

解决方案方向

要解决这个问题，需要确保：

1. 缓冲区配置正确识别内存使用场景
2. 对于工作组内共享的排序临时缓冲区，使用工作组内存
3. 添加必要的内存同步操作
4. 完全避免在内核中使用动态内存分配

技术影响

这个问题不仅影响特定ONNX操作的正确性，还反映了IREE编译器在内存分配策略上的一个重要边界情况。正确处理这类问题对于保证编译器在各种深度学习算子上的可靠性至关重要。

结论

通过深入分析两种不同降级管道的行为差异，我们明确了问题的根源在于不正确的内存空间分配。这一发现为修复编译器行为提供了明确的方向，同时也提醒我们在GPU代码生成时需要特别注意内存分配策略的正确性。