OpenMM：一种与硬件无关的分子模拟框架

当今计算机体系结构的广泛多样性需要一种新的软件开发方法。OpenMM是一个分子力学模拟框架，允许单个程序在各种硬件平台上高效运行。

当今的计算机体系结构正处于快速发展和多样化的时期。十年前，大多数程序都运行在传统的单核处理器上，能够一次执行一个线程。在过去的十年里，那些简单的CPU已经被一系列多核CPU、1专用加速器(如Cell Broadband Engine)、2和所谓的图形处理单元(Graphics Processing Unit)所取代，这些处理器实际上能够进行强大的通用计算。3未来这一趋势可能会持续下去。很难预测短短五年后会有哪些架构可用，也很难预测哪些编程模型最适合利用它们。

这给所有程序员带来了两难境地，尤其是对那些科学和工程领域的程序员来说。一方面，他们的计算需求往往是极端的，涉及只能在巨大的超级计算集群上运行的模拟或其他类型的计算。另一方面，他们用于开发软件的资源通常相当有限。为单一架构编写和优化所有必需的软件是一项挑战。为几个差异很大的体系结构重复该任务是完全不可能的。

在某些方面，这种情况类似于计算机工业的早期，当时程序是用每台计算机的母语编写的。不同的处理器有不同的指令集，因此将程序移植到新的计算机上需要完全重写它。编译器通过在程序员和硬件之间引入抽象层解决了这个问题：程序可以用高级指令编写，这些指令可以自动转换成任何需要的处理器的机器语言。

只要所涉及的处理器在功能和操作模式上都是基本相似的，这个方案就能很好地工作。当处理器差异太大时，它就会停止工作。例如，期望针对单核CPU和大规模并行GPU高效编译单个源代码是不合理的。这些体系结构需要完全不同的算法才能有效地执行相同的计算，这超出了任何现有编译器的范围。

真正需要的是另一个抽象层，将程序员与其代码在其上运行的硬件隔离；不仅是它的指令集，而且是它的基本功能。人们应该能够使用适用于手头问题的高级概念来表达要解决的问题，而不需要指定使用什么特定的算法。然后，抽象层应该自动选择最适合当前可用硬件的那些概念的实现。这种方法已经在许多情况下成功使用，其中两个最突出的例子是用于线性代数的LAPACK4和用于3D图形的OpenGL 5。

为了在特定的分子模拟领域实现这一目标，我们开发了OpenMM，这是一个在高性能计算体系结构上执行分子力学的库。它允许程序员使用高级的、独立于硬件的API编写他们的程序。然后，这些程序无需修改即可在任何支持API的硬件上运行。原则上，这可以是任何东西，从低端的单个CPU核心，到每个节点上都有多个CPU核心和GPU的大型超级计算集群。

让我们详细考虑这些需求，看看它们是如何在OpenMM中实现的。

在任何接口中，选择正确的抽象级别都是至关重要的。目标是确定问题的哪些方面应该由用户确定，哪些应该留给库来确定。太高的抽象级别会使库变得毫无用处：问题描述的某些方面是用户真正关心的，如果接口不允许他们准确地描述这些方面，他们就不能使用它。另一方面，太低的抽象级别限制了库在各种硬件上高效实现问题的能力。如果要求用户指定实际上对他们并不重要的实现细节，例如用于执行计算的特定算法，则不再有自动选择更适合可用硬件的不同算法的选项。

在分子力学中，用户通常想要用势函数、约束、时间积分、温度耦合方法等来描述要解决的问题。他们应该能够指定这些，而不必描述，例如，使用什么方法来评估势函数。换一种略微不同的方式来表达，用户关心的是方程式，而不是算法。理想的界面应该允许用户指定他们想要建模的数学系统，同时让库以适当的方式自由地对该系统进行数值评估。

即使接口没有显式定义如何实现计算，它也可以很容易地限制实际实现的范围。需要非常小心地确保API的任何功能都不会不必要地限制可以在其上使用的硬件平台。

OpenMM中的一个示例是用于访问状态信息的机制。分子力学模拟涉及关于被模拟系统当前状态的各种数据：原子的位置和速度、作用在它们上的力等。传统上，模拟代码将这些值表示为内存中的数组。例如，当程序需要检查原子的位置时，它只需查看适当的数组元素。对于习惯使用这种代码的开发人员来说，访问原子位置的自然而明显的API将是一个获取原子索引并返回“该原子的当前位置”的例程。

不幸的是，该API不可能在许多架构上有效实现。例如，当在GPU上进行计算时，原子位置存储在设备内存中，将它们传输到主机内存是一项相对昂贵的操作。对于原子位置分布在网络上的许多不同计算机的集群来说，问题甚至更严重。任何假定可以随时快速、随机地访问原子位置的程序，在这些系统上的运行速度都会非常慢。

