HPRMAT：核物理耦合通道问题的GPU加速高性能R矩阵求解器

《Computer Physics Communications》：HPRMAT: A high-performance R-matrix solver with GPU acceleration for coupled-channel problems in nuclear physics

　　

## 摘要

本文提出了HPRMAT，一个自包含的高性能R矩阵求解器框架，用于核物理中的耦合道散射计算。它提供了完整的R矩阵传播机制，具有与标准R矩阵包相同的使用者自定义势界面，并且可以作为Descouvemont包的线性代数例程的直接替代品。它采用优化的库进行直接线性方程求解，而非传统的矩阵求逆，实现了显著的性能提升。该包提供四种求解器后端：(1) 双精度LU分解，(2) 混合精度算术配合迭代细化，(3) 利用动能耦合矩阵结构的Woodbury公式方法，以及(4) GPU加速。基准计算表明，GPU求解器在N=25600时相比优化的CPU直接求解器实现约15倍加速，相比传统基于求逆的代码实现41倍加速。混合精度策略在消费级GPU（如NVIDIA RTX 3090/4090）上特别有效，其单精度吞吐量与双精度之比为64:1；通过在单精度下执行分解，并在需要时通过迭代细化恢复完整的双精度精度，HPRMAT克服了消费级硬件糟糕的FP64性能，同时保持了截面计算所需的精度。这使得大规模连续离散耦合道（CDCC）和耦合道计算对使用标准桌面工作站的研究者变得可行，无需昂贵的数据中心GPU。单张消费级显卡已能处理实践中遇到的最大规模问题，对于偶尔超出单卡容量的情况，已测试的多GPU后端将单精度分解分配至多个GPU，并通过主机端迭代细化恢复完整的双精度精度。纯CPU求解器通过优化的库和算法改进，对大矩阵提供约7倍的加速（取决于系统配置）。所有求解器的计算结果与参考结果在实验不确定度范围内高度吻合：双精度和CPU混合精度求解器在机器精度水平上一致，GPU单精度求解器在截面上达到10?3量级（可选迭代细化以提高精度），并已与Descouvemont的参考代码（Comput. Phys. Commun. 200, 199–219 (2016)）进行了验证。HPRMAT提供了Fortran、C、Python和Julia的接口。

## 程序摘要

**程序名称：** HPRMAT

**CPC库程序文件链接：** https://doi.org/10.17632/gcnpfnz3w6.1

**开发者仓库链接：** https://github.com/jinleiphys/HPRMAT

**许可证条款：** MIT许可证

**编程语言：** Fortran 90/95，CUDA C

**问题性质：** 使用Lagrange-Legendre基函数，通过R矩阵方法求解核散射的耦合道Schr?dinger方程。计算瓶颈在于求解由内部区域离散化产生的大型稠密复线性系统。

**求解方法：** 使用优化的LAPACK/cuSOLVER例程直接求解线性系统，并提供混合精度算术和利用动能耦合矩阵结构的Woodbury公式分解选项。

## 引言

R矩阵方法[1]在过去六十多年中一直是核反应理论的基石，为描述核物理中的共振现象和散射过程提供了强大的框架。通过在道半径a处将构型空间划分为内部和外部区域，该方法提供了R矩阵关于共振能量和约化宽度的一种自然极点展开，随后通过与通道半径处渐近解的匹配得到散射矩阵。

在微观R矩阵计算[2]、[3]中，内部波函数在一组完备的?2（平方可积）基函数中展开——通常是从Gauss-Legendre求积推导出的Lagrange-Legendre函数。这使得R矩阵方法成为一种类似于束缚态方法的连续谱问题处理方式[4]、[5]：散射问题在有限的内部区域中表述，在通道边界处添加Bloch算符确保内部区域哈密顿量的厄米性，同时保持与外部连续谱的耦合。散射问题的连续谱性质通过与通道半径处已知渐近解的匹配得以恢复。这种方法结合了束缚态技术的计算优势（有限矩阵对角化、变分原理）与描述散射和共振现象的能力。

虽然Numerov算法[6]等直接积分方法对于简单的单道散射问题在计算上更快，但微观R矩阵方法在耦合道计算中提供了独特优势。对于多道问题，R矩阵方法同时为所有道求解一个线性系统，而Numerov方法需要分别积分每个道并仔细处理道耦合。该方法还通过R矩阵的解析性质自然处理阈值尖点和近阈值行为，而Numerov积分在此处可能需要极细的步长。此外，?2基展开避免了Numerov积分在内部区域深度束缚或高度振荡波函数中可能出现的数值不稳定性。这些特点使微观R矩阵方法特别适合涉及阈值附近强道耦合的问题。需要注意的是，此处比较的对象是Numerov直接积分法。原子和分子散射中广泛使用的其他多道传播器，如对数导数传播器[7]和Light-Walker R矩阵传播器[8]，通过对角化或求解线性系统同时处理所有耦合道，并缓解了Numerov直接积分在阈值附近的某些稳定性问题。

