LAMMPS中Kokkos后端在大规模系统下的性能回归问题分析

问题概述

在LAMMPS分子动力学模拟软件中，当使用Kokkos后端处理大规模系统复制(replicate)操作时，研究人员发现了一个严重的性能退化问题。具体表现为：在512个节点以上的大规模并行计算中，Special.build()组件的执行时间出现了25倍的显著下降。

技术背景

LAMMPS中的replicate操作用于复制分子系统，在此过程中会调用Special.build()来构建1-2、1-3、1-4键列表并进行标记。这一过程在Kokkos后端实现中，特别是在atom_map_kokkos.cpp文件中的map_set函数内，存在性能瓶颈。

问题表现

通过详细的性能分析发现：

1. 在GPU-Kokkos后端下，4096个节点的系统复制时间从16.378秒激增至423.27秒
2. 问题主要集中在rank 0进程上，其执行时间比其他进程高出两个数量级
3. 性能瓶颈出现在map_set函数中的排序操作部分，特别是create_permute_vector和sort调用

根本原因分析

深入调查发现几个关键因素：

1. 标签分布不均：rank 0进程处理的原子标签范围(min/max)远大于其他进程。例如在25亿原子的测试中，rank 0需要处理从1到最大原子标签的完整范围，而其他进程只需处理局部范围。

2. 排序算法效率：Kokkos::BinSort在标签范围极大的情况下表现不佳，特别是当存在极端不平衡的标签分布时。

3. 后端差异：该问题仅出现在GPU-Kokkos后端，CPU版本和Serial-Kokkos版本表现正常甚至有所提升。

解决方案与建议

针对这一问题，开发团队提出了几种解决方案：

1. 临时解决方案：回退到使用CPU构建原子映射表，虽然会牺牲部分性能，但可以保证稳定性。

2. 长期修复：等待Kokkos库的排序算法优化，特别是针对AMD/HIP架构的支持。

3. 代码优化：考虑使用替代排序算法如sort_by_key，避免使用基于分箱(bin)的排序方法。

性能数据对比

测试数据展示了不同配置下的性能表现(单位为秒)：

配置类型	节点数	当前开发版	稳定版(23Jun2022)
默认CPU	1	1.042	1.010
默认CPU	512	1.444	1.296
GPU-Kokkos	1	4.798	3.731
GPU-Kokkos	512	16.194	5.716
GPU-Kokkos	4096	423.27	16.378

结论与展望

这一性能问题揭示了在大规模并行计算中负载均衡的重要性，特别是在处理原子映射这类基础数据结构时。开发团队正在积极寻求解决方案，既包括短期的临时修复，也包括长期的算法优化。对于当前面临紧迫计算任务的用户，建议暂时回退到稳定版本或使用CPU构建原子映射表。

未来随着Kokkos库的持续优化和LAMMPS代码的改进，预期这一问题将得到彻底解决，使Kokkos后端能够充分发挥其在大规模计算中的性能优势。