Kokkos项目中的HIP后端原子操作性能分析与优化

背景介绍

在Kokkos高性能计算框架中，HIP后端作为AMD GPU平台的重要支持组件，其原子操作的性能表现直接影响着应用程序的运行效率。近期开发团队发现，在MI300A架构GPU上执行原子操作时出现了显著的性能下降问题，特别是对于64位整数类型的原子加法操作，执行时间比预期要长得多。

问题现象

开发人员最初在GitLab CI测试中观察到，hip.atomic_views_integral测试用例在MI300A GPU上的执行时间达到了约6.5分钟，而在MI250 GPU上仅需17秒左右。这种巨大的性能差异引起了团队的关注。

问题分析与定位

通过简化测试用例，开发团队创建了一个最小化的复现示例，发现当对64位整数类型执行原子加法操作时，MI300A的性能显著下降。进一步分析表明：

1. 对于double类型，性能表现正常
2. 对于int64_t类型，性能下降严重
3. 问题与输入值无关（无论是加0还是加1）

通过深入研究Kokkos和Desul（Kokkos的原子操作库）的实现，团队发现问题根源在于原子操作的实现路径选择上。对于64位整数类型，系统错误地选择了基于比较交换(CAS)的实现，而非更高效的专用原子指令。

技术细节

在Desul原子库中，原子操作的实现路径选择依赖于以下机制：

1. 对于"总是无锁"的类型（4字节、8字节等），默认使用CAS循环实现
2. 对于特定类型，提供专门的原子指令实现

问题出在HIP后端的实现中，对long类型的专用原子指令支持缺失，导致64位整数类型无法使用最优路径。相比之下，double类型由于有明确的实现，性能表现正常。

解决方案

开发团队提出了以下修复方案：

1. 在Desul的HIP原子操作实现中，显式添加对long类型的支持
2. 确保所有整数类型都能使用最优的原子指令路径

修复后的测试表明，64位整数类型的原子操作性能得到了显著提升，与其他类型和架构的性能表现趋于一致。

性能对比

通过在不同AMD GPU架构上的性能测试，团队收集了详细的数据：

• MI300A架构：
  • 修复前：int64_t原子操作性能极低
  • 修复后：性能提升显著，与其他类型趋于一致
• MI200系列（MI210/MI250）：
  • 需要启用-munsafe-fp-atomics编译选项以获得最佳浮点原子操作性能
  • 修复对整数类型性能影响较小
• MI100架构：
  • 表现与MI200系列类似

结论与建议

本次性能问题的解决过程展示了底层原子操作实现细节对整体性能的重要影响。对于Kokkos用户和开发者，建议：

1. 在AMD GPU上使用最新版本的Kokkos，以获得最佳原子操作性能
2. 对于MI200系列GPU，确保启用-munsafe-fp-atomics编译选项
3. 在性能关键代码中，注意原子操作的数据类型选择
4. 定期进行跨架构的性能基准测试，及时发现潜在问题

这一优化不仅解决了MI300A上的性能问题，也为Kokkos在AMD GPU平台上的整体性能提升做出了贡献。