oneDNN在B60平台上的GPU内存管理问题分析与解决

问题背景

在Intel B60平台（Battlemage架构GPU）上运行基于oneDNN的GEMM（通用矩阵乘法）运算时，开发人员观察到一个异常现象：每次运行都会导致GPU内存使用量小幅增加，且这些内存不会被释放。相比之下，在Arc A770平台上运行相同的代码则不会出现这个问题。

技术环境分析

该问题出现在以下环境中：

• 硬件平台：Intel B60 GPU（Battlemage架构）
• 操作系统：Ubuntu 24.04（使用HWE内核）
• 编译器：Intel icpx 13.2.0
• oneDNN版本：最初测试的是v3.6.1，后续测试了rls-v3.8分支
• 驱动程序：Level Zero 1.6.32961和OpenCL 25.09.32961.7

问题现象详细描述

开发人员使用xpu-smi工具监控GPU内存使用情况时发现，每次执行GEMM运算后，GPU内存使用量都会增加约12MB左右。这种内存增长是累积性的，随着程序运行时间的增加，最终可能导致GPU内存耗尽。

深入技术分析

经过技术团队深入调查，发现这个问题涉及多个层面的因素：

1. 内存池管理机制：oneDNN使用zero-pool（零初始化GPU内存池）来优化性能，这些内存池是按引擎(engine)对象分配的。

2. 引擎生命周期：在原始测试代码中，每次迭代都创建新的引擎对象，导致zero-pool内存不断累积。

3. 版本差异：在oneDNN v3.6.1中存在一个已知的内存泄漏问题，特别是在处理不同引擎对象时。这个问题在rls-v3.8分支中已通过特定提交修复。

4. 平台特性：Battlemage架构的GPU内存管理行为与Arc架构存在差异，这可能放大了内存管理问题的影响。

解决方案验证

技术团队提出了多种解决方案并进行了验证：

1. 设置环境变量：尝试通过ONEDNN_PRIMITIVE_CACHE_CAPACITY=0禁用primitive缓存，但发现这并不能完全解决问题。

2. 升级oneDNN版本：切换到rls-v3.8分支后，内存泄漏问题得到显著改善。测试数据显示内存使用在几次迭代后趋于稳定。

3. 重用引擎对象：保持引擎对象的单例模式，避免重复创建，这可以防止zero-pool内存的重复分配。

4. 驱动程序选择：测试发现使用OpenCL驱动比Level Zero驱动表现出更轻微的内存增长现象。

最佳实践建议

基于此问题的分析，我们建议开发人员：

1. 版本控制：使用oneDNN rls-v3.8或更高版本，其中包含了相关内存泄漏问题的修复。

2. 资源管理：在应用程序中重用引擎对象，而不是每次运算都创建新实例。

3. 监控工具：使用zesMemoryGetState等工具准确监控GPU内存使用情况，而不是仅依赖外部工具。

4. 环境配置：根据实际需求选择合适的运行时驱动（Level Zero或OpenCL）。

结论

GPU内存管理是高性能计算中的关键问题。oneDNN团队通过版本迭代解决了B60平台上的特定内存泄漏问题。开发人员应当注意引擎对象的生命周期管理，并保持oneDNN库的及时更新，以确保获得最佳的内存使用效率和计算性能。

对于使用Battlemage架构GPU的开发人员，特别建议采用上述最佳实践，以避免潜在的内存问题并优化应用程序性能。