OpenMM通过显式地不直接访问状态数据来解决这个问题。相反，用户调用一个例程来创建State对象，预先指定应该存储在该对象中的所有信息。这有两个好处。首先，因为OpenMM事先知道用户将要请求的全套信息，所以它可以通过几个批量操作高效地收集这些信息。其次，用户意识到他们正在执行昂贵的操作，因此将仔细考虑他们何时以及如何访问状态数据。他们不会被看似微不足道的API调用(例如“获得ATOM 5的位置”)所误导，而这些API调用实际上很昂贵。

如果这类库要想成功，它必须将接口和实现之间的划分作为其设计的一个基本方面。程序的作者应该不需要指定要使用什么实现。当程序运行时，它应该自动选择最适合可用硬件的任何实现。同时，它还应该允许程序查询可用的实现并手动选择要使用的实现。例如，程序的用户有时可能希望在主CPU上执行计算，其他时间则在GPU上执行计算。

同样重要的是可扩展性。随着新硬件的出现，新的实施将是必要的。不可能枚举所有可能的实现，并且接口不能尝试这样做。它必须是可扩展的，这样就可以独立于主库编写新的实现，并且现有程序无需修改就可以使用它们。

OpenMM通过插件架构实现了这一点。每个实现(或“平台”)都是作为动态库分发的，只需将其放在特定目录中即可安装。在运行时，将加载该目录中的所有库，并使其可供程序使用。

它还提供了一种不同类型的可扩展性：插件不仅可以实现新平台，还可以向现有平台添加新功能。重要的是要理解，OpenMM不仅仅是一个用于执行某些计算的库；它还是一个旨在统一整个问题域的体系结构框架。虽然该库内置了特定的功能(例如，特定的潜在功能和集成方法)，但它也允许通过插件添加其他功能。我们的目标是提供一个框架，在这个框架内几乎可以实现任何分子力学计算。

我们现在考虑如何创建满足这些目标的体系结构。OpenMM基于分层的体系结构，如图1所示。最高级别是公共API，开发人员在自己的应用程序中使用OpenMM时使用该API进行编程。在任何这样的库中，公共API必须以与问题领域(例如分子力学)相关的术语来表达概念，而不涉及这些概念是如何实现的。在OpenMM中，这些概念是粒子、力、时间积分方法等。例如，Force对象指定粒子之间交互的数学形式，但不规定用于计算它的特定算法。

公共API通过调用较低级别的API来实现，该较低级别的API充当独立于平台的问题描述和依赖于平台的计算内核之间的接口。OpenMM将这个低级API表示为一组抽象的C++类，每个类定义一个要完成的特定计算。请注意这两个接口所扮演的非常不同的角色：公共API由核心OpenMM库实现，由用户调用；低级API由插件实现，由核心OpenMM库调用。

在体系结构的最低层是计算内核的实际实现。这些代码可以用任何语言编写，并使用适用于它们在其上执行的硬件的任何技术。例如，它们可能使用诸如CUDA或OpenCL之类的技术来实现GPU计算，使用Pthread或OpenMP来实现并行CPU计算，使用MPI来跨集群中的节点分配工作，等等。

这就留下了选择和调用低级API实现的关键任务。对于OpenMM，这意味着实例化定义每个内核的抽象类的具体子类。此任务由平台对象协调，该对象充当计算内核的工厂。公共API的每个类查询平台以获得它所需的每个内核的具体实例，然后使用该实例执行其计算。因此，选择使用哪种实现完全包括选择使用哪种平台。程序也可以选择不指定平台，在这种情况下，根据可用的硬件自动选择一个平台。

实际上，安排要稍微复杂一些。平台不直接创建内核，而是将任务委托给一个或多个KernelFactory对象。这就是插件向现有平台添加新功能的方式：它定义一个计算内核，创建一个可以创建内核实例的KernelFactory，并将工厂添加到平台。当平台稍后被要求创建该内核的实例时，它使用新的KernelFactory来执行此操作。

我们现在考虑一个具体的例子，说明这个体系是如何在实践中工作的：分子系统中原子之间非键相互作用的计算。在大多数模拟中，这占了大部分处理时间，因此对其进行很好的优化是非常重要的。

在设计用来在CPU上运行的传统代码中，有一些成熟的技术可以有效地做到这一点。7首先建立一个邻居列表，该列表显式地枚举每一对足够接近相互作用的原子。通过使用基于体素的方法，这可以在O(N)时间内完成。然后循环邻居列表中的所有原子对，并计算它们之间的相互作用。OpenMM的参考平台(编写为在单个CPU线程上运行)就是以这种方式工作的。

不幸的是，由于需要间接内存访问，GPU上的邻居列表效率非常低。对于每个邻居列表条目，必须加载有关所涉及的两个原子的信息(位置、电荷等)。由连续线程处理的原子的索引不需要遵循任何模式，因此不能合并存储器访问。

