基于神威太湖之光的宇宙学多体模拟

刘 旭,张曦煌,刘 钊,吕小敬,朱光辉

(1.江南大学 物联网工程学院,江苏 无锡 214122; 2.国家超级计算无锡中心,江苏 无锡 214072; 3.清华大学,北京 100084; 4.无锡航天江南数据系统科技有限公司,江苏 无锡 214000)

0 概述

近年来,科学家通过超级计算机模拟研究揭示了早期宇宙中星系与暗物质的形成过程。据估计,暗物质约占宇宙物质的85%,占总能量密度的1/4。由于暗物质对整个电磁光谱而言都是不可见的,因此需要在超级计算机上对暗物质演化模型的建立过程进行数值模拟。宇宙学模拟对于科学家研究非线性结构的形成以及暗物质、暗能量等假想形式具有重要作用,并且宇宙N体模拟已成为超级计算机的主要应用。

超级计算机作为大规模高分辨率模拟和应用的重要工具,在过去几十年发展迅速。在20年内,顶级超级计算机的峰值性能已经从每秒1012次浮点运算(TFLOPS)增长到每秒1015次浮点运算(PFLOPS),现在正朝着每秒1018次浮点运算(EFLOPS)迈进。这种快速的增长趋势很大程度上是源于使用GPU或多核芯片等加速器设备的异构架构兴起。我国的神威太湖之光超级计算机由40 960块自主研发的多核芯片SW26010组成,整个系统具有1 000多万个内核,峰值性能为125 PFLOPS。神威太湖之光强大的计算能力使其成为解决超大规模问题的理想平台。因此,自2016年神威太湖之光发布以来,已经在这台超级计算机上进行了大量的前沿研究和多个领域的应用。然而,至今没有任何软件能够在神威太湖之光上模拟具有数万亿个粒子的超大N体系统。

本文在神威太湖之光上基于PHotoNs[1]设计并研发一款高效N体模拟软件SwPHoToNs。针对粒子网格(Particle Mesh,PM)和快速多极子方法(Fast Multipole Method,FMM)的混合算法,提出多层次分解和负载均衡方案、执行树遍历和引力计算的流水线策略以及向量化引力计算和超越函数优化算法。通过以上并行算法设计和优化策略,实施多个大规模并行模拟实验,其中最大算例包含6 400亿个粒子,并扩展到神威太湖之光超级计算系统的5 200 000个计算核心上。

1 背景介绍

1.1 N体问题相关算法

在一个粒子动力学系统中,每个粒子在物理力(如引力)的影响下与其他所有粒子进行相互作用。因此,如果直接按照定义对系统进行模拟,对于每个粒子需要进行O(N)个操作,总计算复杂度为O(N2),其中N为系统中的粒子总数。直接引力模拟算法又称为粒子-粒子(Particle Particle,P2P)算法,具有很高的计算精度,但其随着计算规模的增大,计算效率会降低。

为此,学者们提出了多种以轻微降低计算精度为代价来提高计算效率的方法,其中树形法[2]得到广泛应用,而由BARNES和 HUT在1986年提出的Barnes-Hut分级树算法就是一种典型的树形法。在Barnes-Hut分级树算法中,根据牛顿万有引力定律,两个粒子之间的相互作用力会随着距离的增加而减小,在考虑远程一组粒子对单个粒子的作用时,不是简单地计算组内每一个粒子的作用,而是用它们的总质量和质量中心进行近似计算。树形法的基本思想是使用树形数据结构将所有粒子分为粒子集。在计算受力时,进行分级树的遍历,对节点和粒子的距离进行判定,若距离足够远则利用节点信息进行计算,否则继续细分进行比较,直至遍历到单个粒子。因此,可以极大地减少粒子之间的相互作用,将粒子间的计算复杂度降低至O(NlogaN)。

PM算法[3]通过对空间进行离散化进一步提高计算效率。但由于网格分辨率的限制,PM算法在建模近程力时存在不足,因此研究人员提出多种混合方法,如粒子-粒子-粒子网格(Particle-Particle/Particle-Mesh,P3M)算法和TreePM(Tree+Particle-Mesh)算法等,使用PM算法计算长程力,P2P算法或树形算法计算短程力。这些算法的时间复杂度一般为O(NlogaN)。

