DirectXShaderCompiler中HLMatrixLowerPass模块的内存释放问题分析

问题背景

在DirectXShaderCompiler项目中，当处理包含无限循环和矩阵返回值的特定着色器代码时，编译器会出现ASAN(AddressSanitizer)报告的堆释放后使用(heap-use-after-free)错误。这个问题暴露了HLMatrixLowerPass模块在处理矩阵转换时存在的内存管理缺陷。

问题重现

问题可以通过以下简单的计算着色器代码重现：

float4x2 foo() {
  while (true) {
  }
  return float4x2((0.0f).xx, (0.0f).xx, (0.0f).xx, (0.0f).xx);
}

[numthreads(1, 1, 1)]
void main() {
  float4x2 _e1 = foo();
}

这段代码的关键特征是：

1. 函数foo包含一个无限循环
2. 循环后返回一个矩阵值
3. 主函数调用foo并接收返回值

问题根源分析

HLMatrixLowerPass的工作机制

当编译器处理返回矩阵的函数时，HLMatrixLowerPass模块会执行以下操作：

1. 创建一个"matrix to vector"的转换存根函数(stub function)
2. 在原始矩阵构造函数调用前插入对这个存根函数的调用
3. 使用IRBuilder创建新的CallInst指令，引用这个存根函数

内存管理问题

问题的核心在于内存管理的不一致：

1. 存根函数的创建与销毁：

   • 存根函数由TempOverloadPool临时管理
   • 当HLMatrixLowerPass完成工作后，TempOverloadPool被销毁
   • 销毁过程中调用eraseFromParent()删除了存根函数

2. 指令引用未清理：

   • 虽然存根函数被删除，但之前插入的CallInst指令仍然保留
   • 这些指令继续引用已被释放的函数对象

3. 后续处理冲突：

   • GlobalDCE(全局死代码消除)过程遍历模块中的函数
   • 当检查到包含无限循环的函数时，尝试处理其后继代码
   • 遇到引用已释放存根函数的CallInst指令，触发ASAN错误

技术细节深入

关键调用链

1. HLMatrixLowerPass::runOnModule创建TempOverloadPool
2. 处理矩阵操作时调用getLoweredByValOperand创建转换存根
3. IRBuilder::CreateCall插入引用存根函数的指令
4. TempOverloadPool析构时调用eraseFromParent删除存根函数
5. GlobalDCE处理时访问已释放内存

内存生命周期问题

问题的本质在于对象生命周期的管理不一致：

• 存根函数(Function对象)的生命周期由TempOverloadPool管理
• 引用存根函数的CallInst指令的生命周期由基本块管理
• 缺乏协调机制确保引用关系正确解除

解决方案探讨

短期修复方案

1. 恢复removeFromParent调用：

   • 回退到之前的实现方式
   • 避免立即删除Function对象
   • 可能带来内存泄漏风险

2. 显式清理引用指令：

   • 在删除存根函数前，遍历并删除所有引用指令
   • 需要精确跟踪所有创建的引用

长期改进建议

1. 改进TempOverloadPool设计：

   • 实现完整的引用计数机制
   • 确保资源释放前所有引用都被清理

2. 增强测试覆盖：

   • 添加针对无限循环与矩阵操作组合的测试用例
   • 增加ASAN检测的自动化测试

3. 重构矩阵处理逻辑：

   • 考虑使用LLVM内置的资源管理机制
   • 减少对临时池的手动管理

总结

这个问题揭示了DirectXShaderCompiler中矩阵处理模块的一个深层次设计问题。在复杂的编译器优化流水线中，需要特别注意跨阶段的对象生命周期管理。特别是当自定义Pass创建临时对象并被后续优化Pass处理时，必须确保引用关系的正确性。

对于编译器开发者来说，这个案例强调了：

1. 内存管理一致性的重要性
2. ASAN等工具在捕捉复杂内存问题中的价值
3. 边界条件测试的必要性(如无限循环与复杂类型组合)

该问题的修复需要权衡短期解决方案的简单性和长期架构改进的完整性，同时确保不影响编译器的正确性和性能。