格点量子色动力学蒸馏算法中关联函数的计算优化*

张仁强 蒋翔宇 俞炯弛 曾充 宫明 徐顺

1) (中国科学院高能物理研究所理论物理研究室, 北京 100049)

2) (浙江大学高分子科学与工程学系, 杭州 310058)

3) (浙江大学计算机科学与技术学院, 杭州 310058)

4) (中国科学院计算机网络信息中心, 北京 100190)

5) (中国科学院大学物理科学学院, 北京 100049)

格点量子色动力学(格点QCD)是一种以量子色动力学为基础, 被广泛应用于强相互作用相关计算的理论, 作为一种可以给出精确可靠理论结果的研究方法, 近年来随着计算机能力的提升, 正在发挥着越来越重要的作用.蒸馏算法是格点QCD中计算强子关联函数的一种重要数值方法, 可以提高所计算物理量的信噪比.但用它来构造关联函数时, 同样面临着数据量大和数据维数多的问题, 需要进一步提升计算效率.本文开发了一套利用蒸馏算法产生夸克双线性算符的关联函数的程序, 利用MPI (message passing interface,消息传递接口, https://www.open-mpi.org), OpenMP (open multi-processing, 共享存储并行) 和SIMD (single instruction multiple data, 单指令多数据流)多级别优化技术解决其中计算性能瓶颈问题.对程序进行了多方面的测试, 结果表明本文的设计方案能够支持大规模的计算, 在强扩展性测试下512个进程并行计算仍能达到70%左右的效率, 大大提升了计算关联函数的能力.

1 引 言

在自然界中有四种基本相互作用: 电磁相互作用、弱相互作用、强相互作用、引力相互作用, 它们决定了目前所知世界的物质运动规律.比如摩擦力实际上是电磁相互作用的宏观效果; 而强相互作用决定了强子的基本性质, 如强子的质量[1]、衰变宽度[2]、形状因子[3]等, 同时将质子和中子结合在一起组成原子核的核力, 实际上是剩余强相互作用,这种力决定了原子核的聚变和裂变以及质量等性质[4,5].因此研究强相互作用具有非常重要的意义.描述强相互作用的理论是量子色动力学(quantum chromodynamics, QCD)[6], 这种理论与电磁相互作用不同, 它在低能标是强耦合的, 这意味着无法用微扰展开的方式去近似计算强子的低能性质.目前也发展出了一些研究低能QCD的方法, 如手征微扰论[7]、大Nc展开法[8]等, 这些方法都对理论做了一定程度的修改, 都依赖于模型.而格点QCD则完全不需要任何的模型假设, 它直接从QCD出发, 利用蒙特卡罗方法进行可靠的数值模拟[9].

格点QCD计算应用需要借助大规模计算来实现组态的计算模拟, 需要巨大的高性能计算资源支撑, 其计算密集型和数据密集型特征对高性能计算技术的发展提出了不少挑战.格点QCD是未来美国E级计算机的主要应用之一[10], 格点QCD在我国E级超算的应用还处于探索阶段, 有些算法方面的移植优化工作[11].

在格点QCD中, 通常需要计算如下形式的量[12]:

其中, S [U] 表 示QCD的作用量, U (t,x) 是规范场,Oobserve代表任意可观测量.在欧几里得时空中, 可以将(1)式理解为简单的抽样模型

然后 用 蒙特 卡罗 方 法 去 模拟, 将 ρ (x) 理解 为概 率幅.只要产生一组满足 e-S[U]分布的规范场 U (t,x) ,任意一个可观测量 Oobserve就是在这一组规范场上的统计平均值, 并且可以有效地估计其误差.本工作研究的可观测量为夸克双线性算符的两点关联函数:

为了得到更好的统计精度, 近似地计算全传播子, 有人提出了蒸馏算法[16,17].这种算法在近似计算全传播子时需要的存储空间比传统方法小很多,还可以很方便地进行算符构造, 即便在结束传播子的计算之后, 仍然可以改变涂摩(smear)[18,19]波函数, 提供了算符构造[20]更多的自由空间.蒸馏算法把传播子的计算转换到拉普拉斯算符的本征空间.本征方程表示为其中

Nv表示采用本征矢量的个数, Nv越大, 对拉普拉斯算符的近似越好, 也就越接近于计算全传播子.f(λc) 是一个关于本征值的函数.在蒸馏算法中只需要保存本征值、本征矢量和约化传播子(perambulator):

另外, 蒸馏算法在进行算符构造时, 只需要对本征矢量的内积进行操作, 不需要去更改约化传播子.这意味着, 只需要计算一次约化传播子, 便可以重复地利用; 配合上对本征矢量的各种操作, 就可以构造各种各样的算符.这种灵活性和复用性是传统方法不具备的; 而且非连通部分的计算可以直接利用连通部分的中间结果, 不需要进行额外的计算.因此, 蒸馏算法具有很多传统方法无法企及的优势.

虽然蒸馏算法有很多优点, 但是组合约化传播子和拉普拉斯算符本征系统来构建关联函数过程的计算复杂度正比于将 Nv=10 和Nv=100 进 行 对 比, Nv=100 的 计 算 复 杂 度 是Nv=10 的10000倍.而越是精确的计算越是要求Nv要足够大, 因此并行计算以及相应的优化是很有必要的.我们根据蒸馏算法的特性, 开发实现了一套利用约化传播子和拉普拉斯算符本征系统计算关联函数的程序, 并进行了多方面的计算优化和测试.

本文将探讨蒸馏算法中利用约化传播子和拉普拉斯算符本征系统计算夸克双线性算符的关联函数的实现方法, 结合算法的计算特点提出了代码优化的设计方法, 以提高关联函数的计算效率.首先介绍通过蒸馏算法构造关联函数的基本原理; 其次介绍具体的程序优化方案; 在第3部分, 给出了测试结果和分析; 最后在第4部分进行总结.

2 核心算法

为了降低问题的复杂性, 以关联函数的连通部分为例, 解释如何利用蒸馏算法计算关联函数并且优化程序.非连通部分的关联函数的计算可以利用连通部分的计算中间结果得到, 因此只对连通部分进行讨论.

在格点QCD计算中, 通过蒸馏算法构造介子连通部分关联函数的形式[16]为

其中 c1,c2,c3,c4均为对拉普拉斯算符本征值和本征矢量序号的标记. Nv表示本征矢量的个数.α,β,γ,κ 均为物理学中的旋量指标, 它们的取值范围为 { 0,1,2,3} . A ,B 标记着物理上的不同算符, 每两个具有相同量子数的算符就可以用来构造关联函数, 它们会在不同的哈密顿量本征态上有不同的投影, 对应着不同的能量.记关联函数矩阵[21]为

Nop表示关联函数矩阵的维数.(6)式中Φ和τ的表达式为

x ,ω 均为空间指标.通常τ被称为约化传播子, 它是拉普拉斯算符本征空间的传播子, 对应传统算法的传播子. a ,b,c,d 标记的是物理学中的色指标, 它的取值范围是 { 0,1,2} .Φ的不同定义决定了算符的量子数, 对应到真实世界中的不同粒子.实际上,(7)式中的G就是传统方法中的传播子, 可以看到由于本征矢量与传播子的空间指标和色指标进行了求和收缩合并, 使得τ和Φ只剩下了表示本征矢量的指标、时间的指标和旋量的指标.注意e-σλc1这一项由格点QCD中的涂摩算法给出. λc1表示第 c1个本征矢量的本征值, 其值越大, 该本征矢量对关联函数的贡献就越小, 因此通常选取 Nv个最小的本征值对应子空间来近似计算全传播子.

通过蒸馏算法计算关联函数的步骤如图1, 首先解拉普拉斯算符的本征值和本征矢量; 再利用本征矢量, 计算τ, 并保存以便重复使用; 根据研究的需要利用本征矢量和本征值构造Φ; 最后按照(6)式, 利用Φ和τ计算相应的关联函数.本文工作主要涉及后两步.计算λ, v, τ的过程只需进行一次.在后续的计算中根据不同的研究需要, 对本征矢v进行相应的操作再与约化传播子τ进行缩并可以构造出相应的关联函数.重复利用λ, v, τ.相较于传统方法在每次进行算符构造时都需要进行传播子求解, 蒸馏算法构造算符更加方便, 在有了约化传播子和本征值、本征矢量后, 计算关联函数更经济.同时, 由于产生约化传播子的过程与构造算符的过程相对独立, 蒸馏算法可以在研究课题没有确定的情况下, 不断利用可用的计算资源产生约化传播子和本征值、本征矢量, 合理利用资源.

图1 利用蒸馏算法计算关联函数的流程Fig.1.The procedure of computing correlators via distillation method.

(8)式中Φ的计算实际上可以分解为两部分,Φ=ΦC⊗ΦG. ΦC代表拉普拉斯本征空间的计算,ΦG代表旋量空间相关的计算. ΦC和 ΦG的计算是互相独立的, 并且 ΦG是一些常数矩阵的组合, 不需要额外进行计算.在程序计算时, 将 ΦG固定在程序内部, 可变的 ΦC由外部文件读入, 然后在根据需要将 ΦC和 ΦG通过一定方式组合起来得到所需的具有确定量子数的Φ.

本征矢量的个数 Nv的取值影响着物理信号,通常需要 Nv足够大以便于对尽可能多的情况保持良好的信号, 而蒸馏算法从约化传播子计算关联函数的复杂度虽然相比传统方法计算一个全传播子的关联函数的计算量要小很多, 但仍然需要处理次乘法.通常 Nop的取值是 O (10) 级别, Nv是 O (100) 级 别, Nt也是 O(100)级别.

为加速计算, 在进程级别、线程级别和指令集级别分别进行了算法的计算优化.

3 优化方案

对(6)式进行计算约化, 以减少不必要的计算开销.结合蒸馏算法计算的特性, 选择时间维度上的切片操作, 切割后的数据处理相对独立, 使用MPI多线程方式实现时, 避免了数据通信成为瓶颈的可能; 另外高维约化传播子和Φ的数据被划分到多个计算节点上处理, 也避免了单个计算节点的内存容量限制.为了进一步提高计算并行度, 在MPI进程中, 实现了基于共享内存特性的线程级并行计算, 最大可能地提升了计算节点中资源的利用率.最后, 由于涉及到大量复杂的复数运算操作,开展了数据结构的矢量化适配, 实现了SIMD指令级的并行计算.

3.1 计算约化

关联函数计算表达式(6)式变形成

注意到括号中的 c2和β指标下的部分计算用T矩阵表示, 那么关联函数可以写成

在计算关联函数之前, 可以先计算出T矩阵, 然后利用它来计算不同的关联函数.因为T的计算量∝Nop×Nv, 于是关联函数的计算量可见计算复杂度比原来小很多, 在 Nv和 Nop越大时越明显.实现计算约化后的程序流程如图2.

图2 含计算约化的关联函数计算的流程.其中 T A 和TB 表示两个中间计算量.利用中间量的计算减少了总体的计算量, 让计算量从 变成 极大地减少了计算量Fig.2.The flowchart of computing correlation function.TA and T B are two intermidiate quantities.After introducted intermediate quantities, the computation consumption is highly reducted to

3.2 进程和线程级优化

注意到关联函数的计算关于各个时间片t和 t′独立, 因此选择对时间维度进行切分, 这样就能够在各个进程进行独立的计算, 完成各个t, t′所负责的那部分数据进行收缩合并, 无需考虑通信开销.程序实现是选择MPI并行计算方式, 每个MPI进程负责各个独立的时间片, 在MPI进程中进一步启用OpenMP线程共享内存的并行计算, OpenMP多线程没有跨计算节点的通信, 因此不会增加节点间通信.

在程序实现时, 选择依次对t和 t′切分.先对t进行切分, 当进程数大于 Nt时再考虑对 t′进行切分.以 Nt=128 为例, 如果进程数等于64, 那么第一个进程就要计算 t =0,1 , t′=0,1,···,127 的关联函数, 第二个进程就要计算 t =2,3 , t′=0,1,···,127 的关联函数的数据.如果进程数等于256, 那么第一个进程就要计算 t =0 , t′=0,1,···,63 的关联函数, 第二个进程就要计算 t =0 , t′=64,65,···,127 的数据, 以此类推.这样进行划分的另外一个优势是, 由于关联函数具有随 t′指数衰减的性质,以及对于t具有平移不变性, 可以利用这两个性质来检查计算结果的正确性.

由于约化传播子τ以及Φ的维度 ∝ Nv2, 当 Nv的取值接近1000, Nt=128 时, 按双精度存储,τ的数据量大约1 TB, 因此在MPI实现时使用并行IO的数据读入方式, 并行IO也是按时间片划分方式, 和MPI并行计算方案保持一致.注意到Φ只有一个时间指标, 而τ有两个时间指标.在实际计算中, 当并行进程数 Np≤Nt时, 只对τ其中一个时间指标进行切分, 分配到各个进程上去, 而Φ不切分; 当并行进程数 Np＞Nt时, 同时对τ的两个时间维度进行切分, 而切片τ的计算只依赖部分Φ, 即每个并行进程上都只需Φ的部分矩阵元,因此Φ切片依赖于τ的切分方式.其流程如图3所示.

图3 按照时间切分实现并行计算的方式, 根据 N p 与 Nt的相对大小, 由于数据的特性, 对τ和Φ按情况采用不同的切分方法.Fig.3.Data sgemented according to time.Two conditions are considered which decided how τ and Φ are treated because of the feature of data.

3.3 指令级优化

在计算中, 考虑到形如

的张量缩并过程, 可以将其简化为与缩并指标相对应的矩阵乘法.上述的两个表达式都可以看作形如(4Nv,4Nv) 的矩阵乘法.如将原本按照 (4,4,Nv,Nv)行优先存储的数据转换为 ( 4Nv,4Nv) 行优先存储的格式, 那么可得如下的计算方式:

行优先存储格式下按(13)式作矩阵乘法时, B矩阵无法实现连续的内存访问, 因此将B矩阵转置, 计算变为

则计算中在循环k时可以内存连续地访问矩阵的元素.

在计算中, 每个数字实际上是复数, 而由于复数乘法规则复杂, 在不更改数据排列的情况下,编译器无法保证利用SIMD矢量运算提升复数计算性能.因此将矩阵的实部和虚部分离存储,利用运算法则将1个复矩阵的乘法分解为4个实矩阵乘法和2个实矩阵加法, 形式为 Re(AB)=Re(A)Re(B)-Im(A)Im(B) , Im(AB)=Re(A)×Im(B)+Im(A)Re(B) .

在上述调整下, 变量的内存访问具有连续性,有助于编译器自动完成SIMD矢量计算优化.通过确认编译得到的汇编代码, 编译器确实启用SIMD矢量运算指令.当然, 也根据测试平台手动调用SIMD相关函数和指令实现了计算, 经过测试相比编译器自动优化提升不到10%, 考虑到硬编码的繁琐之处和兼容问题, 我们最终采用的方式是修改数据格式和实现复数计算后, 通过编译器自动优化完成SIMD向量化.

4 测试结果

为测试 加速效果, 在 Nt=128 , Ns=16 的格子上进行测试.组态的海里只包含两味粲夸克.计算过程中采用双精度. a-t1=9.6Gev .分别对SIMD,MPI, OpenMP的加速效果进行测试, 并对物理计算的结果进行展示.

4.1 SIMD计算加速测试

为了验证SIMD计算加速效果, 选择一个含有Nv=10 个本征向量的算例, 按照(9)式来计算关联函数矩阵( 5 ×5 维).测试的处理器支持FMA和AVX SIMD指令集.该CPU处理器包含8个物理核心, 可通过超线程技术实现的16进程并行.计算时对整个流程按步骤分别计时, 具体步骤包括: 初始化、数据IO读取、计算T中间量、利用T矩阵计算关联函数.为了实现使用SIMD优化效果, 还对照测试了直接调用标准库中的复数计算函数的程序设计.按步骤计时的结果如图4, 使用SIMD前后计算性能如图5.

图4 使用SIMD优化前后各阶段耗时对比.I/O代表图1中第一步和第二步的时间, Calc.prepare代表图1中第三步的时间, Calc.result代表图1中第四步的时间, Others代表图1中第五步的时间, Init代表程序初始化的时间.图例SIMD表示启用了AVX形式的SIMD计算性能, 而Complex表示程序直接调用标准库中的复数计算函数(此处未使用SIMD计算).其中16个MPI进程并行计算的结果是在超线程计算状态下获得Fig.4.The cost of time of program's each part to see the effects of SIMD.I/O labels the time of first step and second step in Fig.1, Calc.prepare labels the time of the third step in Fig.1, Calc.result labels the time of the fourth step in Fig.1, Others labels the time of the fifth step in Fig.1.Init labels the time of initialization.SIMD in the picture means SIMD optimization was adopted and Complex in the picture means the stdandard library of complex computation was used.And hyper-threading technology was used for 16 MPI process.

图5 使用SIMD优化前后性能对比.图例SIMD表示启用了AVX形式的SIMD计算性能, 而Complex表示程序直接调用标准库中的复数计算函数(此处未使用SIMD计算).其中, 在SIMD启用时16个超线程计算结果未参与数据拟合Fig.5.The cost of time of program's each part to see the effects of SIMD.SIMD in the picture means SIMD optimization was adopted and Complex in the picture means the stdandard library of complex computation was used.And hyper-threading technology was used for 16 MPI process.

程序在启用SIMD指令后, 在不考虑超线程的时, 运行时间变为原来的1/4.Calc.prepare的时间变为原来的1/6.Calc.prepare是程序的热点, 此处计算程序中间量, 计算时间随Nv的变化最明显.

4.2 OpenMP计算加速测试

现测试OpenMP计算加速效果, 沿用了上一节测试的算例, 处理器开启CPU的超线程, 并关闭SIMD加速.同样按步骤计时的测试结果如图6.

图6 使用OpenMP优化前后耗时对比.图例如图4.图例Serial表示串行版本, 即未开启OpenMP多线程和MPI多进程Fig.6.The effects of OpenMP optimization was showed.Legends are the same as 4.Serial lables the results of serial program which means no OpenMP and MPI was adopted.

从结果可见, 程序采用OpenMP线程加速方式有性能提升, 使用16个超线程加速相对单核获得了约10倍的加速.同时采用OpenMP线程的计算效率与采用MPI进程的计算效率几乎相同.

4.3 MPI强扩展性测试

为了测试MPI并行计算的强扩展性, 选择一个含有 Nv=100 个本征向量的大算例, 此时输入文件体积达到40 GB, 采用不同的计算核数分别按照(9)式来计算关联函数矩阵( 5 ×5 维).测试的处理器支持FMA和AVX SIMD指令集, 最大MPI进程达到512个.同样按初始化、数据IO读取、计算T中间量、利用T矩阵计算关联函数等步骤计时.结果如图7.MPI并行的效果明显, 适合大规模并行.MPI进程数的增加, 效果体现在Calc.prepare部分, 也就是(11)式中间量的计算加速效果最明显.由于各部分加速的效果不同, 当512个进程时,Calc.prepare的计算时间与I/O的时间差别不大.I/O随进程数的变化时间改变并不明显, 耗时可以忽略.

图7 MPI并行强扩展性测试.随着MPI进程数增加, 计算时间成比例减少.图例如图4Fig.7.MPI parallelism in strong scale tests.The cost time decrease with MPI process numbers.Legends are the same as Fig.4.

计算程序在强扩展性测试中的并行计算效率如图8, 从结果可见, 总体并行计算效率很高, 在512个进程的情况下仍能达到70%左右的效率.因此程序适合大规模计算加速.

图8 MPI并行强扩展性测试，不同MPI进程数的测试相对于16个进程的计算效率Fig.8.MPI parallelism in strong scale tests.The efficiency of strong scale tests compared with 16 MPI processes.

4.4 MPI弱扩展性测试

为了测试MPI并行计算的弱扩展性, 选择几个不同规模的算例, 其进程数分别为 Np=2,16,128,1024 , 关联函数矩阵维度仍然为 5 ×5 .测试的处理器支持FMA和AVX SIMD指令集, 对不同算例等比例使用不同规模的并行进程数, 使得理论上每个进程的负载相同.按步骤计时的结果如图9,忽略I/O部分, 其他部分表现良好.但I/O部分当进程数等于1024时, 耗时出现了明显的增加, 并与热点部分(Calc.prepare)的耗时可比.这受限于硬件的性能.在当前测试中, 当 Np≤128 时程序在弱扩展性测试中的并行效率较高, 但 Np=1024 的大规模测试中展现出了一个因系统资源限制产生的问题, 由于磁盘带宽是一定的, I/O请求数激增, I/O部分时间大大增加, 甚至超过计算中间变量T矩阵(Calc.prepare)部分成为了新的热点; 但并行I/O能够解决串行读入数据时内存不足的问题.由于计算量是随着 Nv的三次方增加的, 而数据量大小是随着 Nv平方增加的, 这会导致如果要维持每个进程的计算量不变相应的数据量就会变少, 读入的数据会随着 Nv的增加而变得零碎, 最终受限于I/O.

图9 MPI并行弱扩展性测试. N p 表示使用的并行进程数, N v 表示本征向量数.图例说明同图4Fig.9.MPI parallelism in weak scale tests. N p represents the number of process, N v represents the number of eigenvectors.Legends are the same as Fig.4.

4.5 物理结果

为了检验结果的正确性, 对物理结果进行了测试.构造了三个JPC量子数为 0-+的算符, 分别计算它们的关联函数并分析有效质量,这里 C (t) 表示某种量子数算符的关联函数.通过比较这三个算符所得的关联函数的有效质量, 在时间片比较大的时候趋于同一个平台, 与传统方法所得结果一致, 见图10.

图10 做变分前的结果.在时间比较小时, 三个算符得到的有效质量的行为有很大差别, 证明在不同态的投影是不同的, 意味着变分会有一定的效果.在时间较大时, 三个算符的有效质量趋于同一平台, 说明它们的量子数是相同的,可以用来变分.traditional method表示第一个算符通过传统方法所得到的有效质量, 作为蒸馏算法的参照Fig.10.Results before variation.The behaviors of the effective mass of these three operators are very different and it means variational analysis would give good results.When time is large enough, these three operators approach to the same plateau so that they should have the same quantum numbers.traditial method label the effective mass of first operator throgh traditional mehtod, which can is matched with distillation method.

因为三个算符在不同本征能量态上的投影有很大不同, 如果中间计算没有错误, 通过变分分析[21]可以近似得到三个主要投影到最低的三个能量本征态的算符, 变分方程如下:

变分后, 可以得到本征矢 Vi, 对应第i个态 mi,将 CA,B与 Vi组合, 可以得到对应到相应态的关联函数

并且利用这三个算符所计算出的有效质量要能够相互区别开, 在时间片比较大的时候趋于不同的平台.因此可以此来验证计算的正确性.

利用上述三个算符构造关联函数矩阵, 做变分分析, 得到如图11所示的结果.变分后得到三个不同的态, 分别对应着基态、第一激发态、第二激发态.利用变分本征矢量组合后的关联函数所得的有效质量在时间较大时, 分别趋于不同的平台, 证明变分将原来的三个算符分别组合成了三个主要投影到最低三个态的优化算符.第二激发态的信号略差.变分效果明显, 再次验证了结果的正确性.

图11 做变分后的结果Fig.11.Results after variation.

5 结论与展望

本文围绕格点QCD蒸馏算法中关联函数的计算优化问题, 提出了一套加速程序计算的有效解决方案, 结合理论分析采用了MPI, OpenMP和SIMD多级别的计算优化技术.文中首先介绍了蒸馏算法中关联函数计算中各个模块的优化设计, 并对程序的性能进行了多方面的分析和测试.实验结果表明蒸馏算法计算关联函数效率的显著提升, 并且验证了物理结果的正确性.