由于模拟系统中粒子总数已增长至数千亿甚至达到万亿,上述算法的时间复杂度已无法满足性能要求,因此文献[4]设计了在给定精度下时间复杂度为O(N)的FMM算法,其在某些方面与树形法相似,但其使用的是势而不是力。FMM算法通过层次划分和位势函数的多极子计算各点的位势,其计算精度和划分的层次有关,因此可以达到任意的计算精度。

1.2 N体模拟

尽管随着算法的改进,N体模拟的时间复杂度已降至O(NlogaN)甚至O(N),但从本质上而言,N体模拟具有统计性,这意味着N值越大,相空间的采样越精确和完整,从而可以得到更准确的结果。随着粒子规模的爆炸型增长,传统PC机的计算能力已经远远无法满足N体模拟的性能要求,因此需要大型并行计算机来解决该问题,而超级计算机正因为具有出色的计算和存储能力,所以大规模的N体仿真被移植到这些计算设备上。早期的超级计算机多数为同构架构,这意味着系统中只使用一种处理器或CPU核心。在当时的Caltech/JPL Mark III、Intel Touchstone Delta和Intel ASCI Red超级计算机上,WARREN和SALMON是第一批进行N体仿真的宇宙学家。1992年,他们用树形法模拟了一个包含1 715万个粒子的宇宙暗物质模型,获得了N体宇宙模拟领域的第一个戈登·贝尔奖[5]。在当时的Intel Touchstone Delta超级计算机上,模拟的持续性能达到5.1 GFLOPS～5.4 GFLOPS。为在Intel ASCI Red超级计算机上保持更好的负载平衡,WARREN和SALMON等人进一步改进了树代码,并进行一个包含3.22亿多个粒子的宇宙学模拟,模拟了大爆炸之后宇宙中物质的演化过程。仿真结果达到431 GFLOPS的持续性能,相当于超级计算机峰值性能的33%。在美洲豹超级计算机上,WARREN利用树形法进行暗物质晕的质量函数模拟[6],仿真包含690亿个粒子,单精度持续性能达到1.5 PFLOPS,相当于整个系统峰值性能的43%。ISHIYAMA等人报道了使用TreePM方法进行万亿个粒子的模拟,在K计算机上实现双精度4.45 PFLOPS的性能,达到理论峰值性能的42%[7]。

随着摩尔定律的终结,将配备了GPU、FPGA和GRAPE等加速器的超级计算机用于为科学和工程应用提供计算能力,这些计算系统被称为异构超级计算机。虽然面对仿真软件的复杂性、加速器与主机CPU之间数据传输的困难性以及在异构超级计算机上编程的挑战性等问题,但是应用程序可以从这些加速器中获得更好的高性能计算效果。GRAPE就是一种专门用于引力计算的加速器,MAKINO和FUKUSHIGE等人成功地进行了一系列基于GRAPE的天体物理N体模拟[8-11]。在过去的十几年中,由于GPU一直是顶级超级计算系统的主要组成部分,因此产生了大量面向GPU的N体问题解算代码。HAMADA等人在2009年利用树形法和FMM对256个GPU集群进行42 TFLOPS、包含16亿个粒子的N体模拟[12],仿真结果表明GPU集群在成本/性能、功耗/性能以及物理尺寸/性能方面都优于PC集群。2010年,HAMADA和NITADORI将代码移植到DEGIMA集群中[13],DEGIMA集群是由288个GTX295 GPU组成的PC机集群,可以用于更大规模的天体物理模拟。仿真包含30多亿个粒子,性能达到190 TFLOPS,计算效率为35%。BÉDORF等人于2014年使用树形法对包含2 420亿个粒子的银河系进行了单精度24.77 PFLOPS的N体模拟[14]。泰坦超级计算机在单精度下的理论峰值性能为73.2 PFLOPS,N体模拟的计算效率达到33.8%。

神威太湖之光超级计算机是第一个峰值性能超过100 PFLOPS的超级计算机,其将不同种类的内核集成到一个处理器中,具有处理进程管理、计算加速和内存访问操作的异构架构。虽然目前已在神威太湖之光上进行了许多大气、生物学、计算流体动力学等前沿模拟,但大规模的宇宙学模拟一直未见报道。IWASAWA等人[15]使用粒子模拟平台FDPS对行星环进行N体模拟,算例规模达到80亿个粒子,共扩展到4 096个Sunway节点。实测性能达到理论峰值性能的35%。本文设计的SwPHoToNs是为了利用整个超级计算系统来执行超大规模天体物理N体模拟而设计。在宇宙学模拟中,本文成功地将代码扩展到20 000多个Sunway节点,相当于可用计算资源的一半。

