# cp2k-hip **Repository Path**: cui-rongpei/cp2k-hip ## Basic Information - **Project Name**: cp2k-hip - **Description**: 第二届中国科学院先导杯 CP2K赛题三等奖 CP2K的HIP移植与优化 - **Primary Language**: Unknown - **License**: GPL-2.0 - **Default Branch**: master - **Homepage**: None - **GVP Project**: No ## Statistics - **Stars**: 0 - **Forks**: 0 - **Created**: 2022-08-08 - **Last Updated**: 2024-05-31 ## Categories & Tags **Categories**: Uncategorized **Tags**: None ## README * # cp2k-hip ​ 此项目是第二届中科院先导杯的一道赛题。正是这道赛题让我印象最为深刻，此项比赛一共持续六个多月，从一开始真真正正的小白，连代码怎么看都不知道，更别提移植为何物了，但是也是这道题磨砺了我不屈不挠的精神，想尽各种办法解决问题的途径！ ​ 当时国内几乎没有甚至是搜不到一点移植相关的资料，唯一的资料就是AMD中国发布的全英文的讲解视频，但是里面也只是大致的方法，其中要走的路要踩的坑，对于小白来说太不友好。当时为了尽可能多的了解cp2k与Hip移植相关的内容，自己还关注了cp2k的官方邮箱社区，每天都会给自己发邮件看看有没有Hip的内容。**但是正是这样，我渐渐发现自己爱上了这种没有固定答案，一切要靠自己去研究琢磨最后终于成功的感觉！**从4月份到10月份，光是移植都花费了将近两三个月的时间，在此感谢先导杯比赛官方，感谢曾经帮助过我们的老师，大佬们。 ​ 在此贴出大佬的代码供大家学习参考。 ## 引言 ​ CP2K[1]是马克斯-普朗克研究中心于2000年发起的一项用于固体物理研究的项目，后转由ETH和苏黎世大学维护，成为一个开源的项目。CP2K以其软件跨平台的开放性、卓越的计算速度、优秀的横向扩展（scale out）以处理大体系的能力，获得了一大批计算科学家的青睐，并逐渐构筑起活跃的开源社区。 ## CP2K性能纵览 ​ CP2K与传统的HPC软件类似，支持三种并行方式：只使用MPI并行、只使用单机OpenMP并行、混合MPI+OpenMP并行。并行方式的选择极大地影响着软件性能（我们将在第6节详细讨论这一点），不失通用性，我们以下讨论均为混合MPI+OpenMP并行。 ​ 作为大规模分布式HPC软件，CP2K端到端所用的时间由计算时间、数据交换时间、MPI进程间等待和OpenMP线程间等待时间三部分决定。对于计算，其中热点调用部分（DBCSR模块[2]计算分布式稀疏矩阵乘法）、不可并行部分（GRID模块计算GPW方法的collocate、integrate）将决定计算部分的性能；对于数据交换，CP2K调用MPI的数据交换部分与计算部分异步执行，相对庞大的计算量而言时间被掩盖；对于进线程间等待，除应尽量保证数据的划分尽量均衡外，还应保证MPI的总进程数为平方数避免资源等待浪费：DBCSR模块在计算分布式稀疏矩阵乘法时使用经典的Cannon算法[3]。 ​ CP2K自身已经内置较为完善的profiling机制：对于计算，在重点函数的出入口处有性能埋点，绝大部分热点函数已经涵盖；对于数据交换，会统计各类MPI调用的次数和数据量，可用作初步参考。此外，可以通过TRACE和CALLGRAPH定向地了解CP2K运行过程中指定调用的耗时。 ​ 通过初步profiling，计算部分的热点主要为DBCSR模块执行的分布式稀疏矩阵乘法，其次为GRID模块执行的GPW方法的collocate和integrate，我们将通过将它们offload到异构DCU加速卡上实现大幅性能提升。 ## **国产异构DCU加速卡** ### 硬件架构与计算模型 ​ 国产异构DCU加速卡的硬件架构与NVIDIA GPU类似，大部分模块可以类比。如图1所示，一块DCU计算卡分为多个CU（Compute Unit，类比于CUDA的streaming multiprocessor），每个CU内部有自己的寄存器堆、ALU等硬件资源。每个CU内包含4个独立的SIMD，每个SIMD的宽度为16，即一个指令时钟周期（SCLK）内16个线程在一个SIMD上同步执行。硬件调度的最小单位为wavefront（类比于CUDA的warp），包含64个线程，在连续的4个SCLK中，它们在同一个SIMD上运行。每个SIMD可以容纳至多10个wavefront同时运行，且一个SCLK可以发射至多6条不同种类的指令。 ​ DCU的存储是分层的，从最快到最慢分别为寄存器、LDS（Local Data Share，类比于CUDA的shared memory）、L1缓存、L2缓存、显存（global memory）。每个SIMD有各自的寄存器，其中通用寄存器分为各线程独占的向量通用寄存器（VGPR，64KiB）和一个wavefront的64个线程共享的标量通用寄存器（SGPR，12KiB）。每个CU独享64KiB LDS。L1缓存分为数据缓存和指令缓存，每个CU独享16KiB L1数据缓存，每4个CU共享32KiB L1指令缓存。L2缓存不区分数据和指令缓存，由所有CU共享，且保证CU间的缓存一致性。最后，所有CU共享16GiB显存。 ​ 核函数的执行模型与CUDA类似：执行时三维grid中包含多个workgroup（类比于CUDA的block），每个workgroup也是三维的，包含多个workthread（类比于CUDA的thread）。映射到硬件层面，一个workgroup被分配到一个CU，属于同一个workgroup的多个wavefront可能在多个SIMD上同时执行。 ### DCU硬件参数 表1中列举了大部分重要的DCU硬件参数。 表1　国产DCU的重要硬件参数 | ***\*硬件参数\**** | ***\*描述\**** | | ------------------ | :----------------------------------------------------------: | | CU | 60 CU/DCU，总计5.065T VInst/sec（单精度乘加10.130T Flop/s） | | SIMD | 4 SIMD/CU，16 threads/SCLK/SIMD，即64 Inst/SCLK/CU | | VGPR | 64 KiB/SIMD，4B字，即最多256 VGPR/thread | | SGPR | 12 KiB/SIMD，4B字，即最多128 SGPR/thread | | LDS | 64 KiB/CU，4B字，32 bank，即1/2 wavefront/cycle，128 B/cycle | | L1指令缓存 | 32 KiB/4 CU，64B 缓存行，4路128组，LRU替换策略；约 32 B/cycle/4 CU | | L1 数据缓存 | 16 KiB/CU，64B缓存行，64路4组，LRU替换策略；约64 B/cycle/CU | | L2 缓存 | 4 MiB/DCU，64B缓存行，16路64组，LRU替换策略；约64 B/cycle/channel | | 显存 | 16 GiB/DCU，HBM2，总计2048 bit总线，总计带宽约576 GiB/sec | | 时钟频率 | 指令时钟周期SCLK 1143MHz，有效访存时钟周期 MCLK 2*800MHz | ## DBCSR模块的DCU移植和优化 ​ DBCSR是一个独立的计算分布式稀疏矩阵乘法的模块，主要为CP2K提供高效的、可线性水平扩展的矩阵乘法计算。整体而言，该模块可分为两层：第一层处理矩阵乘法在各个MPI进程间的划分，分布式稀疏矩阵在这一步中由分块CSR格式表示，然后被划分为较大的稀疏子矩阵，这些稀疏子矩阵几乎是同质的，即每个MPI进程的计算量几乎相同，MPI进程之间遵循Cannon矩阵乘算法进行消息传递和规约（2D网格）。第二层处理一个MPI进程内的多个OpenMP线程的并行，在经过单个进程内的一系列进一步的任务拆分和过滤后，同样大小的稠密的小矩阵被堆叠在一起（称为一个stack），由一个小矩阵乘法（SMM，small matrix multiplication）库执行批量的矩阵乘法和累加。 ​ DBCSR现有的任务划分、过滤和规约十分优秀，各MPI进程、OpenMP线程间的计算量基本均衡，递归拆分和任务过滤进一步大大减少了需要的计算量，数据传输也通过双缓冲（double buffer）、多流（multiple streams）等方式被计算较好的掩藏，决定算法最终性能的是小矩阵乘法库的性能。截至CP2K 8.1.0版本，DBCSR支持的小矩阵乘法库包括：为x86 CPU优化的libsmm、libxsmm，为NVIDIA GPU优化的自研的libcusmm[4]，我们为其增加了基于HIP的国产异构DCU加速卡支持。 ​ DCU的支持参考了libcusmm的实现。libcusmm的算法依旧遵照“逐级分块以发挥计算单元性能，逐级缓存以减小访存带宽压力，计算访存重叠以掩盖访存时间”的策略，实现的算法模板如表2所示，并针对不同大小的小矩阵乘法使用不同的算法模板和参数组合。对于某个特定大小的小矩阵乘，算法选择和参数组合可以来自auto tuning（即尝试可能的算法模板和参数组合搜索性能最佳的实现）或prediction（即根据已确定的实现推测一个较好的实现），这与传统矩阵乘法库的实现方式时分类似。DCU的支持基本沿用了libcusmm的算法模板，但将其从CUDA移植到HIP，并针对DCU硬件调整了搜索参数的空间。 表2　libcusmm的小矩阵乘法算法 | 核函数 | 参数数量 | 描述 | | -------- | -------- | :----------------------------------------------------------: | | tiny | 3 | 为特别小的矩阵优化，矩阵大小需能完全放进共享内存；目标矩阵中的每个位置由一个线程计算；内存读写和计算间没有重叠。 | | small | 5 | 每个线程计算目标矩阵中的一片tile，并在共享内存中累加；读取时直接从显存读入共享内存，再从共享存储中读出用于计算。 | | medium | 5 | 每个线程计算目标矩阵中的一片tile，并在共享内存中累加；读取时从显存读入寄存器再写入共享内存，将寄存器用于计算。 | | largeDB1 | 7 | 每个线程计算目标矩阵中的一片tile，并在共享内存中累加；读取时直接从显存读入共享内存，再从共享存储中读出用于计算；显存和共享内存间使用双缓冲。 | | largeDB2 | 7 | 每个线程计算目标矩阵中的一片tile，并在共享内存中累加；读取时从显存读入寄存器再写入共享内存，将寄存器用于计算；显存和共享内存间使用双缓冲。 | ​ 此外，由于HIP工具链相对CUDA更不成熟，我们贴合HIP工具链做了大量处理和优化。其一，CUDA工具链的nvcc编译所得的SASS汇编质量较好，不需要手工优化； HIP工具链的clang编译所得的GCN汇编质量不佳，需要调整较多clang的编译参数以达到较好的性能。其二，libcusmm使用nvrtc[5]运行时编译（JIT，Just-in-time）生成所需的GPU kernel，HIP工具链对应的hiprtc[6] JIT时间大大多于nvrtc，且hiprtc中有全局锁，多线程无法并行编译，故第一次编译时所需时间明显长，在MPI进程数较多时该问题更加严重。对此，我们将部分DCU上的kernel提前编译（AOT，Ahead-of-time）并存放到文件，运行时可直接使用。 ​ 将DBCSR移植到DCU后，DBCSR的计算性能瓶颈基本解决。视不同的矩阵大小，DBCSR在DCU上可达到CPU性能的5 ~ 10倍。 ## GRID模块的DCU移植和优化 ​ GRID是一个独立的计算GPW中collocate和integrate的模块，由于其处在性能关键路径上，高版本的CP2K使用C++重写GRID模块，以便于性能工程师能方便的优化其性能。 ​ GRID C++模块的接口定义清晰，截至CP2K 8.1.0版本，其支持后端包括CPU、NVIDIA GPU（使用CUDA工具链）、CPU和GPU混合计算，我们为其增加了基于HIP的国产异构DCU加速卡支持。 ​ 我们移植和优化的GRID模块主要分为collocate和integrate两大部分，移植时主要参考其CUDA实现，并对应地将cublas库替换为hipblas库。除简单对应移植外，我们进行了如下性能优化： ​ 其一，削减双精度浮点数的原子加（atomicAdd）调用。CUDA的实现中有大量的原子加，包括每个线程在不同位置执行原子加、同一个warp（DCU中的wavefront）的线程对同一个位置执行原子加两种情形。由于DCU上显存和LDS均不支持浮点数的原子加，都是通过回退到循环地读取+原子比较交换（CAS）实现，浮点数的原子加性能十分糟糕。同一个wavefront的线程对同一个位置执行原子加时性能劣化更明显，最后成功的线程至少需要重试至少63次。我们通过代码重构，同时引入DPP指令用于计算一个wavefront内的原子加，大大减少了双精度浮点数的原子加调用数量。 ​ 其二，使用现场计算替换一些查表的函数。GRID模块的计算（无论collocate还是integrate）均为访存密集型，CPU上的一些查表操作（访存换计算）在DCU上不再合适，我们将部分相对计算简单的查表操作用每次计算代替（计算换访存），减轻了访存压力。 ​ 其三，更充分地利用共享内存。我们注意到核函数的参数放在显存里，且在计算过程中会反复访问，而共享内存事实上并没有充分利用。我们在核函数开头将参数复制到共享内存，之后也通过共享存储访问数据，进一步带来了少量的性能提升。 ​ 经过移植优化，GRID模块的计算性能瓶颈得到缓解，integrate计算在DCU上可达到CPU性能的约4倍，collocate计算在DCU上可达到CPU性能的约5倍。 ​ 此外，我们注意到CP2K 8.2.0版本中，GRID模块也新增了HIP工具链的支持。通过简单实验，我们发现CP2K 8.2.0中的实现collocate与我们的实现性能相当，integrate结果错误无法比较，初步推测该错误与HIP工具链版本过低有关。 ## CPU端性能调优 ​ 除将DBCSR模块和GRID模块移植到DCU外，在CPU侧我们也进行了进一步的调优，包括： ​ 其一，MPI与OpenMP并行方式的选择。可用资源限定为32 CPU核，而DBCSR模块在MPI进程间运行Cannon算法，要求MPI进程数量为平方数。我们实验了如下组合：4 MPI进程 ![img](file:///C:\Users\ADMINI~1\AppData\Local\Temp\ksohtml2368\wps1.jpg) 8 OpenMP线程、16 MPI进程 ![img](file:///C:\Users\ADMINI~1\AppData\Local\Temp\ksohtml2368\wps2.jpg) 2 OpenMP线程、64 MPI进程（超配）![img](file:///C:\Users\ADMINI~1\AppData\Local\Temp\ksohtml2368\wps3.jpg) 1 OpenMP线程，最终选择了16 MPI进程 ![img](file:///C:\Users\ADMINI~1\AppData\Local\Temp\ksohtml2368\wps4.jpg) 2 OpenMP线程的并行方式。此外，DBCSR模块的OpenMP支持在HIP工具链下有少量问题，我们对其进行了修复。 ​ 其二，CPU侧计算库的选择。CP2K支持多种CPU计算库，如Intel MKL、OpenBLAS + ScaLapack，以及新增了COSMA库作为DBCSR的替代。在实验比较后，我们选择了使用OpenBLAS + ScaLapack 作为CPU侧计算库，并使用优化过的DBCSR模块。 ## 编译安装步骤 #### 加载所需的环境变量 ``` source ~/cp2k/env.sh ``` #### 切换到toolchain目录，安装toolchain，并复制生成的arch文件 ``` cd ~/cp2k/tools/toolchain ./install_cp2k_toolchain.sh # 不需要额外参数，脚本里默认已经设置好 cp install/arch/* ../../arch/ ``` #### 切换到源码目录，进行编译（可能需要较长时间） ``` cd ~/cp2k ``` 使用文本编辑工具修改 arch/local_hip.psmp，在第7行末尾添加 -D__GRID_CUDA，添加后变为：DFLAGS = -D__HIP_PLATFORM_HCC__ -D__ACC -D__DBCSR_ACC -D__LIBXSMM -D__parallel -D__LIBINT -D__LIBXC -D__SCALAPACK -D__GRID_CUDA ``` make -j 32 ``` ### 运行步骤 ``` cd ~/cp2k source env.sh mpirun -np 16 --map-by node:PE=2 exe/local_hip/cp2k.psmp -i caspra-final/H2O-256.inp -o H2O-256.out mpirun -np 16 --map-by node:PE=2 exe/local_hip/cp2k.psmp -i caspra-final/H2O-dft-ls.NREP2.inp -o H2O-dft-ls.NREP2.out ``` ## 实验性能 经过DBCSR和GRID模块的DCU移植和进一步的CPU侧调优后，我们的端到端性能有较大提升。性能对比如表3所示。 表3(1)　CP2K优化前后H2O-256性能对比 | ***\*指标\**** | ***\*优化前/秒\**** | ***\*优化后/秒\**** | ***\*性能提升\**** | | :-----------------: | :-----------------: | :-----------------: | :----------------: | | 端到端时间 | 1111.454 | 478.313 | ***\*2.32x\**** | | DBCSR时间 | 558.378 | 162.125 | 3.44x | | GRID collocate时间 | 102.488 | 45.346 | 2.26x | | GRID integrate 时间 | 107.045 | 68.579 | 1.56x | 表3(2)　CP2K优化前后H2O-dft-ls.NREP2性能对比 | ***\*指标\**** | ***\*优化前/秒\**** | ***\*优化后/秒\**** | ***\*性能提升\**** | | :-----------------: | :-----------------: | :-----------------: | :----------------: | | 端到端时间 | 91.444 | 62.398 | ***\*1.47x\**** | | DBCSR时间 | 56.350 | 37.054 | 1.52x | | GRID collocate时间 | 9.795 | 5.792 | 1.69x | | GRID integrate 时间 | 10.186 | 4.270 | 2.39x | ## 总结与展望 ​ 总结而言，我们在国产异构平台上充分利用了DCU和CPU的强大算力，对CP2K进行了移植和优化。DBCSR模块的DCU移植上，我们新增支持了HIP工具链，同时支持了更多不同大小的矩阵乘法，搜索了更优的参数组合，部分核函数AOT以减小计算时JIT开销；GRID模块的DCU移植上，我们在DCU上高效地实现了collocate和integrate；CPU侧的性能优化上，我们修复了DBCSR的OpenMP支持，进行了MPI和OpenMP的并行方式调优，基础矩阵运算库选型调优等。在不改变计算方法、不降低计算精度的前提下，优化后的CP2K性能可达到官方版本的3 ~ 5倍。 ​ 另一方面，限于时间原因，目前仍有许多工作没有进行，它们可能对CP2K的性能优化有进一步的益处：DBCSR的核函数是为NVIDIA GPU + CUDA工具链优化的，在DCU + HIP工具链上的性能还有极大的提升空间；DCU上的GPU Direct（即数据直接在Infiniband网卡和DCU显存间交换，不经过CPU主存）能进一步性能；此外，优化后的CP2K在更大规模集群上的性能、扩展能力还有待评估。