商汤高性能计算 (HPC) 团队秉承开源之心，赋能材料计算，与北京科学智能研究院 (AI for Science Institute, Beijing) 的电子结构团队合作，助力 DeepModeling 社区打造源自中国的开源 DFT 软件：ABACUS（原子算筹），帮助科研人员对凝聚态材料进行模拟计算，高效精准预测目标材料的各项物理性质。

在《新年展望》中，我们回顾了过去的业务合作，从手机到娱乐互联网、从智能驾驶到元宇宙，乃至医疗行业，都有 AI 的身影，也见证了 AI 在各行各业中发挥的巨大作用。

在众多的落地实践中，我们敏锐地发现一个趋势：AI 逐渐进入更关键的工业领域以及科学研究场景，并且发挥着硬核的作用。

因此，2023 年商汤 HPC 团队开启了新的发力方向 —— AI for Science (AI4S)。

从 Al 和 HPC 的视角看材料计算

材料计算领域里的第一性原理是指基于量子力学的、不需要经验参数（或只需要少量经验参数），只根据原子的基本信息（例如元素种类和位置等）就能较准确地预测物质性质的一套理论方法。

近年来，密度泛函理论（Density Functional Theory, DFT）在材料计算领域扮演着重要的角色。由于 DFT 较好地平衡了精度和效率，逐渐成为应用最为广泛的第一性原理方法；第一性原理软件常指代的就是 DFT 软件。近三十年来，在欧洲、美国、日本都有一批优秀的第一性原理软件出现，并且在新算法的发展上始终处于领先地位。尽管国内也有相关的研究工作，但总体来说依然落后不少。