1.3 神威太湖之光超级系统

神威太湖之光是我国自主研发的超级计算系统,具有超过1 000万个计算核心,最高性能为125 PFLOPS,持续性能为93 PFLOPS。整个系统由40 960个SW26010处理器组成,这是上海高性能集成电路设计中心设计的一款具有260个异构内核的多核处理器,最高性能为3.06 TFLOPS,性能功率比为10 GFLOPS/W。整个处理器分为4个核心组(CG),每个核心组包含1个管理处理单元(MPE)、1个智能内存处理单元(IMPE)和64个计算处理单元(CPE)。SW26010体系结构如图1所示。

图1 SW26010体系结构Fig.1 Architecture of SW26010

MPE是一个完整的64位RISC内核,频率为1.45 GHz,主要用于处理程序的流量控制、I/O和通信功能。CPE采用更简化的微体系结构,能最大限度地聚合计算能力。每个CPE都有16 KB的L1指令缓存和64 KB的本地数据内存(Local Data Memory,LDM),可以配置为用户控制的快速缓冲区。基于三层结构(REG-LDM-MEM)的内存层次结构由FANG等人[16]提出。CPE可以直接访问带宽为8 GB/s的全局内存,也可通过REG-LDM-MEM内存层次结构获得更高的带宽。每个CPE有两个管道来处理指令的解码和执行。指令重排方案可以降低指令序列的依赖性,从而提高指令执行效率。在每个CPE集群中,一个包含8行通信总线和8列通信总线的网状网络能够在寄存器级上实现快速的数据通信和共享,从而允许CPE之间高效的数据重用和通信。

神威太湖之光采用基于高密度运算超级节点的网络拓扑结构。1个超级节点由256个处理器组成,1个超级节点内的所有处理器通过定制的网络交换机完全连接,实现高效的all-to-all通信,连接所有超级节点的网络拓扑是一个使用自定义高速互连网络的两级胖树,网络链路带宽为16 GB/s,对分带宽达到70 TB/s。

为适应SW26010处理器,本文开发神威太湖之光软件系统,包括定制的64位Linux操作系统内核、全局文件系统、编译器、作业管理系统等。Sunway编译系统支持C/C++、Fortran以及主流的并行编程标准(如MPI和OpenACC),提供了一个轻量级线程库Athread,并允许程序员执行更细粒度的优化。

为神威太湖之光设计的应用程序通常采用两级方法将不同算法映射到管理和计算核心,类似于Summit、Piz Daint、Titan等异构超级计算机,提供了两种使用CPE执行更细粒度优化的应用程序:1)Sunway OpenACC,其是从OpenACC 2.0标准扩展而来,可以在CPE集群上管理并行任务;2)轻量级线程库Athread,可以用于控制每个CPE。Sunway OpenACC利用CPE网格的快捷方式,但其缺少一些用于细粒度调优的重要特性,如寄存器通信、向量化等。相反地,使用Athread接口可以利用这些特性及多核处理器的计算能力。因此,本文采用“MPI+Athread”作为并行化解决方案。

2 SwPHoToNs算法描述及平台实现

2.1 SwPHoToNs框架

PHoToNs是针对大规模并行宇宙模拟[1]而设计,SwPHoToNs在此基础上针对Sunway平台进行设计和优化。PHoToNs软件将PM和FMM两种常用算法相结合,如图2所示。

图2 PM+FMM算法Fig.2 PM+FMM algorithm

本文使用PM算法用于远程引力计算,FMM算法用于短程引力计算,因此将P2P的方程修改为式(1):

(1)

其中,s为分离度,Rs为特征尺度。因此,短程力约在5Rs的范围外消失,修正后的方程引入了PM上的长距离高斯平滑,但却引入了两个计算非常耗时的误差函数(erf)和指数函数(exp)。