因此，我们开发了一种更适合在GPU上运行的替代方法。我们把全套原子分成32个一组的块。然后，如图2所示，这组N2相互作用被分成(N/32)2个平铺，每个平铺都涉及两个原子块之间的相互作用。为了处理瓦片，我们将所涉及的64个原子的数据加载到共享内存中，计算它们之间的所有1024个相互作用，最后将产生的力和能量写出到全局内存。代替传统的邻居列表，我们使用了包含交互的瓦片列表：有效地，指定了32个原子的哪些块交互的邻居列表。其他研究人员也开发了在GPU上计算非绑定交互的算法。-11。

我们的方法是为在NVIDIA GPU上使用而设计的，每个块(32个原子)的大小被选择为与这些处理器的SIMD宽度相匹配。使其适应其他类型的处理器，甚至其他GPU，都需要对算法进行修改。例如，某些AMD GPU的SIMD宽度为64，因此必须以不同的方式在平铺之间分配线程才能获得最高效率。

因此，我们需要几种完全不同的算法来在不同的硬件上高效地实现相同的计算。然而，值得注意的是，算法的选择仅取决于硬件，而不取决于所计算的力的精确形式。在分子模拟中，有许多不同的数学形式用于表示非键相互作用，不同的数学形式包括范德华相互作用的建模方式、截止点的平滑、溶剂筛选效应等。这些都是科学家在运行模拟时真正关心的重要区别。理想情况下，程序员应该能够选择交互的函数形式，并且仍然使用可用硬件的最有效算法来计算它们。用户应指定要使用的方程式，库应确定如何最好地评估这些方程式。

OpenMM通过其CustomNonbondedForce类实现此目标。此类允许用户为原子之间的成对能量指定任意数学函数。该函数可能取决于任意一组原子参数和列表函数，以及各种标准数学函数。例如，以下代码行创建CustomNonbondedForce来计算Lennard-Jones 12-6交互作用：

第一行指定作为距离r的函数的相互作用的能量：

其中，使用Lorentz-Bertelot组合规则合并来自两个相互作用原子的参数：σ的算术平均值和ε的几何平均值。接下来的两行指定参数“sigma”和“epsilon”应该与每个原子相关联。

OpenMM现在的任务是在各种硬件平台上高效地实现这一点。它首先解析用户指定的表达式，然后解析微分能量以确定力的表达式。然后将每个表达式转换为指令序列。为了计算表达式，引用和CUDA实现循环遍历指令并执行每个指令，从而有效地充当内部语言的解释器。

对于基于OpenCL的实现，更好的解决方案是可能的。因为OpenCL允许在运行时从源代码编译程序，12所以可以使用插入到适当的硬件特定算法中的用户定义的数学表达式来合成内核。然后，内核被编译成设备的机器码，消除了解释表达式的成本，并产生了几乎与整个内核都是手工编写的一样快的性能。

我们强调，这种方法在我们的代码中允许极大的灵活性，允许快速开发的强大组合(即，可以很容易地更改粒子之间交互的关键底层方程)，同时仍然保持快速执行(因为底层优化，特别是使用OpenCL，允许最小的开销)。这为我们的代码打开了新用途的大门，特别是在模拟粒子相互作用的新方法的快速发展方面，例如用于分子模拟的新的隐式溶剂模型。

我们实现了与分子模拟中所有最广泛使用的能量项相对应的力类：各种键合力，用于非键相互作用的Lennard-Jones和Coulomb力，用于长程库仑力的Ewald求和和粒子网格Ewald，以及广义Born隐式溶剂模型。OpenMM还包括几种时间整合方法和实施距离约束的能力。这些功能在三个不同的平台上实现：用C++编写的参考平台，用于NVIDIA GPU的基于CUDA的平台，以及用于各种GPU和CPU的基于OpenCL的平台。

我们之前已经发布了在模拟各种蛋白质时CUDA实现的基准。在显性溶剂(共73,886个原子)中模拟318个残基蛋白质时的速度为5 ns/d，在隐式溶剂中模拟33个残基蛋白质(共544个原子)时的速度为576 ns/d。我们还将其与几个广泛使用的分子动力学软件包在模拟显式溶剂中的80个残基蛋白质时的单CPU核心性能进行了比较。结果发现，它比Gromacs快6.4倍，比NAMD快28倍，比Amber快59倍。(GPU计算在NVIDIA GTX280上运行，CPU计算在3.0 GHz Intel Core 2 Duo上运行。)。

OpenMM的最新功能是自定义力，它允许用户指定力的形式的任意代数表达式。除了上述CustomNonbondedForce之外，还有一个用于键合相互作用的CustomBondForce，用于独立施加到每个原子的力的CustomExternalForce，以及支持多种隐式溶剂模型的CustomGBForce。这些在OpenCL平台上最有用，因为它允许在性能损失很小的情况下使用它们。在初步测试中，我们发现使用CustomNonbondedForce实现的Coulomb和Lennard-Jones力仅比标准的手工编码实现慢约4%。这意味着，没有GPU编程经验的科学家仍然可以为他们的非绑定交互实现任意函数形式，并获得几乎与手动调整的GPU代码一样好的性能。

计算机架构的快速、持续变化。

..