作为计算机学科经典分支之一的高性能计算 (High performance computing, HPC）被誉为「科研第三支柱」，是现代科学进步的重要助力。HPC 能够比标准计算机更快地处理大量数据，使企业能够在使用高达 TB 级数据的场景下运行大型分析计算，例如天气预报、石油勘探、药物研发等。

而在传统科学计算领域之外，随着深度学习的崛起，AI 软硬件生态在计算机视觉、自然语言处理等领域都有了长足的发展。经过多年的沉淀，已经形成了一套成熟的工具链和高效的开发模式。

在 AI 的加持下，近几年中 HPC 也从传统方式拓展到全新的领域，诸如：等离子物理、生物计算以及材料计算等 AI4S 的领域。

商汤 HPC 团队在与 DeepModeling 社区深入合作时，也发现了 AI4S 在落地时可能遇到的挑战，并将分享我们实践过程中行之有效的方法。

AI4S 落地问题与挑战

AI4S 的落地面临着诸多问题与挑战，此次合作主要关注硬件与复杂负载的适配、软件生态整合与技术泛化这三个层面。

首先，AI4S workload 内在的复杂计算特性对硬件适配提出了极高的要求。

在硬件方面，异构加速，无论是 AI DSA，GPGPU，或是 CPU 新增的 AI 计算单元，已经成为新一代 AI 超算的标配。可以预见的是，国产 GPGPU/CPU 造芯热潮将会在未来提供更多的硬件选择。而在 AI4S workload 方面，由于传统科学计算的计算模式种类繁多（例如stencil、n-body、spmm等），其负载特性与 AI 计算核心（规整的conv、gemm）有一定差异，导致 AI4S 应用在硬件资源利用、计算/访存负载分配方面呈现出多样化的需求，这往往需要开发者同时理解算法特点，硬件特点，性能优化方法和软件架构设计知识。

因此，在多样的异构计算环境上，部署负载特性复杂的 AI4S 应用，成为了 AI4S 满足高性能、高可用落地标准的一大挑战。

其次，科学计算软件生态的历史积累，与 AI 蓬勃发展的软件栈工具链，在生态上仍未打通。

传统科学计算除了依赖通用的 HPC 计算库（例如 BLAS、LAPACK、FFTW 等标准计算库）外，经过学科发展不断演进，沉淀出了各自学术领域特有的计算软件，例如生物分子动力学软件包 GROMACS、计算流体力学 OPENFOAM、第一性原理 DFT 软件包 ABACUS 等。而面向深度神经网络的 AI 系统软件则包含训练框架（如PyTorch，TensorFlow）、推理框架（例如 TensorRT、OpenPPL）和前后处理算法库等部分。除了部分 BLAS 库基础软件存在共享之外，两者的软件栈生态相对独立。

因此，打通 AI 计算与科学计算的软件生态，共享软件研发环境，甚至统一中间表达，是提高 AI4S 软件研发效率的重要环节。

最后，不同学术领域的科学计算应用，往往围绕各自领域的算法乃至特定的硬件平台，进行高度定制化部署与优化。例如部署在天河二号与神威·太湖之光超算上的同一科学计算应用，由于硬件架构差异较大，其代码设计思路、优化方案不尽相同；即使在同一硬件平台之上，已有的科学应用优化方案往往也难以于另一应用中直接复用。这在一定程度上阻碍了AI4S 的技术推广。

因此，使 AI4S 应用间优化方法可迁移，实现技术泛化，从而降低 AI4S 研发、优化和部署工作的技术门槛，有利于将 AI4S 推广至更多科研领域。

针对这些问题，商汤科技 HPC 团队致力于为 AI4S 提供训（模型训练）、推（模型推理）、算（科学计算）一体的高性能跨平台计算框架。

目前，在「算」方面，商汤 HPC 团队为 DFT 软件 ABACUS 提供了高性能多核扩展性解决方案，作为探索 AI4S 高性能计算框架的第一步。

方法介绍

本次工作中，商汤 HPC 团队主要针对 ABACUS 在 LCAO（原子轨道线性组合）基组下的 OpenMP 扩展性进行深入优化，带来了计算效率更高、内存消耗更少的节点内多线程并行方案，为 MPI + OpenMP 混合方案的大规模扩展做准备。

一、瓶颈分析

优化开始之前，我们需要先分析程序的瓶颈。考虑到本次优化的主要任务是提高OpenMP环境下的性能，且ABACUS已有完善的MPI算法实现，我们首先通过对比MPI与OpenMP的并行计算效率来寻找多线程下的性能瓶颈。

关于并行计算效率，我们可以用以下公式来表达：

下图为 ABACUS 分别使用 MPI 进程并行和 OpenMP 线程并行计算「512 硅原子LCAO」算例的效率曲线。可以看到除了 psiToRho 对应的曲线之外，其他模块在 OpenMP 测试中的多核效率均不如 MPI。

经过代码依赖关系的分析，我们发现 cal_gint 模块并非顶层计算模块，它的耗时可以分散到其他模块中，其中 updateHk 模块的主要部分调用的是 cal_gint 模块。

在当前阶段，HPC 团队主要针对 diag 以及 ForceStress 模块进行扩展性提升，并初步优化了 psiToRho 模块。

二、求原子受力和晶胞应力的模块优化

1. OpenMP 并行化

基于深入的代码分析，ForceStress模块的主要计算部分均进行了OpenMP并行化。

主要方法包括：

a. 改进子模块，使其满足线程安全要求；

b. 解除循环中的稀疏索引依赖；

c. 解除多重循环中，内层循环次数需要在外层循环计算的依赖；

d. 采用动态调度，优化核心间的负载均衡。

除此之外，还引入了多种并行化工具。最终 ForceStress 模块 OpenMP 测试的并行效率基本持平甚至小幅超过原始 MPI 方案。

2. 数学函数优化

通过使用 Intel VTune Profiler 抓取的计算热点，我们发现，求力和应力函数 ForceStress 中，计算复指数函数时使用的 sincos 函数占据了大量的计算时间。考虑到当前测试环境中的标准数学函数库 libm 的版本比较老旧，我们尝试使用最新版本的libm 以及 Intel C++ Compiler 进行测试，并发现 sincos 的性能得到了大幅提升。

为了避免使用更老旧的系统环境的ABACUS用户因数学函数而拖累了性能，我们决定为ABACUS移植最新版本的sincos算法。

移植数学函数后，ForceStress模块部分子函数的性能表现如下图：

三、对角化求解模块优化

我们为GenELPA对角化算法引入了OpenMP支持。对比原始ScaLAPACK算法，GenELPA算法由访存瓶颈转向了计算瓶颈，算法复杂度更高，但可扩展性更优秀，在核心数量较多的大规模并行计算中可以凸显其优势。下图为1~16核MPI和OpenMP测试的相对效率曲线，以ScaLACPACK单核CPU时间为100%基准。

可以看出，ScaLAPACK 在 OpenMP 测试下，多核计算效率下降得非常快，且和MPI测试的效率拉开了显著差距；而GenELPA无论在MPI测试和OpenMP测试中均可以保持相对稳定的计算效率，核心数较多时更能体现其优势。

四、由电子波函数计算电子密度的模块优化

虽然在第一部分的并行效率评估中，效率曲线表明psiToRho模块的OpenMP多核扩展性已略微优于MPI，但在深入分析代码之后我们发现，psiToRho中有时间占用约为80s的子函数仍然存在OpenMP并行化提升空间。对此部分子函数进行优化后，psiToRho模块在OpenMP测试中可以得到比MPI测试更好的可扩展性。

总体优化效果

为了展示这些工作对ABACUS在LCAO基组计算的全方位性能提升，我们使用经过上述优化后的ABACUS再次对「512硅原子LCAO」算例进行测试，分别以各模块并行效率、总内存占用以及总运行时间为指标来评估多核性能的表现。

一、优化后各模块并行效率

二、MPI 与 OpenMP 方案总内存占用对比

三、优化前后总运行时间对比

同时，为了验证这些工作的泛用性和 OpenMP 对热点函数的覆盖率，我们额外测试了 10 个在生产场景中用到的算例在16核心下的总运行时间。测试结果显示，使用OpenMP在大部分算例下的性能都持平甚至超过了 MPI。需要注意的是，这些算例与「512 硅原子 LCAO」算例不同，默认使用的对角化算法为 GenELPA；因此，我们使用的基准版本为 GenELPA 对角化算法进行 OpenMP 支持后的版本。