FMM算法简单而言是将大系统切成小立方体,只计算小立方体之间的作用力。因此,1万个粒子可能只切成10个立方体,从而大幅降低计算复杂度,并且采用树结构代码有利于快速计算和并行划分[17]。在FMM算法中,粒子被组织成K-D树,树的叶子是一个粒子包。树的每个节点存储2个多极矩M和L,将粒子间的力转化为多极矩之间的相互作用,其中计算引力的算符有P2M、M2M、M2L、L2L、L2P和P2P。

P2P算符是两个包中粒子之间的交互,是N体直接求解的核心部分。其他操作符则用于不同树节点的矩之间的交互。本文使用对偶树遍历来确定包和树节点之间的交互关系。对偶树遍历的优点是可以减少节点对的数量[19]。节点或包之间的打开角度(作为控制参数,由两者距离和包的宽度确定)为0.4,如果递归检查其子节点的交互关系,则此节点需满足该打开角度。

为提高编程的并行效果,本文将引力计算分解为一系列任务,对在CPE上处理的6个FMM算符中计算量最大的P2P和M2L构造一个任务池。

2.2 区域分解与并行优化

如图3所示,计算框由K-D树分解。每一个计算节点作为顶层树的一个节点,对应一个连续的空间区域和一个长方体,也是由局部粒子构成的局部树的根节点。因此,所有粒子都链接到一个完整的全局K-D树中。在图3(c)中,点划线圈内粒子间的相互作用由P2P处理,虚线圈内粒子间的相互作用由M2L处理。

图3 区域分解 Fig.3 Domain decomposition

在所有进程中记录顶层树,边界信息需要在进程之间进行通信,在此使用一个额外的遍历根据打开角度估计需要发送给其他进程的本地树节点数量。实际上,引力的长-短分裂保证了边界交换只发生在截止半径内。

在引力计算中,最大的工作量集中在两个部分,粒子间的引力作用P2P和M2L。由于这两个算符包含的计算量大致相同,因此算符的出现次数是负载计算单位。通过测量P2P和M2L的数量(M2L的数量相比P2P的数量少)来估计域之间的工作负载平衡性。

在将交互计算任务分配给每个MPE控制的CPE集群前,通过执行树遍历获取交互任务列表。交互任务在CPE之间均匀分布,以确保细粒度的负载平衡。因此,每个CPE都可以通过DMA在主存和LDM之间传输数据,并独立处理交互任务。

在树形法中,如果需要对质点进行引力计算,则需要遍历树获取交互信息。大多数GPU代码[13-14]是通过GPU进行遍历和引力计算过程,本文提出一种流水线策略,允许在MPE和CPE集群上同时执行遍历和引力计算过程。使用数组存储粒子包的ID,并在遍历过程中判断是否需要计算这些粒子包。一旦具有足够的交互任务(如8 192个包),就启动CPE集群进行引力计算,同时MPE会一直执行遍历过程,通过调整参数使遍历过程和引力计算处于平衡状态并且相互重叠。

根据牛顿万有引力定律,粒子之间的相互作用为all-to-all。因此,在力的计算过程中,目标进程需要与其他所有进程进行通信,这是影响N体求解器可扩展性的主要因素。为有效增强代码的可扩展性,将通信和FFT过程与P2P过程重叠。在MPE上处理通信和FFT时,驱动CPE集群进行P2P计算,此时不对遍历过程和引力计算进行相互掩盖,而是在任务池中存储所有的交互任务后驱动CPE集群,该策略具有较好的可扩展性。如图4所示,此时的掩盖策略与在本地进程进行引力计算时的掩盖遍历方式不同,因此针对两种计算情况构建了不同的任务池。

图4 针对两种情况的掩盖方式Fig.4 Cover-up modes of two situations

3 核心函数的众核优化

3.1 优化策略

如表1所示,P2P算符在整个仿真过程中所消耗的时间是最多的。P2P和P2P_ex函数的作用都是通过粒子之间的距离及粒子质量计算出粒子的加速情况,是天体模拟的核心引力计算函数。区别是两者使用的数据不同,P2P负责计算自身MPE中粒子的影响,P2P_ex负责计算其他MPE中的粒子的影响。因此,通过这两部分的从核计算进行性能优化。

表1 函数占用时间比较Table 1 Comparison of function occupancy time

