适用于结构化网格应用程序的OPS性能便携式伴随算子
《ACM Transactions on Mathematical Software》:Performance Portable Adjoints for Structured Mesh Applications with OPS
【字体:
大
中
小
】
时间:2026年09月11日
来源:ACM Transactions on Mathematical Software 2.0
编辑推荐:
摘要
要查看此由AI生成的摘要,您必须具有高级访问权限。了解更多信息请登录。
**引言**
导数在工程和科学应用中至关重要,例如优化问题和逆问题。尽管有限差分近似可以估计导数,但它们计算成本高且不够准确。算法微分(AD)通过将计算机程序视为数学函数并应用微积分的链式
摘要
要查看此由AI生成的摘要,您必须具有高级访问权限。了解更多信息请登录。
**引言**
导数在工程和科学应用中至关重要,例如优化问题和逆问题。尽管有限差分近似可以估计导数,但它们计算成本高且不够准确。算法微分(AD)通过将计算机程序视为数学函数并应用微积分的链式法则,提供了一种高效且精确的导数计算方法。我们的工作重点是针对结构化网格模板应用的伴随模式AD。我们扩展了Oxford Parallel库(用于结构化网格求解器OPS)的领域特定语言,以使用反向模式的AD来计算导数。OPS使开发人员能够从同一源代码出发,针对多种硬件表达网格算法。利用领域特定的抽象,该扩展在计算循环级别创建了一个紧凑的伴随式磁带,并使用用户提供的原始和伴随模板内核生成特定于平台的(OpenMP和CUDA)并行实现。我们对用OPS编写的三个示例应用进行了微分,证明其性能与原始应用相当,无论是CPU还是GPU。在这些应用中,计算导数所需的时间仅为原始应用的3.7到9.7倍(包括应用程序的评估时间),这与最先进的工具相符。
**AI摘要(实验性)**
此摘要是使用自动化工具生成的,并非由文章作者编写或审核的。它旨在帮助发现相关内容、帮助读者评估相关性,并协助来自相关研究领域的读者理解这项工作。它旨在补充作者提供的摘要,后者仍然是论文的主要摘要。完整文章才是权威版本。点击此处了解更多信息。点击此处评论此摘要的准确性、清晰度和实用性。这样做将有助于改进未来的版本。要查看此由AI生成的简单语言摘要,您必须具有高级访问权限。
**敏感性信息**
敏感性(数学导数)信息有广泛的应用,例如计算流体动力学(CFD)应用中的基于梯度的设计优化 [Albring等,2016;Bischof等,1992;Carle等,1994;Giles和Pierce,2000;Sanchez等,2018]、机器学习 [Baydin等,2018]、逆问题 [Guasch等,2020]、计算金融 [Homescu,2011]、图像处理 [Li等,2018]、医学成像 [Grabner等,2008] 或实时风险管理 [Capriotti和Lee,2018]。算法微分(AD)是获取此类应用中敏感性信息的重要工具:有限差分通常计算成本过高(或不准确),而导出符号梯度信息可能会导致代码冗长且容易出错 [Laue,2022]。AD,也称为自动微分,是一组技术,可以通过反复应用微积分的链式法则来计算函数对其输入变量的精确导数。假设我们有一个具有某些输入变量x∈Rn的应用程序,该程序执行一系列步骤后计算最终结果y∈Rm。如果我们将此程序视为数学函数F:Rn→Rm,AD将计算xˉ=dF/dx·yˉ。手动编写代码计算解析导数可能是最快的解决方案,但它容易出错,并且对于大型应用来说根本不可行。最常见的有三种自动计算导数信息的方法 [Margossian,2019]:计算机代数系统中的符号微分、使用有限差分的数值微分以及AD [Bücker等,2006;Naumann,2011]。符号微分通过操作符号表达式来计算导数,它可以提供精确的导数,但可能较慢,并且需要问题以封闭形式表达 [Baydin等,2018],这在许多情况下是不可能的。使用有限差分估计导数有时最容易实现,但导数受浮点精度的限制,获得精确的二阶或三阶信息可能是一个真正的挑战。最后,AD通过微分组成用户应用程序的基本算术运算和基本函数来计算导数,从而确保准确性。AD相对于有限差分的另一个优势是,伴随模式AD的计算成本与输出数量成比例,而不是输入数量;因此,在应用程序产生的输出远少于输入的情况下,伴随模式AD可以快数百倍地计算梯度。
**AD的工作原理**
AD基于这样一个事实:无论计算机程序多么复杂,它总是由一系列基本操作(如加法、乘法或sin等数学函数)构建而成。这些基本操作的导数是已知且易于计算的。如果将函数视为由K个基本步骤组成的复合函数 F=f1°f2°?°fK,则 y=F(x)=fK(…f2(f1(x))),并且可以通过对所有基本函数应用链式法则来计算y的导数 ?y/?x=?fK/?fK?1?fK?1?fK?2??f1/?x (1)。
**AD的模式**
根据我们从哪个方向开始评估方程(1),AD主要有两种模式:前向(切线)模式AD,其中评估顺序遵循原始计算的顺序;以及反向(伴随)模式AD,其中按相反顺序评估链式法则。设J表示F的m×n雅可比矩阵,jk表示基本函数fk的雅可比矩阵,u∈Rn表示输入空间中的向量。切线或前向模式AD将计算雅可比矩阵对u的作用:Ju=JKJk?1?J1u。前向模式AD按原始函数评估的相同顺序计算导数(如果F在计算机程序中实现,这与程序执行语句的顺序相同)。选择u使得其在输入j处为1,在其他位置为0,单次评估切线模型将计算雅可比矩阵的第j列,即所有输出相对于单个输入的导数。评估切线模型需要两倍的内存,并且通常比原始应用花费大约两倍的时间。然而,通常需要计算多个输入的导数。
**伴随模式AD**
对于伴随模式AD,设w∈Rm,表示输出空间中的向量。伴随模式AD将计算J的转置对w的作用:JTw=J1TJ2T?JKTw。与前一种情况类似,选择w使得其在位置i处为1,其他所有值为0。单次评估伴随模型将产生J的第i行,即单个输出相对于所有输入的导数。伴随模式AD按相反顺序评估导数,相当于反向运行计算机程序。这需要在内存中存储程序控制流的有向无环图(DAG),以及局部导数信息,或者多次重新计算DAG的部分。这可能会导致内存需求大幅增加。记录中间状态和计算图的数据结构通常称为“磁带”。主要挑战之一是在计算程序结果的同时高效地保存所有中间状态和计算步骤(及其局部雅可比矩阵),以便将导数信息从输出传播回输入。
**选择AD模式**
选择前向模式还是伴随模式AD通常取决于目标应用。前向模式对于输入较少而输出较多的应用表现更好,而伴随模式AD在函数具有大量输入但相对较少输出的情况下表现出色,例如基于梯度的优化和机器学习。由于其在这些场景中的高效性,伴随模式AD经常受到特别关注。应用领域包括形状优化任务 [?zkaya和Gauer,2010;Verstraete等,2017;Vitale等,2020]、晶体管建模 [Phipps等,2008],或在CFD求解器中,如NASA冰盖系统模型 [Larour等,2016]、OpenFOAM [Towara和Naumann,2013;Towara等,2015] 或STAMPS [Müller等,2018]。大多数伴随模式AD工具属于两类之一:要么应用源转换来生成函数的伴随值,要么使用操作符重载技术 [Naumann,2011]。源转换工具自动转换原始代码以生成计算伴随值的新程序,而操作符重载工具利用C++和其他面向对象语言的面向对象功能来重载基本算术运算和函数以计算导数。这允许灵活地跟踪复杂的控制流和不同的语言结构。存在多个高效使用操作符重载的工具,如Adept [Hogan,2014]、dco/c++ [Lotz,2016] 或Sacado [Gay,2006;Phipps等,2008]。然而,天真实现的操作符重载工具会导致较大的内存开销,用于存储磁带,这通常可以通过高级技术来缓解,例如CoDiPack中的表达式模板 [Sagebaum等,2018,2019]、在MATLAB中使用ADIMat中的线性代数结构的直接导数 [Bischof等,2002] 或Stan Math库 [Carpenter等,2015],以及在dco/c++的情况下使用额外的混合代码生成扩展。虽然使用操作符重载跟踪原始数据流较为简单,但生成的磁带大小会迅速增加。源转换工具可以减少伴随计算的性能开销,并可能允许更彻底的优化,但代价是更复杂的工具。具有源转换的伴随模式AD工具的示例包括ADIC [Narayanan等,2010]、Tapenade [Hascoet和Pascual,2013] 和OpenAD [Utke等,2008]。总之,AD提供了一种有效且精确的导数计算方法,特别是对于输入变量众多而输出数量较少的情况。我们的方法使用源转换技术,允许我们在较高层次(整个并行循环的层次)存储控制流,同时能够对循环进行转换和优化,以实现并行硬件上的接近峰值性能。
**实现导数**
在并行环境中实现应用程序的导数既困难又容易出错 [Margossian,2019]。理想情况下,用户希望能够用简单的术语描述需要计算的内容(包括任何导数信息),然后有一个工具或框架来生成高效的并行实现,包括选择最适合获取导数的方法。在AD工具中,通过多种方式实现导数计算的并行化。对于原本串行的应用程序,一种相对简单的方法是并行计算多个独立的前向模式AD计算 [Bücker等,2001,2002],或者使用嵌套并行性 [Bücker等,2004;Phipps等,2022]。对于伴随模式AD,存在固有的复杂性。由于数据流被反转,原始代码中的线程安全并发读取在伴随代码中变成了线程不安全的并发写入。高效处理这些竞争条件是一个关键问题。在Heimbach等 [2002] 的研究中表明,结合源转换工具和手写的伴随例程,可以为应用程序获得高效的并行实现。另一种方法是使用与工具无关的扩展库来处理并行性,例如AMPI库 [Schanen等,2010,2012],用于处理MPI构造,并已证明它可以与操作符重载库dco/c++ [Lotz等,2013] 成功集成。OpDiLib [Blühdorn等人,2023年]是另一个与特定工具无关的库,用于支持使用OpenMP或OpenMP与MPI混合代码的共享内存并行性。另一种方法是将对MPI构造的特殊处理集成到AD库中,例如在TAF [Giering,2026年;Giering等人,2005年,2006年]中。OpenMP并行性是通过#pragmas来控制的,这很难用传统的运算符重载工具来实现,因为没有像传统MPI那样的函数重载功能,工具不得不依赖于构造函数和析构函数。最近有许多工具开发了扩展以支持共享内存并行性。PARAD [Kaler等人,2021年]扩展了运算符重载工具Adept以支持OpenMP并行构造,而OpDiLib [Blühdorn等人,2023年]则是通过将指令封装成宏或基于事件的OMPT API来支持OpenMP的。源代码转换工具Tapenade也通过一个理论模型 [Hückelheim和Hasco?t,2022年]和对Fortran的更完整支持(FormAD [Hückelheim和Hasco?t,2023年])获得了对某些OpenMP指令的支持。Enzyme源代码转换工具在LLVM中间表示级别上执行代码生成 [Moses和Churavy,2020年],利用了编译器在代码生成之前的优化步骤,并使用基于规则的转换来实现OpenMP和MPI [Moses等人,2022年]。PerforAD展示了为基于伴随的循环转换生成无竞争条件的OpenMP代码的能力 [Hückelheim等人,2019年]。特别是在模板代码方面,Hückelheim等人 [2018年]开发了一种高效的基于CPU的并行化策略。类似于Stan Math Library,可以使用专门的伴随函数来高效地处理矩阵乘法等操作 [Gremse等人,2016年],也有一些工具支持CUDA。AutoMat [Blühdorn等人,2022年]展示了为CUDA内核进行运算符重载的可能性,但严重依赖于重新计算值,并且缺乏像Revolve [Griewank和Walther,2000年]那样的高级检查点策略。Enzyme通过分析中间表示中的访问模式来生成无竞争条件的CUDA内核,尽可能避免不必要的屏障,并且只在其他线程可以访问的内存位置使用原子操作 [Moses等人,2021年]。
在现代硬件上,为了解决更大、更复杂的问题,需要并行执行算法。有许多框架可以编写具有不同优势的并行算法。长期以来,MPI、OpenMP及其组合一直是针对CPU的最常见方法。相比之下,针对GPU的可用框架要复杂得多,其中供应商特定的框架提供了最高的性能。维护一个能够针对所有这些硬件的代码库是不切实际的。为了解决这个问题,近年来出现了通用框架,可以从单一源代码针对所有平台进行开发,例如SYCL [Reyes和Lomüller,2016年]、KOKKOS [Edwards等人,2014年;Trott等人,2022年]或RAJA [Beckingsale等人,2019年]。然而,由于硬件的差异,像SYCL这样的相对低级别的方法虽然能够实现功能上的可移植性,但在没有大量平台特定调整的情况下很难实现有意义的性能可移植性。另一种方法是使用更高级的抽象来专注于性能可移植性,例如KOKKOS和RAJA,并在表达计算方式上牺牲通用性。另一个例子是使用领域特定语言(DSL)或抽象来交换通用性和性能可移植性。通过在应用领域牺牲通用性,DSL可以使用特定于该领域的更高级概念来表达算法,这些概念对应用程序开发人员来说具有自然的意义。DSL可以利用这种意义及其伴随的限制来执行广泛的优化,以为目标硬件生成高效的代码。
本文重点介绍Oxford Parallel库(OPS)[Reguly等人,2018年]中的模板计算的伴随模式微分,该库嵌入了DSL,旨在弥合高层科学计算与并行和分布式系统最佳使用之间的差距。OPS无缝集成到C、C++和Fortran中,为计算科学家和工程师提供了一套抽象,以适应结构化网格计算。Devito [Louboutin等人,2019年]嵌入在Python中的DSL旨在为有限差分偏微分方程求解器提供抽象。它从问题的高层符号表达开始,并使用模板编译器为应用程序生成并行实现。另一个重要的DSL是Halide [Li等人,2018年],嵌入在C++中,用于表达图像和数组处理任务。通过这些工具的高级表示,它们可以进行更广泛的转换,例如Devito可以制定伴随方程。OPS使用稍微低级别的抽象来处理结构化网格计算,这些计算在计算循环的层面上进行。OPS中的计算表示为一系列模板循环,对结构化网格中的每个元素应用特定的操作,这些元素访问固定的相邻元素模式,称为模板。换句话说,模板指定了如何组合网格中相邻单元的值来计算给定单元的新值。模板计算在许多领域得到广泛应用,包括CFD [Becker等人,2010年]、气候建模 [Heimbach等人,2005年;Singh等人,2023年]和计算金融 [Clevenhaus等人,2021年;Düring等人,2014年;Mollapourasl等人,2020年;Wang等人,2011年]。它们高效且可扩展,适用于并行计算架构上的大规模模拟。模板代码的并行化相对简单,其在现代并行架构上的性能优化已被广泛研究 [Lakshminarasimhan等人,2024年;Mudalige等人,2019年;Rahman等人,2011年;Raut等人,2020年;Reguly等人,2021年]。模板循环的结构与模板编译器等代码生成方法很好地结合在一起。YASK [Yount等人,2016年]、HOSTS [Stock等人,2014年]或Tiramisu [Baghdadi等人,2019年]可以从描述模板计算的更高层次模板抽象自动执行广泛的优化和并行化,并且可以处理不同的调度模式。虽然OPS采用了类似的方法,但它使用更低级别的抽象,使得OPS更容易适应更广泛的数值算法以及现有的代码库。通过将高层应用程序代码与低层并行实现分开,OPS消除了科学家和工程师需要处理并行性和数据移动的细节(如调度和平台特定的编程模型)的麻烦,使他们能够编写单一的高层源代码。同时,库负责目标架构的优化。
OPS定义了几种用于定义计算的关键抽象。块是N维容器,数据集定义在这些块上。数据集是与块相关的数据实体集合,作为计算内核的主要操作数,并通过预定义的模板进行访问。OPS中的计算通常使用并行循环构造来表达,其中用户定义的内核在相同的块上操作的特定迭代范围内和数据集上操作。在自动转换步骤中,OPS会将这些循环构造替换为目标特定并行实现的调用。循环要求用户指定用户内核如何访问每个数据集——例如数据是要读取、写入还是递增,以及控制数据访问的具体模板模式。OPS要求并行操作对各个网格点上的执行顺序不敏感(在机器精度范围内)。此外,内核不得对程序状态产生任何副作用,除了对数据集的更改之外,以及在并行循环参数中标记的归约操作。使用这种抽象,领域科学家可以表达结构化网格算法,而无需指定如何在具有异构内存空间的并行环境中执行它们。
列表1展示了OPS数据集、全局常量及其在OPS并行循环中的使用的示例声明。列表1显示了在Poisson示例应用程序中创建数据集,然后在OPS并行循环中使用这些数据集的API。科学家以这些计算循环的形式表达计算,提供将应用于每个网格点的循环体、块和迭代范围,以及循环体每个参数的描述符。这些描述符包含数据集、访问模式、模板和底层数据类型的信息。每个并行循环和用户定义的内核函数必须遵循一些额外的规则。OPS要求所有由内核使用的数据在使用前都注册到后端。这些数据可以在每个网格点上定义,也可以是全局作用域中的常量或归约操作。计算内核不得使用ops_decl_const在全球范围内注册之外的全局变量,这将在OPS中注册该符号,以确保该符号在所有使用的设备上都可用。可以通过ops_arg_global将网格不变常量传递给内核使用。在内核内部对ops_dats的数据访问是通过访问器对象进行的,这些对象允许对数据集进行多维相对索引。OPS的源到源翻译层利用这些信息生成针对特定硬件的低级并行实现。OPS支持并为不同的目标硬件生成代码,这些硬件具有不同的并行编程模型,如OpenMP、OpenACC、CUDA及其与MPI的组合。用户不需要修改OPS应用程序源代码即可使用这些范式。然后使用传统编译器编译转换后的源代码,并将其与OPS后端库链接。需要注意的是,所有数据集的所有权都交给了库,只能通过OPS API访问,这使得OPS能够跟踪所有数据集的状态以及何时需要移动数据(例如,在CPU和GPU之间)。能够在整个应用程序中跟踪数据移动和依赖性已经被证明是OPS执行复杂优化(如循环分块和惰性执行 [Reguly等人,2018年])的关键因素,并且在提供逆模式AD方面也起着至关重要的作用。
OPS通过允许研究人员开发单一的高层应用程序源代码,然后自动访问OPS提供的所有优化,从而提高他们的生产力。尽管将应用程序转换为OPS抽象会带来一些成本,例如重构循环为概述的循环体和ops_par_loop调用,但通过维护统一的代码库,应用程序可以无需开发任何平台特定的优化或应用程序代码,即可无缝地访问各种硬件的高效实现。OPS负责调度应用程序及其数据,这为在循环级别和循环链级别提供额外的优化机会。扩展像OPS这样的DSL以支持AD,使得一类新的应用程序能够充分利用DSL提供的性能可移植性和优化。为了为OPS应用程序启用AD,原始应用程序代码只需要进行最小的更改(例如种子和访问导数)和循环体的伴随。用户在OPS的抽象中定义应用程序,提供循环体及其伴随作为函数。OPS库本身可以根据领域的额外信息和限制生成代码,这使得诊断和优化成为可能,这些优化可能需要广泛的代码分析或对通用工具来说是不可能的。OPS API的描述性质结合源代码生成,可以精细控制循环和导数中的并行性。上述许多工具依赖于大型磁带来跟踪应用程序的控制流。在OPS中,并行实现的代码转换和优化步骤在编译时进行,并附加构建磁带的操作。这使得OPS能够通过使用计算步骤的高级表示来减少伴随AD的内存开销。在运行时,OPS构建了一个高度紧凑的模拟磁带,其中磁带中的每个条目代表一个并行循环或一个外部函数(见第2.3节)。本文基于Balogh和Reguly(2020年)的早期工作,该工作初步介绍了使用OpenMP的OPS伴随模式AD,并通过对OPS抽象的扩展以及针对CPU和GPU架构的并行实现进行了详细描述来进行扩展。总体而言,我们的文章做出了以下贡献:
1. 我们提出了一个模型,用于使用OpenMP和CUDA从同一个在OPS DSL中表示的源代码中对复杂模板应用程序进行逆模式求导。
2. 我们介绍了将计算循环的高级描述映射到伴随循环的并行实现的方法。我们展示了利用DSL提供的信息,可以对伴随循环的并行实现进行广泛的转换和优化。
3. 我们在OPS DSL中引入了一种AD磁带,其中的元素在计算内核级别进行处理,由于OPS抽象的存在,这使得存储更加紧凑。我们的实现表明,这种磁带使得实现诸如通过Revolve [Griewank和Walther, 2000]提供的最优AD检查点或任意“外部”(对OPS而言)伴随函数(例如线性求解器)等功能变得相对简单。
我们展示了使用OPS实现的三个应用程序的实验结果:一个简单的泊松方程求解器、Cloverleaf流体动力学迷你应用程序[Mallinson等人,2013年]以及一个对流扩散反应方程(CDE)求解器(实现了计算金融中的随机局部波动性模型[Wyns和Du Toit, 2017年])。我们提供了OpenMP和CUDA实现的详细性能比较。我们展示了两种计算泊松方程伴随值的实现方式:一种使用传统实现方式,另一种将应用程序表示为固定点迭代。此外,我们还展示了Revolve作为Cloverleaf和CDE求解器的检查点算法的性能。
#### OPS中的逆模式AD
我们对OPS的扩展使用了逆模式AD,不支持前向模式AD;从现在开始,讨论将集中于逆模式AD。
OPS通过ops_par_loop调用遵循控制流(如清单1所示)。这些调用指的是计算内核,列出了所有访问的数据集和访问模式。这使它们成为构建OPS控制流图的理想候选者。由于OPS在循环级别拥有所有所需的数据依赖性信息,我们可以在计算循环级别而不是表达式级别构建有向无环图(DAG),从而减少内存占用。OPS将在原始评估过程中缓存中间状态,并在逆序中将这些状态重新加载到伴随循环中。
#### OPS中的逆模式AD
进行逆模式AD所需的唯一组件是内核体的伴随函数。OPS已经可以从内核函数和模板提供的数据生成原始内核。如果OPS有关于内核体伴随函数的引用(如清单2中的函数),它可以使用相同的信息来生成内核伴随函数的并行循环。
在OPS循环中,每个数据集都是通过具有指定访问模式的模板进行访问的:读取、写入或增量。OPS中的所有前向循环必须仅使用收集模板,这意味着读取可以用任何具有任意点数的模板来实现,但增量和写入模板必须具有一个单点访问且偏移量为零。在伴随循环的逆模式AD中,数据流将从收集模板变为分散操作。由于这种反转,数据集上的写入和增量模板将变成伴随数据上的单点读取模板,而读取模板将变成具有多个点的增量模板。图1显示了五点模板原始循环和伴随循环之间访问模式的变化。上述情况有两个含义:首先,原始循环没有数据竞争,并行化很简单;其次,伴随循环在伴随数据上会有数据竞争。为了处理这些数据竞争,我们需要了解哪些伴随函数会被递增以及递增的模式或位置。然而,OPS已经根据并行循环描述符知道了所有数据的访问模式,伴随上的访问模式遵循相反的模式,因此OPS已经掌握了所有关于这些数据竞争发生位置的信息。
#### OPS中的其他循环参数类型
OPS还有两种循环参数类型:网格不变常量和全局归约。我们将全局常量视为被动数据,它们不会有导数,并引入了一种新的ops.scalar类型来处理具有伴随数据的全局读取参数。对ops_scalars的读取将在伴随循环中转换为全局归约。最后,全局归约在伴随循环中变为对伴随数据的全局读取。
#### OPS中的示例并行循环
例如,使用活性网格不变标量数据的并行循环。OPS使用ops_scalar数据类型来处理网格不变数据及其伴随的生命周期。清单3展示了一个具有活性标量值的内核。要在内核中使用opsscalars,必须使用带有ops_read访问方式的ops_arg.scalar包装器传递它。这个参数会通知OPS生成适当的伴随代码。请注意,伴随循环中的标量导数具有全局归约,这自然会对伴随循环的执行性能产生显著影响。另一方面,对于原始循环中的归约,可以从伴随循环中完全移除,并在伴随内存中产生全局读取。
#### 循环体的伴随函数
OPS处理伴随循环的代码生成方式与原始循环类似,需要一个指向循环体的函数指针。OPS依赖于用户提供的具有+_adjoint后缀的函数作为循环体的伴随函数。这些函数通常遵循与原始循环函数相同的规则,但有一些例外。带有_a1s后缀的参数被用作参数的伴随函数,并允许对其值进行非零偏移量的增量操作。为了帮助提供这些函数,我们实现了一个辅助脚本,使用Tapenade [Hascoet和Pascual, 2013]来生成用户内核的伴随函数,这可以帮助生成相对简单循环体的伴随函数,并为编写更复杂函数的伴随函数提供一个良好的起点。该工具充当调用Tapenade的包装器,将OPS API和一些C++结构映射为Tapenade可以生成的格式。该工具目前将ops_dat和ops_scalar参数视为活性参数。生成伴随函数后,OPS将它们用作用户提供的伴随函数来生成特定于后端的并行循环。在我们的实验中,我们使用此工具为Cloverleaf中的87个不同内核生成了伴随实现。一些内核需要进一步修改才能使其与CUDA兼容。在这些情况下,我们只需要进行局部修改以移除Tapenade用于处理分支的栈对象,但对于较长、复杂的核,则可能需要大量的编码工作来生成与OPS兼容的有效伴随函数。
#### 在OPS中获得导数
清单4展示了一个示例AD工作流程。代码在OPS的五个主要步骤中计算伴随值。第1-20行初始化数据集和网格不变值(ops_scalars),然后像传统的OPS应用程序一样计算原始值(第12-21行)。接着在第27行设置输出变量的导数种子,在第30行执行逆向传递。最后在第37行获取导数的副本。在原始执行过程中,OPS会将所有与导数传播相关的事件(循环、归约)保存到自己的磁带中,并在必要时保存中间状态。图2展示了OPS如何在前向传递过程中处理并行循环调用。OPS在内存中保存并行循环的小描述符以创建计算的表示。基于此描述符,OPS将执行原始循环并存储被覆盖的原始数据。原始执行完成后,用户设置输出值的伴随值(见清单4的第27行),最后,ops_interpret_adjoints调用将以逆序运行每个计算步骤的伴随函数,累积导数信息。这一步骤将分配并初始化在函数调用之前未设置的所有所需伴随变量。对于DAG中的每个事件,OPS将执行相应的伴随函数,例如如图2所示调用伴随循环。清单4还展示了访问数据集导数的两种主要方法。API类似于OPS中的数据访问API。用户可以直接访问ops_dat变量下的原始指针,也可以在某些内存空间中创建副本。
#### OPS中的逆向循环与高级磁带的交互
黑色箭头表示与并行循环相关的控制流,而蓝色箭头表示相应的数据移动。每个活跃的前向循环都会将一个小的描述符推入DAG,然后执行原始循环。在执行过程中,它将把被覆盖的数据保存到类似栈的存储中。在逆向传递过程中,OPS将调用生成的伴随函数,该函数将传播导数信息并加载保存的状态。
为了最大化性能并限制内存使用,识别哪些变量参与了导数信息的传播是非常有帮助的。现有的工具使用不同的约定——例如,大多数运算符重载工具将这些信息嵌入到它们的类型中。一些源转换工具(如Tapenade)将所有指向浮点值的指针视为活跃的,而所有按值复制的变量视为被动常量,或者依赖于活跃输入和输出的显式列表。在OPS中,所有用opsDECL_dat声明的数据集都被视为活跃的,用户可以使用opsDECL_DAT_passive函数定义被动数据集以避免分配伴随变量;同样,用户可以使用ops_par_loop_passive函数为不参与导数传播的循环创建被动循环。对于活跃循环,OPS可以从循环体的伴随函数签名中检测活跃变量。清单3中的示例内核显示了一个所有输入都是活跃的循环,它们的导数存在于伴随函数的参数列表中。通过简单地将输入的导数从给定内核函数的签名中省略,OPS将把该内核中的相关数据视为被动数据。
#### OPS之外的计算
在某些情况下,纯模板计算无法提供足够的性能,或者需要无法在OPS抽象中表示的方法。一个重要的例子是直接线性求解器,它们用于实现高阶紧凑离散化方案或现代方向分裂时间步进器(如Douglas、Modified Craig-Sneyd或Hundsdorfer–Verwer [Hundsdorfer等人,2003])。在非AD环境中,标准解决方案是请求访问OPS数据集背后的原始数据,在OPS之外进行计算,然后将数据返回给OPS。我们将这种在OPS抽象之外进行的计算称为外部计算。OPS提供了一个外部伴随API,以支持AD-aware方式实现这一功能。当OPS停止录制磁带时,用户需要提供一个函数来填充磁带中的空白部分。对于这些外部计算,OPS接受一个指向原始函数的函数指针、一个外部函数中访问的数据集的显式列表(类似于OPS循环),以及一个由用户提供的指向外部函数伴随函数的函数指针。在这些函数内部,用户可以使用传统的API甚至OPS循环执行任意计算,只要在函数结束前将所有数据所有权返回给OPS。重要的是,该函数不得修改不属于OPS的任何全局程序状态,因为如果使用检查点,可能会导致错误的值。为此,为了避免在重新计算过程中出现错误结果,OPS提供了一个ops_execute_external_function版本,它总是会保存受影响的OPS数据,并加载这些状态,而不是重新执行外部函数,尽管这会增加额外的内存开销。清单5展示了一个来自CDE应用程序的外部函数,该函数访问四个数据集(形成一个三对角线系统),求解这些系统,并将所有权返回给OPS。ad函数包含外部函数的伴随函数,其中包括转置系统的三对角线求解和对导数的OPS循环。
2.4 控制内存需求
伴随模式AD按照原始函数中每个指令的逆顺序执行其伴随函数,这需要变量的中间状态。大多数工具结合使用了两种方法来生成这些中间状态:在原始函数执行过程中保存状态并从内存中重新加载,或者重新计算。默认情况下,OPS在每个循环开始时会保存所有被写入的数据。在逆向传递过程中,OPS会加载这些数据,将所有变量恢复到原始循环之前的状态。通过这种策略,OPS不需要重新运行原始循环,从而提高了整体运行效率,但这带来了显著的内存需求,因为需要为所有这些中间状态创建检查点。现有工具使用检查点策略来管理由于重新计算而产生的额外计算步骤与伴随模式AD的高内存需求之间的权衡。OPS使用了针对时间推进方案开发的Revolve [Griewank和Walther, 2000]算法的改进实现。
2.5 重新录制磁带和反向累积
OPS提供了对其磁带对象的访问权限,以实现更细粒度的控制(如果需要的话)。多次重用磁带允许用户创建和组合磁带层次结构,从而支持诸如固定点迭代的反向累积等高级技术 [Christianson, 1994],这些技术可以大幅减少磁带和缓存状态的大小。清单7展示了一个示例代码,用于在原始函数中执行固定点迭代,并使用反向累积来计算伴随函数。在外部函数的伴随函数中,用户可以创建本地磁带来使用。与其他外部函数类似,用户可以使用被动循环进行误差计算、重播种或其他目的。
3 协调
如第2节所讨论的,OPS中的所有原始计算循环都会在模板中心写入数据或对全局变量进行归约操作。OPS会保存所有循环的归约状态,以避免重新计算相同的归约操作。对于数据集而言,并行的OPS循环已经确保了写入操作不存在竞争条件;因此,在调用用户内核之前,会在循环内部复制状态。在伴随循环中,OPS会在循环体开始时加载这些副本,并将后续对这些数据集的写入操作重定向到本地存储。伴随内核会反转数据流,从而在只读数据集的伴随数据上引入数据竞争。主要有四种数据访问模式需要考虑:归约操作、全局读取(ops.scalar)、网格点上的数据访问以及对低维数据的数据访问。前两种模式在逆向传递过程中会交换角色:归约操作的伴随数据将成为只读数据,而循环会对ops.scalar参数的伴随数据进行归约。在其他两种情况下,数据竞争会在访问多个点或低维(但仍依赖于网格)数据的读取模板中出现。在我们的实现中,我们选择单独处理低维数据的数据竞争问题,这将在第3.1节中讨论。
如果不使用高级转置来将伴随模板转换为无竞争的采集模板,有一些方法可以在存在数据竞争的情况下确保正确结果。可以使用特定于模板的着色 [Lülfesmann和Kawarabayashi, 2014] 技术,仅并行运行相同颜色的点。另一种常用的技术是对非线程局部的数据使用原子操作,或者对具有间接增量的数据集使用稀疏或局部归约。在这些解决方案之间做出选择的因素包括创建着色的成本、多块运行循环的开销、原子操作的开销,以及执行顺序对处理点的局部性和内存访问的影响。从原始循环中的模板,OPS可以获得关于在伴随循环中可能产生数据竞争的模式的信息,并可以生成不同配置的伴随循环。为了比较这些方法在伴随循环上的效果,我们使用了CDE应用程序中hv predictor步骤0循环的伴随函数的修改版本。该循环使用12点模板来增加2D数据集的伴随变量,所有低维数据集被视为被动数据。作为基准,我们使用贪婪着色算法创建了一个图着色方案,该方案不会出现数据竞争。在CPU上,我们使用了Spray [Hückelheim和Doerfert, 2021] 稀疏归约库创建了三个版本。第一个实现(标记为spray dense)使用与OpenMP内置归约类似的密集归约。这些实现会为每个线程分配一个线程局部数组,收集增量,然后在循环结束时汇总这些增量。第二个版本(spray atomic)在对伴随变量具有数据竞争的循环中对每个增量使用原子操作。最后一个Spray实现(spray block)使用BlockReduction归约器,将数组分割成大小为4,096的块,并仅在线程实际写入数据时才分配块。这些方法相对于着色方法的优势在于它们可以保持原始循环中的缓存友好网格点处理顺序,但同步开销较高。
在没有对代码进行高级转换以将伴随模板转换为无竞争的采集模板的情况下,有几种方法可以在存在数据竞争的情况下确保正确结果。可以使用特定于模板的着色技术 [Lülfesmann和Kawarabayashi, 2014],并仅并行运行相同颜色的点。另一种常用技术是对非线程局部的数据使用原子操作,或者对具有间接增量的数据集使用稀疏或局部归约。在这些解决方案之间做出选择的因素包括创建着色的成本、多块运行循环的开销、原子操作的开销,以及执行顺序对处理点的局部性和内存访问的影响。从原始循环中的模板,OPS可以了解在伴随循环中可能产生数据竞争的模式,并可以生成不同配置的伴随循环。为了比较这些方法在伴随循环上的效果,我们使用了CDE应用程序中hv predictor步骤0循环的伴随函数的修改版本。这个循环使用12点模板来增加2D数据集的伴随变量,并将所有低维数据集视为被动数据。作为基准,我们使用贪婪着色算法创建了一个图着色方案,该方案不会产生数据竞争。在CPU上,我们使用了Spray [Hückelheim和Doerfert, 2021] 稀疏归约库创建了三个版本。第一个实现(标注为spray dense)使用密集归约,类似于OpenMP内置的归约。这些实现会为每个线程分配一个线程局部数组,收集增量,然后在循环结束时汇总这些增量。第二个版本(spray atomic)在对伴随变量具有数据竞争的循环中对每个增量使用原子操作。最后一个Spray实现(spray block)使用BlockReduction归约器,将数组分割成大小为4,096的块,并仅在线程实际写入数据时才分配块。这保持了循环中聚合的内存访问操作,但对伴随数据的写入性能取决于模板本身。3.1 处理低维数据集如前所述,我们分别处理对低维数据集的访问,以减少大量线程共享数据访问的影响。对低维数据集原生数据的读取将在伴随循环中导致大量写入操作,这些写入操作是针对该低维数据的伴随数组。本质上,循环沿着低维数据集不变的维度进行求和。例如,在列表8中的2D循环中,循环会访问两个1D数据集x和v,其中x沿Y轴不变,v沿X轴不变。图5中,阴影部分使用了顶部的标记方法来处理正常数据集的伴随数组上的数据竞争问题,而对低维数据集则使用X轴上的标记方法。着色标记的组对正常数据集使用通用图形着色,稀疏还原组使用Spray [Hückelheim和Doerfert, 2021]稀疏还原库并配合原子操作,最后Stripes组使用模板在正常数据集上创建条纹着色效果。我们看到的情况与图4类似,但值得注意的是Spray密集方法由于对低维数组相同位置的写入次数增加而性能更优。此外,在Stripes实现中,Spray密集方法将成为性能最好的稀疏还原版本。这里,循环在最外层(非连续)维度上进行并行化。如果数据集在该维度上不是不变的(例如示例中的v),那么原始的并行化已经处理了数据集中的竞争问题。然而,在x的情况下,由于数据集在最外层维度上是不变的,所有线程会写入伴随数组的相同值,这会导致进一步的数据竞争。在这种情况下,线程将写入这些数组的所有元素,这就要求Spray块也分块分配整个数组(同时在这一过程中损失性能)。密集方法(在OPS中使用)是对Spray密集方法的一种轻微修改,通过内存池化来减少不断重新分配线程私有数据集副本的影响。图5显示了不同调度方法在CDE应用程序的hv预测器步骤0的伴随数组上的效果,该数据集具有较低维度的数据访问,并且网格大小分别为1,024^2和4,096^2。测量是在2.9 GHz的Intel(R) Xeon(R) Gold 6226R CPU上进行的,没有使用超线程,并使用了Intel(R) oneAPI DPC++/C++编译器2024.0.2。顶部的标签显示了网格上的并行化方法,而X轴上的标签显示了对低维数据伴随数组增量的处理方式。在GPU上,理论上原子操作足以处理数据竞争问题。然而,在存在低维数据集的情况下,由于冲突数量众多,原子操作的性能会大幅下降。如果可能的话,OPS将使用32×nwarps形状的2D CUDA线程块来减少同一位置内的写入次数。这种形状保持了高维数据集原生数据的聚合内存访问。然而,warp中的线程会对列表8中的v[j]和v[j±1]写入相同的值(一个在X轴上不变的数据集),因此当warp执行原子写入时,warp内的内存地址冲突会导致写入操作序列化。对于在Y轴上不变的数据集x,warp内的线程只有在存在非零模板时才会出现地址冲突,但地址冲突会在块内的warp之间发生。为了减少这些数据竞争的影响,OPS会将每个线程的增量累积到寄存器中,然后在内核的最后一步单独处理竞争问题。OPS处理四种重要情况:首先,对于像v这样的始终由warp共享的数据集,OPS将执行warp级别的还原,以便每个warp中只有一条线程对全局内存进行原子写入;其次,对于像x这样的始终在块内的warp之间共享的数据集,增量由每个线程在最后直接使用原子操作写入,因为单个warp中的线程在模板中心没有竞争问题。最后两种情况出现在三维或更高维度的循环中。如果数据集在X和Y两个维度上都是不变的,块内的所有线程将共享所有访问的数据,OPS将对所有写入的伴随值进行块级别的还原。图6(a)显示了这种优化在CDE测试应用程序中针对低维数据集的核的伴随数组上实现的实际加速效果,与仅使用纯原子操作相比;图6(b)显示了优化后的伴随循环的伴随因子。这些核使用了列表8中显示的访问模式,在两个不同的数据集上运行。CDE应用程序是一个在非均匀网格上的2D模板代码[Wyns和Du Toit, 2017]。计算有限差分估计需要知道X轴上的左右和Y轴上的上下网格坐标,以及当前网格点的位置。这些X和Y网格坐标存储在两个1D数据集中。Warp对X网格坐标有聚合读取操作,类似于列表8中的x,而从Y网格坐标的每次读取会像v一样在整个warp中广播一个值。对这两个数据集的三点访问可以用每个网格点访问12个单独的数据集来替代,但使用一点模板而不是三点,从而减少了伴随核期间的数据竞争。这种优化提高了原生核和伴随核的性能。仅使用一点模板时,相比在1,024^2和4,096^2网格上,分别实现了8.8–15.4倍和18.7–31.3倍的加速效果。这种加速使得伴随循环的运行速度不到原始原生循环的10倍。图6显示了我们在CDE应用程序的核上使用这种考虑访问模式的伴随还原和累积操作在低维数据集上的性能,网格大小分别为1,024^2和4,096^2,使用的是NVidia A100 GPU。(a)相对于仅使用原子操作的加速效果;(b)伴随因子。4 评估为了展示OPS中伴随模式AD的性能,我们在三个不同的应用上进行了测试。我们测试了OPS示例中的Poisson迷你应用、Cloverleaf [Mallinson等人, 2013] CFD应用,以及计算金融中的2D抛物线CDE求解器[Wyns和Du Toit, 2017]。在我们的测量中,passive表示运行原始应用所需的时间,forward表示运行应用并记录数据带的时间,adjoint表示计算所有输入的导数所需的时间。Poisson应用通过Jacobi迭代使用2D有限差分来解方程?2u=f,其中更新迭代可以表示为:ui,j(n+1)=(ui?1,j(n)+ui+1,j(n))dy2+(ui,j?1(n)+ui,j+1(n))dx2?dx2dy2fi,j2(dx2+dy2)。我们使用了一个4,096^2的网格和100次迭代来评估性能。这些测量显示了使用OPS记录数据带和累积导数信息的性能。我们引入了两种处理导数传播的方式:第一种是记录所有迭代并计算伴随值;第二种是利用代码可以表示为定点迭代的事实。因此,我们可以使用反向累积[Christianson, 1994]并通过数据带进行单次迭代的定点迭代来计算导数。Cloverleaf [Mallinson等人, 2013]是一个在笛卡尔网格上使用显式二阶精确方法解决可压缩欧拉方程的迷你应用。Cloverleaf是英国Mini App Consortium的一个迷你应用,广泛用于模板计算的性能基准测试。我们在两个示例问题上使用了OPS实现的Cloverleaf,一个问题的规模为960^2,有2,955个时间步长,另一个更大的问题为3,840^2,有87个时间步长。CDE应用程序是计算金融中的对流-扩散-反应代码[Wyns和Du Toit, 2017]。它使用Heston随机波动性模型通过线法计算欧式看涨期权的价格。CDE应用程序使用中心有限差分的有限体积方案,并在v=0边界处进行一阶迎风处理。它使用Hundsdorfer–Verwer ADI时间步进方案,并需要求解批处理的三对角线系统,这是由tridsolver [Balogh等人, 2022]库完成的。我们使用这个应用来突出两个方面的问题:低维数据集对性能的影响以及OPS中的外部伴随节点。4.1 测量设置我们在CentOS Linux版本7.9上使用单个Socket的Intel(R) Xeon(R) Gold 6226R CPU(2.9 GHz,无超线程)和376 GB RAM进行了OpenMP测量。CUDA测量是在CentOS 8上使用的AMD EPYC 7F72 24核处理器和NVidia A100 GPU(40 GB RAM)上进行的。所有代码均使用GCC 11.2和CUDA 12.0编译。对于Poisson应用,我们没有使用Revolve,但保存了100次迭代的全部中间状态。除非另有说明,否则我们报告使用29个和985个Revolve检查点的Cloverleaf的运行时间,分别用于960^2和3,840^2测试案例,以及10个检查点的CDE。所有测量都重复了10次,平均值作为最终结果显示。对于每个应用,我们测量了原始应用和使用组合前向和反向传递的AD版本。4.2 结果与许多模板应用一样,Poisson和Cloverleaf应用都是内存带宽受限的代码,所实现的网络带宽是衡量给定硬件性能的常用指标。CDE应用程序在大多数计算核中使用低维数据集,导致大多数循环的计算受限。表1显示了所有测量代码的绝对运行时间和峰值内存使用情况,以及通过评估伴随模型计算出的敏感性数量。最简单的应用是Poisson,它执行列表4中的计算,其中每个迭代都应用列表1中的模板。第一个版本直接执行伴随计算,在前向传递过程中缓存整个应用的所有中间状态。Rev.Acc.版本在执行100次迭代时不缓存任何状态,在伴随评估期间为反向累积创建一个新的数据带,为单次迭代记录一个小得多的数据带,并在该数据带上执行100次伴随评估。图7显示了应用程序内存使用的组成部分。Rev.Acc.方法通过仅缓存一次迭代(以及计算误差后的一个循环)并重用该数据带来减少了所需的内存占用,而不是使用100次迭代的数据带存储空间。这种方法的权衡是在伴随评估之间重置原始和伴随内存时引入的开销。尽管如此,这些开销对总运行时间的影响比我们在前向传递过程中缓存所有中间状态的影响要小。图7显示了应用程序的内存需求。Primal表示原始应用的内存使用情况,而Adjoint表示为数据集分配的伴随内存。Checkpoint大小标记了Revolve检查点的大小。Cloverleaf的First CP和CDE的Second CP表示由于数据集仅在开始时写入而只有一个更大的检查点。所有其他检查点使用统一的大小,标记为Checkpoint。Tape表示解释一个时间步长所需的保存状态。最终,结果部分在前向传递的末尾额外标记了用于保存数据集的副本。表1显示了使用OpenMP的Intel(R) Xeon(R) Gold 6226R 16线程的性能。Revolve在运行时间、被动内存使用量(GB)、运行时间伴随量(GB)和伴随内存使用量(GB)方面的表现,以及敏感性数量。不同算法的详细数据如下:
| 算法 | 迪利克雷过程(Poisson) | 反向累积(Rev.Acc.) | Cloverleaf | MCoverleaf | MCDE | Nvidia A100 |
|---------------|--------------|--------------|-----------|---------|---------|
| 4,096线程 | 2,100次迭代 | -0.677 | 0.54 | 2.97 | 4.39 | 4 |
| 4,096线程 | 2,100次迭代 | -0.677 | 0.54 | 2.51 | 4.97 | 3.72 |
| 960核心 | 2,955次迭代 | 985 | 64.07 | 0.194 | 5.02 | 6.31 | 7.03 |
| 960核心 | 2,955次迭代 | 100 | 64.07 | 0.194 | 6.94 | 4.25 | 7.33 |
| 3,840核心 | 2,87次迭代 | 293 | 12.98 | 20.29 | 30.56 | 6.19 |
| 3,840核心 | 2,87次迭代 | 103 | 293 | 12.98 | 21.85 | 19.79 | 6.66 |
| MCDE | 1,024线程 | 101.39 | 90.19 | 13.52 | 11.15 | 9.67 |
| MCDE | 4,096线程 | 103 | 44.25 | 93.09 | 26.14 | 13.44 | 7.63 |
表2比较了在Intel(R) Xeon(R) Gold 6226R(16线程)上使用OpenMP和在Nvidia A100上使用CUDA进行计算时的原算法和伴随算法的运行时间和内存使用情况。伴随因子表示伴随代码的运行时间相对于被动算法的比率。注意:对于迪利克雷过程,Rev.Acc.表示使用反向累积方法,即在固定点迭代中计算伴随值,并且仅存储最后一次迭代的数据。
在Cloverleaf算法中,每次时间迭代的核心数量显著增加(超过150个),这导致缓存内存大幅增加,如图7所示。由于这些数据的大小,创建包含每个时间步长数据的磁带是不可行的。因此,我们的Cloverleaf实现必须使用Revolve方法。
此外,CDE应用程序还使用了带有三对角线求解器调用的外部函数API。我们测试了两种问题规模:1,0242个网格和4,0962个网格,每种情况都有100个时间步长。在GPU上,三对角线求解器的调用占前向传递时间的约90%;而在CPU上,1,0242个网格的问题占23%,4,0962个网格的问题占46%。
我们使用了BabelStream [Deakin等人,2016]中的Triad内核来测量测试平台的可实现带宽。该内核在Intel(R) Xeon(R) Gold 6226R的单个NUMA节点上实现了60.22 GiB/s的带宽,在Nvidia A100 GPU上实现了1,287.70 GiB/s的带宽。可用内存带宽是迪利克雷过程和Cloverleaf应用程序的主要性能瓶颈。
表2(a)显示了整个应用程序在CPU上运行的带宽。对于所有测试用例,由于条带之间的额外同步,伴随版本实现的带宽值略低于原始版本。在Cloverleaf中,9602网格的带宽差异最大,为4.3 GiB/s,导致实现的带宽仅为BabelStream的76.4%。使用反向累积计算迪利克雷过程的导数的带宽与原始应用程序相当,且由于前向传递期间的开销较低,运行时间更短。
表2(b)显示了在Intel(R) Xeon(R) Gold 6226R的单个插槽上使用OpenMP测量得到的带宽值(未启用超线程):
| 应用程序 | 测试用例 | 带宽(GiB/s) | 峰值(%) |
|---------------|--------------|-------------|------------|
| 迪利克雷过程(Poisson) | 4,096线程,200次迭代 | 55.54 | 92.28 |
| 迪利克雷过程(Active) | 4,096线程,200次迭代 | 54.82 | 91.03 |
| 迪利克雷过程(反向累积) | 4,096线程,200次迭代 | 57.07 | 94.77 |
| Cloverleaf | 960核心,2,955次迭代 | 50.33 | 83.58 |
| Cloverleaf(Active) | 960核心,2,955次迭代 | 46.02 | 76.43 |
| Cloverleaf | 3,840核心,2,87次迭代 | 59.42 | 98.67 |
| Cloverleaf(Active) | 3,840核心,2,87次迭代 | 58.22 | 96.68 |
| BabelStream(Triad) | 数组大小:1.0 GiB | 60.22 | 100 |
| | | | |
表2(b)显示了在Nvidia A100 GPU上使用CUDA测量得到的带宽值。9602网格的尺寸太小,无法充分利用GPU的带宽。与OpenMP版本类似,伴随实现的带宽通常较低,但在CUDA的情况下,原子操作的影响导致带宽下降更大。
表3显示了CDE应用程序中计算受限内核的吞吐量。使用OpenMP时,伴随循环的吞吐量与原始循环相当。然而,在GPU上,原子操作的成本增加导致性能差异更大。表3显示,伴随循环的吞吐量仅为原始循环的55%-82%。
表3中的OpenMP吞吐量(GFLOPS/s)和CUDA吞吐量(GFLOPS/s)分别对应于不同的网格规模。
图8显示了原始应用程序和伴随计算之间的运行时间关系。表1中的运行时间和带宽值是使用相同的高级OPS源代码测量的,未进行任何修改。
图8清楚地显示了缓存中间状态对性能的影响。对于迪利克雷过程,OPS通过缓存每个网格点的中间数据集值,将每个内核中的写入次数加倍,从而减少了OpenMP前向传递所需的时间。由于CUDA中与缓存状态相关的内存操作成本较高,开销略大。然而,Rev.Acc.版本在减少前向传递的开销方面明显具有优势,因为它在固定点迭代过程中不需要缓存中间状态。当然,由于需要在反向累积过程中重置原始数据和重新种子伴随值,Rev.Acc.在评估伴随值时花费了更多的时间,导致开销分别增加了2.5倍和2.7倍和4.0倍。然而,在迪利克雷过程中,这些开销通过减少开销得到了补偿。
在Cloverleaf中,开销主要来自缓存最终迭代和处理Revolve检查点。对于每个问题,我们选择了迭代次数的三分之一作为检查点数量,这导致了9602问题的相对开销增加。然而,由于分组的检查点分配,实际差异被摊销了。所选的检查点数量导致重新运行前向内核部分的时间增加了2.1-2.5倍。就实际伴随内核而言,OPS在CPU上花费了3.3倍(9602)和2.8倍(4,0962)的时间,在GPU上花费了2.12倍(9602)和3.2倍(4,0962)的时间。
对于CDE应用程序,我们将检查点数量固定为10个,这在前向传递中导致了较小的开销,并且重新计算前向内核(包括带有缓存的前向内核)花费了2.05-2.87倍的时间。OpenMP和CUDA实现之间的伴随计算相对运行时间差异很大。主要原因在于三对角线求解器调用的时间:原始CUDA应用程序在小网格上花费了92%的运行时间,而在大网格上花费了88%;而OpenMP版本这一比例分别为23%和46%。正如我们所展示的,伴随计算中的三对角线求解是另一个三对角线求解加上一个小的导数循环,因此开销较小。在小型网格上,OpenMP和CUDA版本的伴随计算性能相似;而在大型网格上,OpenMP的开销更大。
图8还显示了Revolve对总运行时间的影响,图7显示了其对内存消耗的影响。迪利克雷过程应用程序显示了简单实现AD(缓存所有中间状态)的影响。原始应用程序使用了538 MB的内存(4个数据集,每个数据集大小为134 MB),此外还分配了额外的内存。图7中显示的磁带值对应于一次迭代的时间,其中循环写入了一个数据集。因此,OPS在每次迭代时存储了一个数据集的值,每次迭代增加了134 MB的内存,总共13.4 GB的内存用于100次迭代,以及前向传递的运行时间开销,但评估伴随值的时间更快。这突出了反向累积的主要优势,因为它一次只为一个迭代缓存数据,大大减少了计算伴随值所需的内存,使得内存需求保持不变,而不是随着迭代次数的增加和计算复杂性的增加而增加。对于相对复杂的应用程序,单次迭代所需的磁带大小显著增加,使得一次性缓存所有迭代变得不可行,许多应用程序不是固定点迭代,因此不适合使用反向累积。对于这样的应用程序,Revolve提供了一种解决方案,可以在额外的重新计算步骤的代价下控制内存占用。
图7还显示了每个Revolve检查点存储的数据所导致的内存成本,这决定了随着Cloverleaf和CDE中使用的Revolve检查点数量的变化,内存需求如何变化。双重效应是:OPS每次只缓存一个迭代(将这一类别的峰值内存减少了迭代次数的倍数),并将内存使用量固定在图中显示的磁带值。OPS在CPU上花费了3.3倍(9602)和2.8倍(4,0962)的时间,在GPU上花费了2.12倍(9602)和3.2倍(4,0962)的时间。
总之,Revolve通过额外的重新计算步骤和实际保存检查点的成本增加了运行时间。为了更清楚地了解对运行时间的影响,图9显示了在不同最大检查点计数下Cloverleaf和CDE的运行时间。如果应用程序在前向传递期间存储所有Revolve检查点,理论上的最小运行时间开销会出现。然后,每次迭代都会被重新运行一次,同时保存所有状态。另一方面,虽然使用较少的Revolve检查点可以降低内存需求,但会导致更多的重新计算步骤。然而,图9显示,随着检查点数量的增加,重新计算步骤的开销会迅速下降——因为只需要重新计算少量的前向步骤,并且允许更多的检查点存储在内存中带来的收益逐渐减少。此外,在重新计算循环过程中,OPS可以通过跳过那些在首次执行时已经保存了结果的计算来降低执行成本,这些计算会显著影响循环性能。图9展示了使用Revolve评估应用(包括原始应用)的旁导数的运行时间,随着可用Revolve检查点数量的增加,恢复中间状态所需的额外前向步骤也随之减少。
**5 结论**
在这项工作中,我们提出了对OPS DSL的扩展,以支持伴随模式下的优化(Adjoint Mode AD)。OPS允许从高级源代码表达结构化的网格应用,并处理并行性、数据传输以及对不同硬件后端的优化。这种扩展为CPU和GPU上的这类应用提供了性能可移植的导数计算能力。我们的AD扩展利用了OPS的领域特定抽象进行优化,这些优化使用更通用的工具将更为困难。我们利用OPS抽象中的计算循环描述来处理伴随循环中的数据依赖性和数据竞争问题,并将伴随循环映射到OpenMP和CUDA实现中。我们的方法在OPS中构建了一个紧凑的AD代码流,通过在循环级别处理导数而不是表达式来实现这一点。像支持外部伴随函数和通过Revolve进行最优检查点设置这样的扩展,使用这种高级表示法非常容易实现。
在两个受内存带宽限制的应用程序上的实验结果表明,伴随传播的带宽与原始原始应用程序相当。在CPU上,这些应用程序在大型网格上实现了超过90%的峰值带宽,在较小的测试案例中仅下降了7%。同样,在GPU上,活跃版本在能够饱和GPU的测试案例中实现了75-80%的峰值带宽。与原始应用程序相比,OpenMP版本在所有三个测试应用程序上的计算开销减少了10倍以内,在GPU上减少了4.6-6.6倍。我们已经将Revolve算法集成到OPS中,并证明可以按照预期控制伴随评估的内存需求。这使得大型模拟能够在伴随模式下运行,否则由于需要通过重新计算前向传递的部分来牺牲内存需求,这将无法实现。
我们在OPS中的伴随模型为性能关键的结构化网格应用程序提供了一种高效的导数计算方法。使用像OPS这样的DSL可以从相同的高级源代码开始,在两种硬件上解锁导数计算,使用户能够使用具有非常不同能力的多种硬件,而无需任何额外的开发成本。