论文分享：Moirae: Generating High-Performance Composite Stencil Programs with Global Optimizations

1 概述

模板计算：

7点模板（3D）：在三维空间中，访问中心点及其上下、左右、前后的6个邻居。

27点模板（3D）：访问一个 3×3×33\times3\times3 立方体内的所有点。

由于物理模型越来越复杂，现在的计算不再是跑一个简单的公式，而是由一系列相互依赖的模板操作构成的有向无环图（DAG），这就是所谓的复合模板（Composite Stencil）

Moirae框架：将图优化（如算子融合）与内核优化（如分块、流处理）结合在一起，处理由多个相互依赖的模板构成的有向无环图（DAG），认识到图优化和内核优化之间存在复杂的相互影响，因此通过搜索全局优化空间来寻找最优组合，避免陷入局部最优。该框架能够生成具有最优全局优化的复合模板高性能代码。

论文将复合模板的性能优化分为两类：图优化和核函数优化。图优化将多个模板融合到单一GPU内核中，利用模板间局部性，而内核优化则高效地将融合后的模板适配到GPU上。

核心贡献：

1. 核心技术突破：扩展全局优化空间

不仅仅是纵向融合：传统的编译器通常只关注垂直融合，即生产者-消费者关系 。

引入水平融合：Moirae 创新地引入了水平融合，将具有相同输入数据的独立模板操作合并 。

数据复用最大化：这种做法能显著提高缓存命中率，减少对全局内存的重复访问，因为多个计算任务可以共享同一次加载的数据 。

2. 性能评估的基石：细粒度代价模型

轻量化预测：为了在成千上万种优化组合中快速筛选，Moirae 提出了一个轻量级的性能预测模型 。

数据移动表示法：这是该模型准确的核心。它不再模糊地估算开销，而是描绘了复合模板在特定优化策略下精细的内存访问行为 。

捕捉硬件特性：通过这种表示法，模型能够识别出访存冲突、数据对齐等影响 GPU 实际效率的关键底层细节。

3. 搜索效率优化：进化搜索与剪枝

自动决策：面对庞大的优化组合空间，Moirae 使用进化搜索算法（Evolutionary Search）自动寻找最优解 。

智能剪枝：搜索并非盲目进行。Moirae 会基于算子融合收益分析先生成一批“有前途”的候选图，排除掉明显低效的组合，从而极大地缩小了搜索范围 。

领域知识驱动：利用模板计算的领域知识，确保搜索过程既高效又具有物理意义上的合理性 。

4. 实施与验证

自动化流水线：所有上述理论最终集成为 Moirae 框架，实现了从计算图输入到高性能代码生成的全自动化流程 。

实战表现：通过对合成模板和真实世界（如气象模拟）模板的测试，Moirae 证明了其性能超越了目前最先进的（SOTA）模板编译器，如 Artemis 和 Open Earth Compiler 。

2 背景及相关内容介绍

2.1 复合模板的三个关键特征

1.复杂度要求 (D1)：包含至少三个算子

内容：复合模板总是包含至少三个模板算子 。

背景：在简单的教学示例中，模板计算可能只有一个操作（如平滑处理）。但由于现实中的物理系统（如大气动力学）非常复杂，模拟过程通常涉及多个连续或并行的计算步骤，因此定义其规模至少为三个算子起步 。

2.关联性要求 (D2)：算子间存在直接关系

内容：每个模板操作与其他操作之间都有直接的关联 。

逻辑：这种关联性源于不同物理网格之间的相互作用 。例如，在气象模拟中，计算“风速”的模板可能需要“气压”模板的输出结果作为输入，这种不同物理量（网格）之间的耦合决定了算子不是孤立存在的 。

3.结构性要求 (D3)：静态且无环的依赖关系

内容：算子之间的依赖关系是静态且无环的 。

  • 意义：

    静态：意味着在程序运行前，计算的先后顺序和数据流向就已经确定了，不会在运行时随机改变。

    无环：这意味着数据流是单向的，不会出现 A 依赖 B、B 又依赖 A 的死循环。这种特性保证了计算过程可以被表示为一个有向无环图（DAG） 。

举例：

2.2 模板计算优化

分块（tiling）

