MUSA移植案例系列 - 分子动力学模拟应用GROMACS
2024-06-25
1. 介绍
Gromacs 是一款用于模拟生物大分子的软件，其核心是一个高效的分子动力学模拟程序。Gromacs 通过使用 GPU 加速来提高计算效率，目前 Gromacs 的 GPU 加速功能已经支持 NVIDIA、AMD 的 GPU 以及能够运行 OpenCL 程序的设备。为了让 Gromacs 能够在 MUSA 平台上运行，我们需要将 Gromacs 移植到 MUSA 平台上，从而实现利用 MUSA 平台的高性能计算资源来加速 Gromacs 计算的目的。
目前，Gromacs MUSA 加速版，包括利用 OpenMP 支持单节点多卡的版本，和利用 MPI 支持多节点多卡的版本都已经功能完备，并在 KUAE 集群上验证通过。本文将介绍将 Gromacs 移植到 MUSA 平台上的过程。
2. 准备工作
2.1. MUSA 环境准备
本文使用的硬件平台是 KUAE 千卡集群，其配备的 GPU 类型是 S4000，架构为曲院2代（后称 QY2）。理论上 QY1/2 的架构差别不大，因此本文的移植过程也适用于 QY1 架构的 GPU，如 S80。
本文使用的操作系统是 Ubuntu 20.04，GPU 驱动和 MUSA Toolkits 尽量使用最新的版本；MUSA 环境的搭建请参考 MUSA 官方文档。
为了实现单节点多卡以及多节点多卡的加速，系统中还应根据需要分别安装 OpenMP 和 OpenMPI。
2.2. Gromacs 源码准备
Gromacs 的源码可以从官方网站下载，也可以从 GitHub 上的 Gromacs 仓库获取。本文使用的 Gromacs 版本是 2023.3。更新的版本应该也适用于本文的移植过程。
2.3. 回归测试
Gromacs 本身包含了非常全面的单元测试，但如果希望进行更严格的测试，Gromacs 也提供了一些回归测试的数据集。在编译 Gromacs 前，如果给出 REGRESSIONTEST_DOWNLOAD=ON 的选项，Gromacs 会自动下载回归测试数据集。但是该数据集比较大，下载时间较长，在移植阶段需要频繁更改代码和测试的场景下，会降低效率。因此推荐下载对应版本的回归测试数据集，然后在编译 Gromacs 时使用 REGRESSIONTEST_PATH 指定回归测试数据集的路径。回归测试的下载路径可以在 CMake 脚本中找到，其相对路径为 tests/CMakeLists.txt。
2.4. 真实数据集
本文使用了常用的两个对 Gromacs 进行性能比较的数据集，即 STMV 和 BenchPEP-h。其下载地址如下：
3. 基于 Musify 的 Runtime API 移植及编译脚本移植
Gromacs 已有的支持 GPU 加速的框架已经非常完备，可方便地编译到NVIDIA GPU 或者 AMD GPU 以及支持 OpenCL 的设备上运行。由于 MUSA 编程模型跟 CUDA 较为相似，所以主要思想是我们以 Gromacs 之中已有的 CUDA 代码为基础，利用 Musify 等手段将其转换为 MUSA 代码，同时进行一些必要的修改，使其能够正确、高效地在 MUSA 平台上运行，具体步骤如下：
3.1. API 迁移: 基于 Musify
对于 CUDA kernel 调用相关的代码，可以分成两种情况来处理：
- 仅仅包含 CUDA Runtime API 调用的文件，后缀可能为
.cpp或者hpp，以及包含有 CUDA kernel 的文件，后缀可能为.cu或者.cuh，以及名为cuda的文件夹，这些文件可以直接使用 Musify 进行复制和转换，转换后的 MUSA kernel 文件后缀为.mu或者.muh，具体方法可参考 Musify 的相关博客； - Gromacs 框架 C++ 代码，包含有对各个加速库的调用封装，这里应该手动修改，复制 CUDA 相关代码块，加入
GMX_GPU_MUSA宏定义，然后修改其中的调用函数等。目前 MUSA Toolkits 暂时不支持 Graph，所以这方面相关的选项请一直保持关闭。
3.2. 编译脚本的修改
与代码类似，我们也仿照 CUDA 在 Gromacs 中的编译选项，来新增 MUSA 的编译选项，比较关键的有：
GMX_GPU：新增MUSA选项；GMX_GPU_FFT_LIBRARY：新增muFFT选项。- 新增
MUSA_ROOT选项，指定 MUSA 的安装路径。
其他与 CUDA 相关的编译脚本和选项也应该复制一份，然后针对 MUSA 的具体情况修改，如将 MUSA 源代码加入编译等等。
经过上述修改后，我们应该可以使用 CMake 来编译 Gromacs 了。如果编译过程有错误，可以根据错误信息来进一步修改代码和编译脚本。以下给出一个简单的编译脚本示例：
cmake -DGMX_OPENMP=ON \
-DCMAKE_INSTALL_PREFIX=build/install \
-DGMX_FFT_LIBRARY=fftw3 \
-DGMX_GPU=MUSA \
-DGMX_GPU_FFT_LIBRARY=muFFT \
-DREGRESSIONTEST_PATH=/home/mingxu/Downloads/gromacs/gromacs-regressiontests-release-2023 \
-DGMX_BUILD_OWN_FFTW=ON \
-DGMX_THREAD_MPI=ON \
-DGMX_PHYSICAL_VALIDATION=OFF \
-DCMAKE_BUILD_TYPE=Debug \
-DMUSA_ROOT=/usr/local/musa \
-B build
cmake --build build -j 6
我们推荐在测试开发阶段使用 Debug 编译模式，因为 Gromacs 代码中包含大量的 assert 函数调用，我们也可以加入自己的 assert 函数调用来帮助我们调试代码，这些都只在 Debug 模式中生效。在测试通过后，我们可以使用 Release 编译模式来编译 Gromacs，以获得更好的性能。
4. MUSA kernel 参数适配与微调
从上述步骤可以看出，MUSA 对于 CUDA 的兼容性是非常高的，甚至只需要更改 Runtime API 的函数名前缀，复制 GPU kernel，就能够通过 MUSA Toolkits 编译。但是目前 MT GPU 与 CUDA GPU 的架构仍存在差异，体现在 QY 上面，就是每个 warp 有 128 个线程。我们不需要更改 GPU kernel 的整体结构和调用流程，但是当任务分配、数据同步等操作涉及到线程数量相关的计算时，我们就需要对代码进行微调，以适应 QY 架构的特点。
4.1. 代码正确性测试
在编译好代码后，就应该用单元测试来指导我们后续的工作了。如果你按照上述命令编译 Gromacs 成功，则编译运行所有测试的的命令如下：
cd build
make tests -j 6
make check
由于 Gromacs 所有编译过程都是由 CMake 完成的，所以我们可以用 ctest 来查询和运行单个测试。列举所有测试的命令如下：
ctest -N
运行单个测试的命令如下：
ctest -R <test_name> --verbose
4.2. GPU kernel 移植
QY 架构与 CUDA 架构的最大差别，是 QY 架构每个 warp 有 128 个线程，而 CUDA 架构每个 warp 有 32 个线程。这里影响比较大的地方有两个：
- 与 warp size 相关的常量计算需要修改；否则有可能导致程序运行不正确；
- warp vote 和 shuffle 需要考虑 128 线程，需要修改。
下面我们分别来讲解这两个问题。
4.2.1. warp size 相关的常量计算
在 src/gromacs/nbnxm/musa/nbnxm_musa.mu 文件中，有专门读取 deviceProperties 对象中 warpSize 的语句。这个结果在 QY 架构上必然是 128：
/* assert assumptions made by the kernels */
GMX_ASSERT(c_nbnxnGpuClusterSize * c_nbnxnGpuClusterSize / c_nbnxnGpuClusterpairSplit
== deviceInfo->prop.warpSize,
"The MUSA kernels require the "
"cluster_size_i*cluster_size_j/nbnxn_gpu_clusterpair_split to match the warp size "
"of the architecture targeted.");
那这里的 c_nbnxnGpuClusterSize 和 c_nbnxnGpuClusterpairSplit 也需要做相应修改，以满足这个条件。我们取 c_nbnxnGpuClusterSize 为 16，c_nbnxnGpuClusterpairSplit 为 2，这样就满足了 warp size 为 128 的条件。
但是，c_nbnxnGpuClusterSize 的取值会影响到 c_clSize 的取值，具体请参照 src/gromacs/nbnxm/musa/nbnxm_musa_types.h 文件。大量的 MUSA kernel 的 block 维度都是由 c_clSize 决定的，所以 kernel 中跟维度相关的计算都有可能需要修改，具体见下文。
在 src/gromacs/gpu_utils/musa_arch_utils.muh 文件中，有一些 warp size 相关的常量定义，需要修改：
// static const int warp_size = 32;
// static const int warp_size_log2 = 5;
static const int warp_size = 128;
static const int warp_size_log2 = 7;
综上，Gromacs 中的 MUSA kernel 请过一遍，看看有没有 warp size 相关的常量计算，然后根据具体问题进行修改。如果相关 kernel 并没有涉及 warp size 的计算，那么就不需要修改。
4.2.2. warp vote 和 shuffle 相关的修改
在 QY 架构上调用 warp vote 和 shuffle 的函数有一定的复杂性，因为调用的函数参数与 CUDA 一致，每次调用最多有 32 个线程参与；而 QY 上每个 warp 有 128 个线程，所以可能需要额外的同步才能得到正确结果。例如在 src/gromacs/nbnxm/musa/nbnxm_musa_kernel.muh 中，我们应该有如下的修改：
# if GMX_PTX_ARCH < 300
// shared memory for 128 threads warp shuffle and vote synchronize
// (num_threads / warp_size) warp lanes
constexpr int num_warp_vote = THREADS_PER_BLOCK >> warp_size_log2;
__shared__ unsigned int warp_vote_sync[num_warp_vote];
const int gidx = tidxz * blockDim.x * blockDim.y + tidx;
const int lane_id = gidx & 31;
const int wide_lane_id = gidx & (warp_size - 1);
const int warp_id = gidx >> warp_size_log2;
# endif
// other codes ................
# ifdef PRUNE_NBL
# if GMX_PTX_ARCH >= 300
/* If _none_ of the atoms pairs are in cutoff range,
the bit corresponding to the current
cluster-pair in imask gets set to 0. */
if (!__any_sync(c_fullWarpMask, r2 < rlist_sq))
{
imask &= ~mask_ji;
}
# else
if (wide_lane_id == 0)
{
warp_vote_sync[warp_id] = 0;
}
__syncwarp();
auto tmp = __any_sync(c_fullWarpMask, r2 < rlist_sq);
if (lane_id == 0)
{
atomicOr(&warp_vote_sync[warp_id], tmp);
}
__syncwarp();
if (!warp_vote_sync[warp_id])
{
imask &= ~mask_ji;
}
# endif
# endif
上述代码中，我们将原有的代码和加入的代码用 GMX_PTX_ARCH 进行了区分，因为 QY 架构的版本都是小于 300 的。在 QY 架构上，我们单单只做 __any_sync 是不够的，因为这样只是 128 个线程中，每 32 个一组的线程知道了 __any_sync 的结果，但是这个结果并没有同步到其他的线程。所以我们需要额外的同步，这里我们使用了 shared memory 来存储临时结果，以及原子操作 atomicOr 来得到最终包含整个 warp 所有线程计算的结果。
4.2.3. 综合修改
这里仅举一个例子进行分析。在 src/gromacs/nbnxm/musa/nbnxm_musa_kernel_utils.muh 中，有一个 reduce_force_i_warp_shfl 函数，这个函数与 CUDA 实现一致：
/*! Final i-force reduction; this implementation works only with power of two
* array sizes.
* c_clSize = 8
*/
static __forceinline__ __device__ void reduce_force_i_warp_shfl(float3 fin,
float3* fout,
float* fshift_buf,
bool bCalcFshift,
int tidxj,
int aidx,
const unsigned int activemask)
{
fin.x += __shfl_down_sync(activemask, fin.x, c_clSize);
fin.y += __shfl_up_sync(activemask, fin.y, c_clSize);
fin.z += __shfl_down_sync(activemask, fin.z, c_clSize);
if (tidxj & 1)
{
fin.x = fin.y;
}
fin.x += __shfl_down_sync(activemask, fin.x, 2 * c_clSize);
fin.z += __shfl_up_sync(activemask, fin.z, 2 * c_clSize);
if (tidxj & 2)
{
fin.x = fin.z;
}
/* Threads 0,1,2 and 4,5,6 increment x,y,z for their warp */
if ((tidxj & 3) < 3)
{
atomicAdd(&fout[aidx].x + (tidxj & 3), fin.x);
if (bCalcFshift)
{
*fshift_buf += fin.x;
}
}
}
在 CUDA 环境下，上述函数运行的 block 维度是 8X8，每个 warp 有 4 行，每行 8 个线程。这样就可以安全地使用每个 warp 的前三行来分别对 float3 的 x, y, z 成员进行 global memory 上的原子求和操作。而在 QY 架构上，因为 c_clSize 为 16，所以参与 warp shuffle 操作的线程仅有 2 行，每行 16 个线程，这样就没法完全照搬 CUDA 的实现了。我们需要对这个函数进行修改，使其能够在 QY 架构上正确运行。
// c_clSize = 16
// fshift_buf: float3 array in shared memory
static __forceinline__ __device__ void reduce_force_i_warp_shfl_qy(float3 fin,
float3* fout,
float3* fshift_buf,
bool bCalcFshift,
int tidxj,
int aidx,
const unsigned int activemask)
{
__shared__ float3 fout_buf[c_clSize];
if (threadIdx.y == 0)
{
fout_buf[threadIdx.x] = make_float3(0.0F);
}
__syncthreads();
fin.x += __shfl_down_sync(activemask, fin.x, c_clSize);
fin.y += __shfl_down_sync(activemask, fin.y, c_clSize);
fin.z += __shfl_down_sync(activemask, fin.z, c_clSize);
if (!(tidxj & 1))
{
atomicAdd(&(fout_buf[threadIdx.x].x), fin.x);
atomicAdd(&(fout_buf[threadIdx.x].y), fin.y);
atomicAdd(&(fout_buf[threadIdx.x].z), fin.z);
}
__syncthreads();
if (tidxj == 0)
{
atomicAdd(&(fout[aidx].x), fout_buf[threadIdx.x].x);
atomicAdd(&(fout[aidx].y), fout_buf[threadIdx.x].y);
atomicAdd(&(fout[aidx].z), fout_buf[threadIdx.x].z);
if (bCalcFshift)
{
fshift_buf[threadIdx.x] += fout_buf[threadIdx.x];
}
}
}
5. 通过 MPI 实现多节点多卡并行计算
如果所有的单元测试都能够通过，那么我们就可以编译安装 Gromacs 为最终的仿真做准备了。这时我们需要用 Release 模式编译 Gromacs，然后安装到指定目录。同时，只要系统中 MPI 环境可用，稍作修改就可以运行多节点多卡的 Gromacs MUSA 版本了。编译安装过程示例如下：
cmake -DGMX_OPENMP=ON \
-DCMAKE_INSTALL_PREFIX=<install-dir> \
-DGMX_FFT_LIBRARY=fftw3 \
-DGMX_GPU=MUSA \
-DGMX_GPU_FFT_LIBRARY=muFFT \
-DREGRESSIONTEST_PATH=../gromacs-regressiontests-release-2023 \
-DGMX_BUILD_OWN_FFTW=ON \
-DGMX_THREAD_MPI=ON \
-DGMX_PHYSICAL_VALIDATION=OFF \
-DCMAKE_BUILD_TYPE=Release \
-DMUSA_ROOT=/usr/local/musa \
-DCMAKE_PREFIX_PATH=<openmp-dir> \
-DGMX_MPI=ON \
-DMPIEXEC_NUMPROC_FLAG=-np \
-B release
cmake --build release -j 6
cmake --install release
6. 真实数据集测试
6.1. 测试常量的修改
为了运行 STMV 和 benchPEP-h 数据集，我们需要修改一些常量，具体如下：
src/gromacs/ewald/pme_gpu_constants.h文件中，应加入如下代码：
#if GMX_GPU_MUSA
constexpr int c_spreadMaxWarpsPerBlock_musa = 4;
constexpr int c_solveMaxWarpsPerBlock_musa = 4;
constexpr int c_gatherMaxWarpsPerBlock_musa = 4;
//! Spreading max block size in threads
static constexpr int c_spreadMaxThreadsPerBlock = c_spreadMaxWarpsPerBlock_musa * warp_size;
//! Solving kernel max block size in threads
static constexpr int c_solveMaxThreadsPerBlock = c_solveMaxWarpsPerBlock_musa * warp_size;
//! Gathering max block size in threads
static constexpr int c_gatherMaxThreadsPerBlock = c_gatherMaxWarpsPerBlock_musa * warp_size;
//! Gathering min blocks per MUSA multiprocessor (determined empirically to give best performance)
static constexpr int c_gatherMinBlocksPerMP = GMX_MUSA_MAX_THREADS_PER_MP / c_gatherMaxThreadsPerBlock;
#endif // GMX_GPU_MUSA
src/gromacs/ewald/pme_gpu_internal.cpp文件中，应加入如下代码：
#if GMX_GPU_MUSA
constexpr int c_pmeAtomDataBlockSize = 128;
#elif !GMX_GPU_SYCL
constexpr int c_pmeAtomDataBlockSize = 64;
#else
// Use more padding to support 64-wide warps and ThreadsPerAtom::Order
constexpr int c_pmeAtomDataBlockSize = 128;
#endif
6.2. 基于 OpenMP 的测试
以 STMV 为例，可以使用如下命令来运行单节点多卡测试：
cd <stmv_dir>
source /path/to/gromacs/install/bin/GMXRC
gmx mdrun -ntmpi 8 -ntomp 16 -nb gpu -pme gpu -npme 1 -update gpu -bonded gpu -nsteps 1000 -noconfout -dlb no -pin on -v -gpu_id 4,5,6,7
由于我们一直是用的 OpenMP 的模式来编译的 Gromacs，所以完全可以使用同一节点上的多个甚至全部 MT GPU 来运行任务。
6.3. 基于 MPI 的测试
单节点测试：
mpirun -np 16 gmx_mpi mdrun -nb gpu -pme gpu -npme 1 -update gpu -bonded gpu -nsteps 3000 -noconfout -dlb no -pin on -v -gpu_id 0 -s /path/to/stmv/topol.tpr
双节点测试：
mpirun -np 8 -hostfile ../hostfile gmx_mpi mdrun -nb gpu -pme gpu -npme 1 -update gpu -bonded gpu -nsteps 3000 -noconfout -dlb auto -pin on -v -s /path/to/stmv/topol.tpr -gputasks 0123