在计算粒子间的相互作用时,由于力的作用是相互的,因此计算结果应累加到两者之上,但是该做法会导致所有粒子需记录已受过哪些粒子的作用,进行的判断与记录会耗费大量的时间和空间。因此,本文将计算结果只累加到一个粒子上,将计算中累加计算结果的一种粒子称为计算粒子,而另一种称为影响粒子。使用如图5所示的伪代码描述计算粒子包与影响粒子包之间的P2P计算和从核化过程。

图5 P2P函数伪代码Fig.5 Pseudo-code of P2P function

对于众核计算程序的优化,通过提高DMA带宽以降低DMA通信时间是一个关键策略,而另一个关键策略是将DMA通信时间隐藏在计算时间下,通过测试发现DMA通信时间比计算时间短,因此建立同等大小的数据缓冲池,将DMA通信时间进行隐藏,异步DMA方案的具体实现过程如图6所示。

图6 异步DMA方案Fig.6 Asynchronous DMA scheme

3.2 超越函数的优化

由于Sunway处理器在CPE上没有对超越函数进行向量化的指令,因此只能将向量存储到数组中,运用查表+多项式逼近的方法对exp函数和erf函数逐一进行求解。由于向量无法直接查表,因此针对exp函数的多项式(2)～式(4)和erf函数的定义式(5)、多项式(6)和多项式(7)展开,只采用多项式逼近的方法实现这些函数的SIMD版本,虽然会在erf函数的部分输入域时产生误差(最大误差为1.5×10-7),但加速效果较好。

exp(x)=2k×exp(r)

(2)

(3)

c(r)=r-(P1×r+P2×r+…+P5×r)

(4)

(5)

erf(x)≈1-(a1t+a2t2+…+a5t5)e-x2

(6)

(7)

其中,计算式(2)中的2k时,本文针对浮点型数据的存储形式,对于整型数据的指数k加偏差后位移,再保存为浮点型数据,即为2k的结果,以此避免了幂函数计算。

3.3 众核优化的加速效果

优化策略产生的加速效果与P2P算符的运行时间如图7所示,其中柱状表示一次迭代的运行时间,基准是一个MPE上对于执行P2P原始代码的运行时间。将计算任务分配给其控制的CPE集群,在解决了部分访址非连续问题后获得了38.27倍的加速。虽然理论上SIMD优化可将计算性能提升4倍,但是实际上除了加载和存储的开销外,exp函数和erf函数是计算中主要耗时的部分,约占总计算量的90%,且不存在相应的SIMD函数,只能将向量封装入数组中分别进行计算。因此,将函数循环展开并进行向量化后,只获得1.30倍的加速。在使用多项式逼近方法实现向量化超越函数后,既避免了向量的封装与加载,又加快了两个函数的计算。虽然产生了误差,但是误差大小是可接受的,且能通过泰勒公式更高阶的展开进行控制。经过函数计算方法的改进,获得了2.20倍的加速。将DMA通信时间隐藏后,获得了1.51倍的加速。最终实现了165.29倍的加速,在该规模的模拟中,对于P2P函数单次计算的时间从6 804.800 s减少到41.170 s。

图7 优化策略产生的加速效果与P2P算符的计算时间Fig.7 The acceleration effect produced by the optimizationstrageties and calculation time of P2P operator

4 实验结果与分析

4.1 宇宙学模拟效果分析

宇宙大规模结构是由初始的微小随机密度波动演化而来。根据Zel’dovich近似方法[20],通过粒子迁移建立仿真的初始条件。该方法是由初始功率谱编译卷积,早期可以用微扰理论进行估算。本文主要研究一个具有临界密度ΩM=0.25、Ω∧=0.75、σ8=0.8和哈勃常数参数H0=70 km/s的平面模型。在红移时用一个巨大的周期盒(z=49)进行模拟,该尺度可以满足宇宙学天空测量体积的统计量要求。具体而言,体积接近于8h-3Gpc3(其中,pc为秒差距,是天文学上的一种长度单位,h是以100 km/s/Mpc为单位的哈勃常数),从z=49模拟演变到z=0,形成的宇宙密度场如图8所示。本文共使用109个粒子参与到宇宙演变研究的实际模拟过程中,最大的基准测试算例包含6 400亿个粒子,总质量接近于太阳质量的108倍,这是首次在Sunway平台上实现同类宇宙学模拟。