根据需要设定为任何维度和大小（形状）。不同的分块策略会直接影响程序的并行度、数据局部性，以及是否能进一步应用其他优化方法 。

流式处理 (Streaming)

流式处理会在缓冲区中保存一个或多个二维切片，并沿剩余维度扫描网格以完成模板计算。流式传输可以在模板计算过程中有效提升数据的重用和数据定位。

预取 (Prefetching)

预取常与流式传输结合使用。在流式优化中，通过片计算与网格扫描之间的必要同步，避免了计算和内存的重叠。因此，预取会提前将扫描点加载到缓冲区中。

循环展开 (Unrolling)

循环展开允许每个线程同时更新多个连续的模板计算，且这些计算所需的点可以被寄存器重用。

寄存器优化 (Register Optimization)

包括公共子表达式消除（CSE）或将中间结果存储到共享内存中 。

2.3 研究动机

2.3.1 前人工作

1.大多数现有工作只优化了复合模板优化空间的一部分，要么只做图优化，要么只做内核优化。前人的研究（如 Wahib 等人）将图优化视为组合优化问题并使用启发式搜索，但在内核优化方面非常有限，仅仅使用了分块（Tiling）等基础手段 。

2.图优化和内核优化不是正交的，这意味着它们之间并不是互不干扰的独立关系 。

3.不同的图优化方案会生成具有不同算术强度和访存模式的模板操作。由于算术强度变了，原本针对特定访存特征设计的内核优化（如流式处理或循环展开）的效果也会随之改变 。

4.结论：局部最优不等于全局最优

2.3.2 实验设计

论文选择了三个带有不同DAG的复合模板。例如，s10-i10-o4-e28 指的是其 DAG 中包含 10 个模板作、10 个输入网格、4 个输出网格和 28 条边的复合模板。

根据重用数据的大小，挑选出9个最有前景的图优化方案。我们利用 Artemis [16]，一款 SOTA 模板编译器，用于内核优化（包括铺砖、流式、预取、展开和寄存器优化）和代码生成（利用所有核的寄存器优化）。基于经验，我们挑选出20种最有效的内核优化配置。我们评估了复合模板在NVIDIA V100 GPU上的性能。所有模板的网格大小为256×256×256。

2.3.3 观察结果

观察 1：内核优化的效果取决于图的结构

没有哪一种内核优化（如特定的 Tiling 或 Streaming 配置）能在所有图优化方案中都表现最好。即使是同一个复合模板，当你改变它的图结构（例如改变了算子融合的方式）时，原本最快的内核配置可能就不再是最优的了。

观察 2：图优化的效果取决于内核的配置

最佳的图优化方案也会随着内核优化的改变而改变。在特定的内核配置下，某种算子融合方案可能表现优异；但换一套内核配置，该方案可能就变慢了 。

观察 3：全局组合优化带来额外增益

只有同时探索图和内核的组合，才能找到真正的性能最高点 。

3 研究内容

3.1 Moirae框架概览

核心处理步骤：

图候选方案生成 (Graph Candidate Generation)： 由于全局优化方案的组合数量呈爆炸式增长，Moirae 首先会对输入的复合模板进行分析 。它会预先筛选并生成一系列“有前途”的候选图结构（即不同的算子融合方案），以缩小最初的搜索范围 。

轻量级代价模型 (Lightweight Cost Model)： 尽管初步筛选了图结构，但每个图对应的内核优化（Kernel Optimization）空间依然巨大 。为了提高效率，Moirae 引入了一个轻量级的代价模型来预测性能，而无需对每种组合都进行实际的运行测试 。

数据移动表示法 (Data Movement Representation)： 为了让代价模型更快地收敛（即更准确地学习性能特征），Moirae 提出了一种数据移动表示法 。该方法直接提取代码中内存访问行为的特征，并将其作为代价模型的输入 。

进化搜索 (Evolutionary Search)： 在代价模型的指导下，Moirae 采用进化搜索算法来自动寻找性能最优的全局优化配置。

3.2 全局优化空间

Moirae 会针对每一种图优化方案中的每一个内核，去搜索其最优的内核优化方法及参数 。这种“图+内核”的同时搜索构成了其庞大的全局空间，如图 5 (a) 所示 。