耦合道方法的现代应用越来越多地需要处理许多道。例如，重离子融合反应中的转动和振动耦合[9]、[10]，以及弱束缚核的弹性散射和 breakup反应，后者用连续离散耦合道（CDCC）方法描述[11]。在此类计算中，耦合道数量很容易超过50，导致矩阵维度N=nch×nlag达到数千，其中nch为道的数量，nlag为每道所需的Lagrange基函数数量。

CDCC方法代表了一种特别具有挑战性的应用。为了描述弱束缚弹射体的 breakup连续谱，连续谱必须离散化为直到截止能量的"bin"态，每个bin贡献多个分波。晕核如11Be（以1?Be为核的单中子晕核）和?He（双中子晕核）由于其延展的空间分布和低 breakup阈值而特别具有挑战性。包含核激发的扩展CDCC（XCDCC）计算[12]尤为复杂。典型的11Be散射XCDCC计算[13]需要：(1) 每个弹射体自旋宇称Jπ从1/2 ±到15/2 ±的束缚态和连续谱bin；(2) 每个Jπ多个能量bin（通常低能时10-15个bin，高能时4-5个）；(3) 核-价核相对运动的分波高达lmax≈9。这可产生nch≈200-400个弹射体伪态。若收敛需要nlag=40-80个Lagrange基函数，则通常需要N=10,000-30,000的矩阵维度。这促使本工作聚焦于大矩阵性能。

R矩阵计算中的计算瓶颈是求解由此产生的复线性系统。传统实现，如广泛使用的Descouvemont代码[14]，采用矩阵求逆从哈密顿量矩阵获取R矩阵。虽然数学上等价于求解线性系统，但矩阵求逆有几个缺点：它需要计算和存储完整的逆矩阵，在数值上不如直接线性求解器稳定，且无法利用为线性系统优化的现代数值算法。

现有的若干R矩阵代码适用于不同领域。在原子和分子物理中，英国PRMAT代码[15]采用MPI并行化和ScaLAPACK进行分布式计算，UKRmol+套件[16]是分子与电子、正电子和光子碰撞的广泛使用实现；GPU加速也开始出现在该领域，例如最近关于H?O + H非弹性碰撞的R矩阵处理[17]。然而在核物理中，据我所知，没有公开可用的R矩阵代码利用GPU加速或现代混合精度技术。本工作填补了这一空白。

在本文中，我提出了HPRMAT（High-Performance R-MATrix），一个自包含的高性能R矩阵求解器框架，用于耦合道散射计算。HPRMAT提供完整的传播机制（Lagrange-Legendre网格、动能和Bloch算符、带库仑渐近解的碰撞矩阵构建以及波函数重构），连同四种可互换的求解器后端和Fortran、C、Python及Julia接口。它采用与Descouvemont标准参考包[14]相同的使用者自定义势界面：物理设置（势构建、道耦合、边界条件、必要的Buttle修正等）由用户提供，与上述包完全一致。作为额外的部署便利，HPRMAT与Descouvemont的包完全兼容，其现有用户无需修改应用程序代码即可采用新求解器。HPRMAT代表了当前核物理应用中公开可用代码在R矩阵求解器性能方面的最先进水平。在我的使用消费级NVIDIA RTX 3090 GPU的基准测试中，GPU求解器在N=25600时相比优化的CPU直接求解器实现约15倍加速（相比传统基于求逆的代码实现41倍加速）；在具有更高FP32吞吐量的更强大消费级GPU上，这一优势将更加显著。HPRMAT通过以下几项创新解决计算瓶颈：

1. **采用直接线性求解器。** 虽然直接求解Ax=b而非计算x=A?1b是数值线性代数中的标准做法，但许多传统核物理代码——包括Descouvemont广泛使用的包——仍采用显式矩阵求逆。HPRMAT实现了直接LU分解，将这一成熟的最佳实践引入R矩阵社区。

2. **优化的BLAS库集成。** 我使用了OpenBLAS[18]，一个高度优化的开源BLAS实现，利用多线程和SIMD向量化，在现代多核CPU上相比参考LAPACK提供显著加速。