图8 z为0且边长为2 Gpc时的宇宙密度场Fig.8 Cosmic density field when the z is 0 andthe side length is 2 Gpc

4.2 计算性能与可扩展性分析

由于程序执行过程中的绝大部分浮点操作都集中在P2P算符中,因此忽略程序中其他的浮点计算,仅统计P2P核心中的浮点运算次数作为宇宙模拟的总浮点计算次数。使用该方法计算出的总浮点计算次数略小于程序的真实浮点计算次数。在本文实现中,每次P2P过程中包含3次减法运算、12次乘法运算、6次加法运算、1次指数函数(exp)运算、1次误差函数(erf)运算和1次倒数平方根函数(rsqrt)运算。为便于与其他工作进行性能比较,本文对rsqrt函数使用文献[13-14]中提到的4次浮点运算,对exp和erf函数分别使用21次和15次浮点运算,因此1次P2P交互共产生61次浮点运算。

程序总浮点运算次数的计算公式为:

FFLOPS=I×C×SSteps

(8)

其中,I表示每次迭代计算中粒子间的相互作用(P2P)总次数,C表示每次P2P过程中包含的浮点运算次数,SSteps表示程序运行的总迭代步数。

程序持续浮点运算性能的计算公式具体如下:

PPerformance=FFLOPS/T

(9)

其中,T是程序运行的总时间。本文采用的最大算例中包含6 400亿个粒子,每次迭代计算中的粒子相互作用总数约为3.2×1016,每次P2P中包含61次浮点操作,总迭代步数为64,程序运行总时间约为4 243.2 s。将以上数据代入式(8)和式(9),计算得到程序的持续性能达到29.44 PFLOPS,计算效率为48.3%,并行效率约为85%,说明大部分通信时间与计算时间相重叠。

图9(a)、图9(b)为弱可扩展性性能与并行效率测试结果。对宇宙学天体演变测量中的统计量进行一系列以盒径为2h-1Gpc标定的宇宙学模拟。本文从1 024个节点开始,逐渐扩展到20 000个节点,每个节点包含4 000万个粒子(每个CG有1 000万个粒子),其中最大的粒子总数为6 400亿,实现了29.44 PFLOPS的持续计算速度和84.6%的并行效率。图9(c)为强可扩展性加速比测试结果,粒子数量为549亿,在20 000个节点上的并行效率达到85.6%。

图9 大规模实验性能测试结果Fig.9 Performance test results of large-scale experiment

4.3 计算性能对比

表2给出了HAMADA & NITADORI[16]、ISHIYAMA[7]、 BÉDORF[14]、FDPS[15]和SwPHoToNs大规模宇宙学N体模拟平台的性能比较,同时介绍了所有模拟平台使用的计算机名称、计算节点体系结构以及内存带宽浮点比率(B/F)。虽然SwPHoToNs硬件平台在B/F方面处于劣势,但相比其他平台具有更高的计算效率和浮点性能,并且本文对SwPHoToNs扩展到神威太湖之光全系统的计算结果进行了预估。

表2 大规模宇宙学N体模拟平台的性能比较Table 2 Performance comparison of large-scale cosmological N-body simulation platforms

5 结束语

本文设计了在神威太湖之光上进行大规模宇宙学N体模拟的软件SwPHoToNs。为将SwPHoToNs算法映射到神威太湖之光计算系统的5 200 000个计算核心中,提出多层次分解和负载均衡方案、树遍历和引力计算的流水线策略以及向量化引力计算算法等多种实现并行可扩展性和提高计算效率的技术,并在神威太湖之光上的5 200 000个计算核心上进行了宇宙学模拟,模拟算例包含6 400亿个粒子,弱可扩展性并行效率为84.6%,持续计算速度为29.44 PFLOPS。与BÉDORF[14]宇宙学N体模拟平台相比,SwPHoToNs将最大模拟算例中的粒子数提高了2.6倍,浮点计算速度由单精度24.77 PFLOPS提高到双精度29.44 PFLOPS,浮点计算效率高达48.3%。由于本文模拟工作建立在半机神威太湖之光系统上,因此还没有实现真正意义上的万亿级天体物理模拟,预估整机神威太湖之光系统能够支持包含多达1.2万亿个粒子的宇宙学模拟,持续计算速度将达到56.00 PFLOPS。