为了便于搜索和代码生成，Moirae 使用计算图来表示复合模板的优化结果，使用输入/输出节点代表网格数据，使用**融合节点（FN）**代表一个或多个模板操作 。

FN 的生成：每一个融合节点最终都会被编译为一个独立的 GPU 内核 。

FN 的组成：包含三部分：输入的参数（FN.args）、计算主体（FN.body）和返回的结果（FN.rets） 。

递归融合：内核融合是基于 FN 进行的，融合后的节点还可以继续进一步融合 。

3.3 代价模型构建

高层表示（High-level）：仅嵌入优化方法和模板形状，缺乏足够的性能细节来支持复杂的“复合效应”分析 。

底层表示（Low-level）：仅计算寄存器使用量等指标，不足以在巨大的搜索空间中区分不同的全局优化方案 。

Moirae 的设计基于两个关键的性能洞察：

洞察 1（访存受限）：由于模板计算本质上是访存受限型的，内存访问行为对最终性能的影响具有决定性意义 。

洞察 2（多维影响因子）：影响内存性能的关键因素包括存储层级、局部性和访问次数 。这些因素又与程序的循环结构和并行方式紧密相关 。

由此，文中提出了一个数据移动表示法（dm representation）

不再使用间接的指标，而是直接记录程序在特定并行方法和循环结构下的细粒度访存行为 。考虑到循环结构和并行层级，对每一次内存访问行为进行分析和简化。

数据移动表示法

1）从程序中提取 dm 表示 (Trans 算子)

  1. Trans 算子构成：记录了一次数据从源（src）到目的地（dst）的转移 。
  2. 核心字段：每个 Trans 算子包含数组名称（name）、分配形状（shape）、内存层级（global/shared/register）、访问的起止索引（start/end index）以及访问数据量（num） 。
  3. 寄存器处理：由于寄存器不涉及复杂的局部性分析，其形状固定设为 1024，索引设为 0 。
  4. 算子合并：如果相邻的 Trans 算子访问的是同一个数组，且索引连续或有固定步长，它们会被合并以简化分析 。
  5. 结构保留：在提取时，Moirae 会保留程序原始的并行层级（Block、Thread 算子）和循环结构（For 算子）。

2）表示法的分析与简化

为了让代价模型能处理这些信息，Moirae 采用三层（循环级、线程级、线程块级）处理流程，将动态变量转化为静态特征 ：

  1. 消除变量：目标是移除所有的 For、Thread 和 Block 算子，并将索引中的变量（如 threadIdx、循环变量 k）替换掉 。

  2. 添加属性 (attr)：将消除的变量影响转移到 attr 字段中 ：

    • 循环级：将循环的上下界带入起止索引，计算总访问量（num），并记录偏移量（offset）（即相邻两次迭代之间的索引差值） 。

    • 局部性识别：偏移量的大小能反映数据的局部性（Locality）。例如，较小的偏移量意味着较好的局部性 。

    • 线程与块级：使用类似方法处理 threadIdx 和 blockIdx 。这种方法能有效区分不同的并行方式，并支持分析 Bank Conflict（银行冲突） 和 Memory Coalescence（访存合并） 。

从程序中推导出的 dm 表示：

DM表示的分析与简化：

3）代价模型构建

Moirae 使用了一个轻量级的模型进行性能预测 ：

  1. 特征张量化：由于简化后的 dm 表示不包含编译时未知的变量，索引、访问量、偏移量等直接转换为常数；字符串信息（名称、内存层级）采用 One-hot 编码 。
  2. 对齐与填充：对不同维度的线程和循环进行对齐，维度不足时进行填充（Padding）。
  3. 模型架构：采用 四层多层感知机（MLP），并以均方误差（MSE）作为损失函数。
  4. 复杂度控制：为了加快预测速度，模型仅输入访问量（num）最大的前 kk 个 Trans 算子。

3.4 全局优化搜索

3.4.1 图候选方案生成

核心权衡：虽然图优化的主要收益来自数据复用，但过度激进的融合（Aggressive Fusion）会导致 GPU 资源（如寄存器、共享内存）不足，反而降低性能。