3. **通过NVIDIA cuSOLVER实现GPU加速。** 我提供了GPU加速求解器，在消费级GPU上对大矩阵实现高达约20倍加速，优于优化的CPU求解器。

4. **混合精度算术。** 我利用了有利的FP32:FP64性能比（RTX 3090上为64:1），通过在单精度下执行LU分解并通过迭代细化保持双精度精度。

5. **Woodbury公式优化。** 我开发了一个利用耦合道哈密顿量矩阵特定块结构的求解器，其中对角块为满矩阵（动能），非对角块为对角矩阵（耦合势）。

6. **多语言接口。** 我提供了C、Python和Julia的绑定，便于无缝集成到超越Fortran的现代科学工作流中。

HPRMAT的一个关键设计原则是与Descouvemont的R矩阵包[14]的即插即用兼容性。子程序接口、调用约定和数据结构尽可能保持与原始实现一致。Descouvemont代码的用户只需替换相关源文件并重新链接即可切换到HPRMAT，无需修改现有应用程序代码。这一设计选择确保了性能提升可立即为现有用户社区所用。所有求解器均已使用Descouvemont包的全部五个标准测试用例与参考实现进行了验证：α+2??Pb光学模型散射、采用Reid软核势的核子-核子散射、1?O+??Ca耦合道散射、12C+α非弹性散射，以及非局域Yamaguchi势。

本文组织如下：第2节简要回顾R矩阵形式化和由此产生的线性系统的结构。第3节描述HPRMAT中实现的四种求解器算法。第4节讨论包括Fortran-CUDA接口和BLAS库集成在内的实现细节。第5节展示三种不同硬件平台上的基准测试结果。第6节总结发现并讨论未来方向。

## 各节节选

### 理论背景

R矩阵方法[2]、[3]通过将在道半径r=a处划分构型空间来解决散射问题的Schr?dinger方程，提供了一个框架。在内部区域（r < a），波函数在完备基集中展开。在外部区域（r > a），波函数用已知渐近解表示，并在边界处与内部解匹配。

对于具有nch个耦合道的系统，径向Schr?dinger方程取如下形式

### 数值方法

HPRMAT实现了四种求解器后端，允许用户根据硬件和精度需求选择最优算法。所有求解器以矩阵M和右端向量b为输入，返回解向量c。

CPU性能的一个关键因素是BLAS库的选择。HPRMAT中所有CPU求解器都设计为使用OpenBLAS[18]，一个优化的开源实现，提供多线程BLAS/LAPACK例程。与参考LAPACK库不同，

### 代码结构和API兼容性

HPRMAT使用Fortran 90/95实现，GPU支持部分使用CUDA C扩展。一个主要设计目标是保持与Descouvemont原始R矩阵包[14]的API兼容性。主子程序rmatrix保留与原始代码相同的接口签名、参数顺序和数组约定。这使得调用Descouvemont包的现有代码无需修改源代码即可使用HPRMAT——只需重新链接即可。

主要组件包括：

- **rmatrix_hp.F90**

### 性能结果

我在三种代表典型使用场景的硬件平台上对HPRMAT进行基准测试：高端桌面CPU（Apple M3 Ultra）、配有消费级GPU的工作站（Intel Xeon + RTX 3090）以及中端CPU系统（Intel i9-12900）。所有基准测试均使用与物理R矩阵问题具有相同块结构的合成测试矩阵。

这使求解器的性能得以独立评估，并允许矩阵维度在任何特定反应之外进行独立扫描；图4的物理验证结果确认了求解器的……

**总结与展望**

我提出了HPRMAT，这是一个自包含的高性能R矩阵求解器框架，用于耦合道散射计算。它提供了完整的传播机制，与标准R矩阵软件包一样，支持用户自定义势的接口，同时还与Descouvemont软件包中的线性代数例程即插即用兼容。据我所知，HPRMAT代表了核物理领域R矩阵求解器性能的当前最先进水平，是首个公开的……

**关于生成式人工智能及AI辅助技术在稿件撰写过程中的声明**

在本文准备过程中，作者使用了Claude（Anthropic公司出品）对稿件文本进行润色并调试程序代码。使用该工具后，作者对内容进行了审阅和编辑，并对已发表论文的全部内容承担全部责任。

**CRediT作者贡献声明**

金磊：写作——审阅与编辑，写作——初稿，可视化，验证，软件，资源，项目管理，方法论，调查研究，资金获取，形式化分析，概念化。

**利益冲突声明**

作者声明不存在任何可能影响本文所述工作的已知竞争性经济利益或个人关系。

金磊

同济大学物理科学与工程学院，中国上海，200092