渐进式融合策略：为了避免错过最优解，Moirae 不会一次性完成所有融合，而是从原始计算图开始，逐步增加融合程度，生成一系列具有不同融合深度的候选图。

收益评估指标：使用 benefit=Vhor+Vverbenefit = V_{hor} + V_{ver} 来量化数据复用机会，其中 VhorV_{hor} 代表水平融合（共享输入）的收益，VverV_{ver} 代表垂直融合（生产者-消费者）的收益。

算法逻辑 (Algorithm 1)：

  1. 计算所有内核对之间的 benefit 。
  2. 尝试融合 benefit 最高的两个内核 。
  3. 资源约束检查：评估融合后的内核是否会导致寄存器溢出，或者共享内存是否足以支持流式处理和预取等优化 。
  4. 若满足约束，则执行融合并加入候选列表；若不满足，则尝试收益次高的组合 。

3.4.2 进化搜索

这一阶段是在代价模型的指导下，为候选列表中的所有内核寻找最优的参数配置。

  1. 跨图联合搜索：由于相邻候选图之间存在许多相同的内核，Moirae 不会对每个图进行独立搜索，而是统一搜索所有内核，以消除重复搜索的开销。

  2. 基因编码：将不同的内核优化方法（如分块、展开等）及其参数编码为“基因”。

  3. 搜索流程：

    • 初始化：为每个内核随机生成初始种群。

    • 预测与测量：使用代价模型预测性能，挑选出预测表现最好的 pp 个个体进行实际的硬件测量。

    • 迭代进化：从表现最好的 qq 个个体中产生下一代。

    • 模型反馈：在每一轮进化完成后，利用实测数据对代价模型进行微调，使其预测越来越准。

3.5 代码生成与优化

3.5.1 技术栈与基础工具

框架基础：Moirae 采用 Relay 框架（一种深度学习图表示）来实现图优化 。

后端生成：它将每一个融合节点（FN）生成为对应的 GPU 内核，并调用现有的模板编译器 Artemis 来处理内核层面的具体优化和最终代码生成 。

为了最大化局部性，Moirae 采取了更激进的策略：

3.5.2 核心挑战：循环间距过大

如果只是简单地把 FN.body 里的模板操作直接生成代码，融合的收益会大打折扣：

数据被驱逐：如果每个模板操作都在各自独立的循环中运行，两个循环之间的间隔可能太长，导致第一个操作产生的数据还没被第二个操作用到，就已经从缓存中被“踢”出去了 。

案例分析：以图 11 为例，如果 %25 和 %26 模板在不同循环里，中间产生的数据 D 在被 %26 读取前可能就已经失效了 。

3.5.3 极致的循环融合

单循环执行：尽可能将一个 FN 内的所有模板操作都挤进同一个循环里执行 。

水平融合优化：对于符合水平融合模式的操作，只要它们访问的网格范围相近，就会被放入同一个循环中 。

3.5.4 解决融合顺序的难题

当一个 FN 包含两个以上算子时，融合的先后顺序会严重影响性能：

额外开销风险：如果先融合了某些算子（如 %25 和 %26），可能会因为循环内引入了过多的网格数据，导致后续再融合其他算子（如 %27）时的性能增益反而下降 。

策略借鉴：Moirae 借鉴了算法 1 的思路，通过收益分析（Benefit Analysis）来确定优先级 。

优先级与约束：优先融合那些收益更高的操作，并一直尝试融合直到触碰到硬件资源限制（如寄存器用量），从而榨干局部性优化的加速潜力 。

4 结论

论文将 Moirae 与 Artemis 和 Open Earth Compiler 这两款最先进的（SOTA）复合模板编译器进行了对比实验，证明其框架级的全局优化策略优于这些编译器原有的局部优化 。

充当了全局优化指挥官的角色。它的工作流程体现了其“框架”的特性：

  1. 输入：接收复合模板的计算图 。
  2. 图优化：分析数据复用机会，生成一系列融合节点（FN）候选方案 。
  3. 性能预测：通过其独创的 dm representation 代价模型预测不同组合的性能 。
  4. 搜索与决策：利用进化搜索算法在巨大的空间中锁定最优参数 。
  5. 输出：最终调用底层编译器生成优化后的 CUDA 代码 。