【AI 编译·II】Optimizing Deep Learning Inference Efficiency through Block Dependency Analysis

本文选自 ASPLOS’25，来自计算所谭光明研究员和孙凝晖院士。领域为 AI 编译→算子优化→算子融合→线程块级算子优化

摘要：
深度神经网络（DNN）中的算子间优化依赖于精确的数据依赖关系分析。传统的机器学习编译器（MLC）通常在元素级或算子级执行静态数据依赖分析，然而这种方式存在两大关键局限性：其一，复杂的依赖关系阻碍了高效的跨算子优化；其二，可并行计算部分未被识别，导致 GPU 资源未被充分利用。为解决上述问题，我们提出了 BlockDepend，一种全新的机器学习编译框架，通过 块级（block-level）依赖分析打破现有瓶颈。BlockDepend 通过分析编译流程中的低层阶段，提取关键的块级依赖信息，从而简化算子之间的复杂关系，发掘被隐藏的并行计算机会。这一机制使得编译器能够采用更具针对性的优化策略，提升内存访问效率，并增强 GPU 的利用率。实验结果表明，BlockDepend 在多个工作负载上表现出显著性能优势，分别相较于 NVIDIA TensorRT 与 AMD MIGraphX 实现了 1.71× 与 2.88× 的加速效果。

引言

在深度神经网络（DNN）推理任务中，有效利用 GPU 硬件资源对于实现最优性能至关重要。机器学习编译器（MLC）在这一优化过程中发挥着核心作用，其目标在于充分挖掘计算单元（如 NVIDIA 的 Streaming Multiprocessors 以及 AMD 的 Compute Units）及存储资源（如内存带宽与多级缓存层次）的潜能。

内核融合（kernel fusion） 是 MLC 所采用的一项关键优化技术。该技术通过将多个算子整合为一个高度优化的融合内核，减少频繁内核调度所带来的运行开销。根据算子之间的数据依赖关系，内核融合可沿两种主要方向展开：

  • 一方面，将具有数据依赖关系的连续算子融合可显著提升数据局部性，从而降低离片（off-chip）内存访问频率，并减轻 CPU 与 GPU 之间的上下文切换开销；
  • 另一方面，融合彼此之间无依赖关系的并行算子则可增强融合内核内部的并行性，从而提升 GPU 计算资源的利用率。

当前最先进的 MLC 系统依赖对算子之间数据依赖的分析，以发现并实施上述融合机会。==这种分析通常基于对单个算子处理逻辑的建模，以及其在计算图中的交互关系进行静态理解。==然而，当面对复杂数据依赖结构时，现有方法面临两难困境：要么依赖冗余计算或代价高昂的全局同步以实现融合，严重削弱优化效果；要么完全放弃融合，从而导致大量宝贵的 GPU 资源无法被有效利用。

为什么依赖冗余计算或要全局同步？
冗余计算：通过重复部分计算，避免处理复杂依赖，从而实现融合。但这会增加不必要的计算开销。
全局同步：为了确保数据一致性，系统需要进行代价高昂的全局同步，这会显著拖慢整体执行速度。

为克服上述局限，本文提出了 BlockDepend，一种全新的机器学习优化框架。与以往基于张量表达式和计算图进行静态依赖分析的 MLC 不同，BlockDepend 聚焦于不同内核中的线程块（thread block）之间的数据依赖关系，开展跨块级（block-level）的依赖分析。

除了本文提出的块级（block-level）依赖分析外，常见的级别还包括线程级、张量级和算子级。线程级关注单个线程内部的依赖，粒度最细；张量级和算子级以整个张量或算子为单位进行分析，粒度较粗，常用于传统 MLC 系统。相比之下，块级分析在粒度和性能之间取得更好平衡，更适合挖掘跨线程块的融合机会。

该新颖分析方式带来两项关键优势：

  • 首先，块粒度分析可将复杂的多对多元素级依赖关系简化为不同内核块之间的一对一关系。这种简化允许在融合内核中利用更快速的共享内存（shared memory）实现同步，从而避免冗余计算与高开销的全局内存操作。
  • 其次，BlockDepend 可识别出算子内部尽管存在数据依赖但仍具可并行性的块组。通过对这些块组进行战略性划分与重组，BlockDepend 能够发掘此前被视为不可并行的算子间潜在并行性，从而进一步提升 GPU 资源利用率。

本方法不仅有效应对了现有编译器在处理复杂数据依赖时的缺陷，同时也为 DNN 推理任务的性能优化开辟了新的方向，使得 GPU 硬件能力能够被更加高效地释放。

本文的主要贡献总结如下：

  • 提出 块依赖抽象（Block Dependency Abstraction），用于建模神经网络中不同内核之间并行任务单元的潜在依赖关系，并对现有编译器中的主要问题进行深入分析；
  • 设计一种 编译缺陷检测工具，可基于块间依赖类型识别当前编译策略中的低效行为；
  • 实现一套 编译优化与代码生成工具链，包括四种面向不同低效场景的优化方法，协同提升数据局部性与任务并行性；
  • 在 NVIDIA 与 AMD GPU 平台上完成系统实现，并通过广泛实验验证 BlockDepend 在提升数据访问效率与计算资源利用率方面相较于主流方案的显著优势。

背景与动机

算子间优化（Inter-operator Optimization）是提升深度神经网络（DNN）推理性能的关键，而数据依赖关系的分析则是实现此类优化的前提与核心组成部分。本节简要回顾数据依赖的基本概念、当前基于依赖信息所采取的优化方法，以及现有方案的局限性。

2.1 静态数据依赖分析（Static Data Dependency）

高效的 DNN 推理在很大程度上依赖于对计算图中算子之间的跨算子优化。**在计算图中，每个节点代表一个算子（例如卷积、GEMM），而边则表示张量数据在算子间的传递路径。**数据依赖关系源于这些连接，决定了算子必须遵循的执行顺序，以确保结果的正确性。

通过分析这些依赖关系，编译器可以识别潜在的并行计算机会，减少内存通信负载，并提升带宽利用率。图 1(a) 展示了 DNN 中常见的一种场景，描述了两个连续 GEMM 算子之间的数据依赖关系。GEMM0 的输出张量CC 是 GEMM1 的输入张量之一，且CC 中的每个元素均被多个EE 张量元素所使用，具体计算如下：

E0=C0⋅D0+C1⋅D2E1=C0⋅D1+C1⋅D3E2=C2⋅D0+C3⋅D2E3=C2⋅D1+C3⋅D3\begin{aligned} E_0 &= C_0 \cdot D_0 + C_1 \cdot D_2 \\ E_1 &= C_0 \cdot D_1 + C_1 \cdot D_3 \\ E_2 &= C_2 \cdot D_0 + C_3 \cdot D_2 \\ E_3 &= C_2 \cdot D_1 + C_3 \cdot D_3 \\ \end{aligned}

如图 1(b) 所示的数据依赖关系所揭示，编译器可在元素级或算子级执行融合优化。TVM 与 XLA 等编译器通过将逐元素操作进行内联，实现基于寄存器的数据直接计算；而 TensorRT 则采用算子级依赖分析，识别可并行融合的算子（如共享输入的同级算子），从而支持横向层级融合（horizontal layer fusion）。

这类传统的数据依赖分析方法主要基于张量表达式与计算图结构的静态理解，并在块调度与代码生成之前完成。因此，这些分析与优化过程可归类为静态分析范式（Static Analysis Approach），其在一定程度上限制了更细粒度的调度与优化空间。
image-20250711084557.webp

2.2 低效的 Kernel 融合来源

尽管 kernel 融合在优化 DNN 推理方面具有显著潜力，但现有方法在提升数据局部性和任务并行性上存在多种低效，尤其是在面对多样化工作负载时。

来源一：低效的数据访问与冗余计算

当前的算子融合技术能够高效地处理一对一的元素级依赖（例如 elementwise 操作），通过每个线程使用寄存器进行操作。但如表 1 所示，存在更复杂的算子类型，包括多对多、一对多或多对一的依赖关系，这些仍然难以实现高效融合。
image-20250713182710.webp
当前最先进的方法，如 AStitch 和 TVM，要么依赖全局内存进行同步，要么引入冗余计算以减少内存访问。在 TVM 中，当一个由生产者生成的单个元素被多个消费者元素所依赖时，就会出现多对多或一对多的元素级依赖。这种情况下，每个消费者线程会独立地重新计算已有的值，造成显著的计算冗余。

图 2(b) 展示了 TVM 尝试融合两个 GEMM 操作时产生的冗余。线程 0 的对应计算表达式如下：

E0 = C0 · D0 + C1 · D2  
   = C0 · D0 + (A0 · B1 + A1 · B3) · D2

在该场景中，GEMM1 中的每个线程都要计算一个输出元素，而这需要来自 GEMM0 的两个中间结果。但是由于线程局部寄存器彼此隔离，编译器不得不重复计算像 (A0 · B1 + A1 · B3) 这样已经在 GEMM0 中计算过的值。这种冗余在处理大规模张量时带来巨大的开销。

为确保数据正确性，一些方法（如 AStitch）会将中间结果写入全局内存以供共享。尽管这种方式减少了 kernel 启动次数及其开销，但却未能充分利用融合 kernel 中更快的本地内存。

来源二：并行性提升机会的缺失

在 GPU 架构中，要实现最佳性能，必须充分利用固定数量的执行单元，这依赖于调度器分发的线程块（blocks）。尽管一些先进方案如 Welder 和 Roller 能够自动构建高效的 tile 实现，但它们忽视了一个关键问题：tile（block）数量与硬件可用执行单元之间的不匹配。

这种不匹配导致部分执行单元在运行时处于空闲状态，特别是当元素级或 block 级的数据依赖过于复杂，迫使 kernel 顺序执行时。在最后一波（wave）执行过程中，这一问题尤为明显，硬件资源利用率严重不足。

我们的实验工作负载验证了该问题的严重性，平均浪费 block 槽（slot）达到 21.6%，这凸显了需要更智能的调度方式以更好地利用现有硬件资源。

图 3 展示了 Welder 和 TVM 等编译器中常见的低效情况，示例中两个连续的 GEMM 操作，每个包含 6 个 block，而 GPU 上只有 4 个执行单元。这意味着每个 GEMM 需要分两波执行。由于两个 GEMM 之间存在复杂的数据依赖，当前的 MLC（多层编译器）难以对其进行跨 kernel 的优化，导致顺序执行：GEMM1 必须等待 GEMM0 完成。
image-20250712173724.webp
在每个 GEMM 的最后一波中，只有 2 个执行单元在处理 block，其余 2 个保持空闲，导致每个 GEMM 有 2 个 block 槽被浪费。在这个例子中，浪费占了可用槽位总数的 25%。

这一低效对计算密集型工作负载影响显著，尤其是如 GPT-3 这样的模型。我们对 GPT-3 在 NVIDIA A100 GPU 上（batch size 为 768）的分析表明，约 35% 的总执行时间消耗在“非满载”的波段中，这些波段中平均有 38% 的 block 槽未被使用。因此，该模型的 MLP 层的平均 SM（Streaming Multiprocessor）效率仅为 73%。

这些发现表明，在现有方法中存在大量尚未开发的性能优化空间。

2.3 动机与挑战

与直觉相反，诸如 TVM 等多层编译器（MLCs）在融合复杂 kernel 时，并未利用显著更快的共享内存来同步中间结果。其根本原因在于：共享内存仅能在同一线程块（block）内的线程之间访问。 因此，只有当两个 kernel 在 block 级别具备一一对应关系时，才能利用共享内存进行优化融合。此时，对应 block 可以被融合，其复杂的元素级依赖关系可借助共享内存进行解耦，如图 1© 所示。

此外，正如图 3 所示，GEMM1 中的 Block0 和 Block1 的计算并不依赖于 GEMM0 中的 Blocks 2 至 5，这表明存在潜在的并行执行机会。要实现此类优化，显然需要获取 block 级的数据依赖关系。

然而，正如第 2.1 节所讨论的，当前的数据依赖分析是静态的，基于张量表达式和计算图进行推导。而这一分析过程及其对应的融合策略通常在 MLCs 生成 block 结构之前完成，因而无法获得精确的 block 级依赖信息。这些观察结果促使我们提出一种新的优化框架，旨在挖掘 block 级数据依赖，从而识别更多的优化机会，提升资源利用率并增强推理性能。

新的优化策略主要面临以下两大挑战：

挑战一：实际工作负载中跨 kernel 的 block 级数据依赖关系具有高度不规则性

在真实的工作负载中，算子之间及其融合后的元素级依赖结构复杂多变，识别其对应的 block 级依赖尤为困难。以 TVM 为代表的编译器在调度和代码生成之前执行算子间优化，缺乏对 block 级依赖关系的分析能力。此外，非多面体（non-polyhedral）编译器（如 TVM 和 Roller）以及依赖模板库的编译器（如 CUTLASS）也亟需一种通用的依赖分析方法。该方法应能够在 kernel 间分析并确定具体的 block 级依赖关系。我们将在第 4 节中详细阐述该依赖分析方法。

挑战二：跨 kernel 的 inter-block 优化

利用 block 级依赖关系实现跨 kernel 优化，对现有编译器而言构成了重大挑战。与传统的等价替换、循环变换等常见优化策略不同，inter-block 优化要求对多个 kernel 在 block 层面进行灵活的重写与组合。然而，当前编译器通常依赖统一的中间表示来生成 kernel，缺乏对 block 级调度与组合的支持。

为应对上述挑战，我们提出了针对不同依赖场景的优化方法，这些方法将在第 5 节中进行阐述。我们的目标是在保持正确性的前提下，通过优化 kernel 之间的交互，提升整体数据局部性与并行度，从而实现更高效的推理性能。

Overview

图 4 展示了 BlockDepend 的整体架构。BlockDepend 的目标是通过分析并优化不同内核之间的块级（block-level）依赖关系，从而提升深度神经网络（DNN）推理的效率。该系统由多个关键组件构成，这些组件协同工作以提升推理性能。
image-20250713183416.webp
整个流程从 ONNX 模型 开始，首先使用 NNFusion 在图级别基于静态数据依赖进行优化，并将其转换为中间表示 TE（Tensor Expression）。随后，Roller 构建算法对 TE 进行处理，生成初步的内核代码。

在“识别阶段”（见第 4.1 节），系统对内核代码进行性能剖析与依赖分析。第 4.2 节进一步根据具体的依赖类型对存在的效率问题进行分类。在此基础上，第 5 节则聚焦于根据分类结果对内核进行针对性的优化。

优化阶段中，BlockDepend 依次应用了一系列针对特定块级依赖问题的优化策略：

  • 首先是 BlockFusion（第 5.1 节），将图中的操作节点合并；
  • 接着是 ParallelFusion（第 5.3 节），用于增强并行执行能力；
  • 在图优化完成后，应用 Splitting（第 5.2 节） 技术以改善资源利用率；
  • 最后通过 Prefetching（同为第 5.3 节） 提升最终内核的性能。

这些优化技术均依赖于以 块 ID 映射（block ID mapping） 形式表示的块级依赖关系。例如，算法 2 就是基于 算法 1 所生成的依赖信息（dep）执行优化操作的。

最后一个阶段是 调优与生成，如第 6 节所述，涵盖了生成最终可执行代码并调整参数以匹配硬件架构的过程。

识别和利用块间依赖

识别各类块级依赖关系有助于发现操作符之间的优化机会。本节首先通过内核分析与代码剖析方法，阐述如何确定操作符之间的所有类型的块级依赖关系。随后，我们讨论每种块级依赖类型在操作符间优化中所带来的不同优化潜力。
image-20250713094643.webp

4.1 块级依赖关系的确定

本节阐述如何通过内核分析与源代码剖析来确定内核之间的块级依赖关系。为表述清晰，在计算图中，我们将数据的源操作符称为生产者（producer），将目标操作符称为消费者（consumer）。如算法所示，块级依赖关系分析过程可划分为两个阶段：首先，BlockDepend 对生产者进行分析，建立输出数据元素与生产者块之间的映射关系；随后，BlockDepend 重写并剖析消费者代码，以获取每个消费者块所使用的输入元素（即生产者的输出）索引；最后，根据这些索引与初始映射，建立消费者块与生产者块之间的映射关系。

为说明分析一对内核之间块级依赖关系的具体过程，我们以 BERT 中的前馈神经网络（由两个 GEMM 组成）为例，并采用较小的数据规模进行展示。

4.1.1 内核分析

BlockDepend 首先分析生产者，以获得输出张量中元素与其所在生产者块之间的映射关系，并将该映射生成一个元素索引计算函数。由于在计算图与张量表达层面难以获取块级信息，BlockDepend 从更低的层级提取所需的块结构信息。具体而言，BlockDepend 构建与硬件对齐的块结构，并采用高效的 tile 实现以贴合硬件单元的关键特性。该过程确保了高效的数据处理流水线，并建立了张量表达与硬件结构之间的清晰联系。

接下来，BlockDepend 基于构建的块信息生成 CalculateBlockID 函数（见第3行），用于将每个输出数据元素的索引映射到其所属的生产者块 ID。该函数利用 TensorShape、TileShape 与栅格化（Rasterization）等细节，构造出一个将元素索引映射到块 ID 的计算表达式。例如，对于一个采用行主序块划分的二维张量，其块 ID 计算函数如下：

block_id=⌊⌊indextensor width⌋tile height⌋×(tensor widthtile width)+⌊index mod tensor widthtile width⌋\text{block\_id} = \left\lfloor \frac{\left\lfloor \frac{\text{index}}{\text{tensor width}} \right\rfloor}{\text{tile height}} \right\rfloor \times \left( \frac{\text{tensor width}}{\text{tile width}} \right) + \left\lfloor \frac{ \text{index} \bmod \text{tensor width} }{\text{tile width}} \right\rfloor

该方法使我们能够确定输出数据中每个元素所对应的生产者信息。

4.1.2 kernel源代码剖析

在获取元素与生产者块之间的映射关系后，BlockDepend 还需进一步确定这些元素与消费者块之间的对应关系，从而建立生产者块与消费者块之间的依赖关系。

在该阶段，BlockDepend 采用不同的方法将输入数据映射到消费者块的标识符。在前一阶段，块的构造基于输出数据，每个块负责特定输出区域的计算，因此可通过构造参数直接建立输出数据与块之间的映射关系。而在当前阶段，输入元素需映射至消费者块的标识符，而这一映射关系依赖于具体的计算表达式。

在此过程中，BlockDepend 向消费者内核中插入代码，通过执行时的信息收集，获得生产者块与消费者块之间的依赖关系。这种方法能够适用于多种张量表达形式下的依赖关系提取。首先，BlockDepend 自动解析消费者内核，识别出针对特定张量的所有数据读取位置。例如，在图中所示的情况中，BlockDepend 找到对张量 I0 的访问语句，即从 I0 的某一特定索引读取元素的操作。随后，BlockDepend 解析该访问操作，并提取出被访问元素的索引值；然后在访问语句所在作用域内插入 CalculateBlockID 函数，以计算该元素对应的生产者块标识符。同时调用 RecordDep 函数，记录当前消费者块 ID 及其所依赖的生产者块 ID。

对于诸如三元操作符等特殊语句，BlockDepend 亦作特殊处理，以避免记录那些在实际执行中未被访问的索引。最后，BlockDepend 对插入代码后的程序进行编译并运行，所插入的代码会将每个消费者块 ID 及其所依赖的生产者块 ID 记录到块 ID 映射表中。如图中所示，记录项 [0, 1] -> 0 表示消费者块 ID 为 0 的计算依赖于生产者块 ID 为 0 和 1 的数据。至此，内核之间的块级依赖信息提取过程完成。
image-20250713095406.webp

4.2 基于依赖关系的优化分析

传统的静态分析方法（如前文所述）通常侧重于元素级的数据依赖关系，往往忽略了跨操作符的块级关系，从而在存在复杂依赖的场景下，可能错失重要的优化机会。例如，在 BERT 的前馈神经网络中，两个 GEMM 操作之间的元素依赖极为复杂：消费者 GEMM 的每个元素依赖于生产者 GEMM 的 3,072 个元素，而生产者 GEMM 的每个元素又参与计算消费者 GEMM 中的 768 个元素。

通过对块级依赖关系的分析，BlockDepend 能够在这些复杂的依赖中揭示潜在的规律，从大量无序的元素依赖中识别出有序的块级依赖模式。以图中所示的两个神经网络为例，基于块依赖分析所揭示的优化机会可归纳为以下四种典型模式：

一对一块依赖（One-to-One Block Dependency）。当每个消费者块仅依赖于一个唯一的生产者块时，即构成一对一的依赖关系。如图所示，该模式下，生产者块计算所得数据被写回至全局内存，而消费者块再将其重新加载至共享内存或寄存器中，造成冗余的数据访问。
image-20250713095419.webp
多对多块依赖（Many-to-Many Block Dependency）。在此类依赖模式中，每个消费者块依赖于多个生产者块，而每个生产者块也可能被多个消费者块所依赖。如图所示，该类型的依赖结构使现有编译器难以进行高效融合，导致资源浪费。实际上，两组内核的块形成了具有内部依赖的分组结构，而各组之间则保持相互独立，从而为并行执行创造了优化空间。

块的部分独立性（Partial Block Independence）。该模式表示，在某些数据区域上，当前内核的块不依赖于前一内核的输出。如图所示，这使得相关数据可在前一内核执行期间预取至 L2 缓存，从而减少当前内核执行过程中因访问全局内存而造成的停顿。

块的完全独立性（Full Block Independence）。当一个内核中的块与另一个内核中的块不存在任何直接或间接依赖时，即构成完全独立关系。如图所示，此类结构偏离固定计算模式，令现有如 TensorRT 等编译框架难以灵活地重排内核块的执行顺序。

基于四类块依赖的代码优化

5.1 数据复用优化

针对一对一的块级依赖，BlockDepend 对生产者进行一次计算，并将其结果存储于共享内存中，供消费者复用。为实现这一目标，BlockDepend 基于原始的生产者与消费者内核生成一个新的融合内核，如图 7所示。
image-20250713183806.webp
首先，BlockDepend 构建新的融合内核的执行逻辑。在该内核中，每个线程首先执行生产者内核的计算，并将中间结果写入共享内存而非全局内存（阶段 0）。随后，内核在块内进行线程同步，以确保数据一致性。接着，线程执行消费者内核的计算，从共享内存中读取数据以实现复用（阶段 1）。

在此过程中，BlockDepend 对 blockIdx 和 threadIdx 进行重映射，以融合具有不同形状的内核。新内核中每个块的线程数设置为原始内核中线程数的最大值，并按需执行不同阶段的计算任务。最终，在该融合内核中，中间结果通过共享内存复用，从而显著减少对全局内存的访问次数以及内核的调度开销。

5.2 内核划分与重构优化

该优化策略旨在解决由于线程块数量与GPU执行单元数量不匹配所导致的硬件利用率低下问题。BlockDepend 利用前期提取的块级依赖信息，对生产者与消费者内核之间的依赖关系进行细粒度分析，从而识别出可并行执行的潜在机会。具体而言，BlockDepend 将一对生产者和消费者内核划分为四个独立的内核，并将其调度至不同的流（stream）中执行，从而实现部分并行化执行。

以 BERT 的前馈神经网络（FFN）中的两个 GEMM 为例，BlockDepend 按照块级依赖对内核进行划分，并使存在依赖关系的两个内核部分重叠执行。如图所示，BlockDepend 首先将原始内核拆分，并将执行效率低下的块单独划分到新内核中。例如，GEMM0 被划分为 GEMM0_P0 和 GEMM0_P1。其中，GEMM0_P0 的块数与硬件资源匹配，能够充分利用 GPU 执行单元；而 GEMM0_P1 由于块数不足，无法独立形成完整波次（wave）。

随后，BlockDepend 在 GEMM1 内核中识别出可与 GEMM0_P1 并行执行的子部分。例如，当 GEMM1_P0 仅依赖于 GEMM0_P0 而不依赖 GEMM0_P1 时，可将 GEMM1 内核划分为 GEMM1_P0 和 GEMM1_P1。依据上述依赖映射，BlockDepend 生成相应代码，确保：

  • GEMM1_P0 在 GEMM0_P0 之后执行；
  • GEMM1_P1 在 GEMM0_P0 和 GEMM0_P1 均完成后执行。

BlockDepend 通过 GPU 的事件机制（Event mechanism）严格保证上述执行顺序。借助多流并行执行，GEMM0_P1 的执行可以与 GEMM0_P0 的部分执行过程重叠，从而提升整体资源利用率。

5.3 其他优化策略

并行内核融合（Parallel Kernel Fusion）
对于计算图中存在并行性的操作符，BlockDepend 可进一步利用其块级依赖分析能力进行优化。尽管某些编译器（如 TVM 和 TensorRT）能在操作符级别识别并行操作，但其缺乏块级抽象和灵活重组机制，限制了优化潜力。例如，TensorRT 可以将具有共享输入的兄弟 GEMM 融合为一个新操作符，但对于如 Mixture of Experts (MoE) 中不同尺寸的并行融合操作，则因算子库有限无法生成等效替代操作。

BlockDepend 通过对操作符内部块的精细分析与管理，灵活地组织这些块，构建新的融合内核。在该融合内核中，所有原始块被统一调度执行，每个块根据其编号选择对应的实现路径，从而实现块级融合。

数据预取（Data Prefetching）
尽管 L2 缓存中的数据在内核之间是可复用的，但当前编译器无法对其进行显式编程控制。TVM 等框架依赖硬件机制隐式管理缓存复用，而 TensorRT 的 AutoScratch 机制虽能提升中间结果复用，但未考虑权重数据的访问优化。

为提升数据访问效率，BlockDepend 引入了权重数据预取机制，并对中间结果进行持久性管理。由于权重数据仅在操作符内部有效且不可跨内核共享，BlockDepend 能调整其访问时机。在当前内核执行期间，BlockDepend 从全局内存中读取下一内核所需的权重数据，并将其存入 L2 缓存，以减少下一阶段访问全局内存引发的等待。
image-20250713095825.webp
BlockDepend 向内核中插入预取指令，建立跨内核的预取依赖关系；并利用指令级并行性（ILP），将预取操作与当前内核的计算和存储操作并发执行。

此外，BlockDepend 的块级分析还支持 感知分区的预取（partition-aware prefetching），这一点对于采用 L2 缓存分区机制的 GPU（如 NVIDIA A100、H100）尤为关键。BlockDepend 确保发起预取与消费数据的块位于同一分区，从而避免跨分区访问所带来的延迟与缓存容量压力。基于块 ID 和分区信息进行数据选择，有效减少跨分区访问，从而提升内存访问效率。

该预取策略引入了当前主流编译器尚未具备的 L2 缓存管理与分区感知预取机制，显著提升了深度学习推理过程中的内存访问效率。
image-20250713095846.webp

实现

我们的实现基于两个开源编译工具：Roller 和 NNFusion。整个编译流程以 ONNX 模型为输入起点。首先，使用 NNFusion 基于静态数据依赖进行计算图优化，如元素级操作的融合。随后，利用 Roller 的构建算法生成与硬件对齐的单操作符内核实现。

在此基础上，我们对内核中的块进行分析与管理，并基于块级数据依赖执行跨操作符优化（如第 5 节所述），以生成优化后的内核代码。最终，我们构建出可用于端到端部署的项目。

在代码生成方面，我们实现了支持跨操作符优化的内核模板。这些模板会根据块级数据依赖，对 Roller 所生成的内核代码进行重写与重组。该过程包括重构内核执行逻辑、数据索引方式、blockIdx 与 threadIdx 的重映射、共享内存的统一管理、代码分析以及代码插入等。

图 9 展示了一个内核模板，能够在 MoE（Mixture of Experts）模型中融合多个 expert 的实现，从而实现并行融合优化。生成的内核会基于块 ID 将来自多个 expert 或子程序的块调度为并发执行。每个子程序在公共缓冲区中分配自身的共享内存，计算其私有线程与块 ID，并执行原始内核的计算逻辑。

除该示例外，ID 重映射在内核划分中也起到了关键作用。例如，如图 8 所示，GEMM0_P1 通过块 0 和 1 执行原本由块 4 和 5 执行的任务。此外，如图 7 所示，内核中还集成了同步原语，以保证数据依赖的正确性。通过对原始内核的自动重写，BlockDepend 实现了新内核中块的高效调度与管理。

为避免优化带来的额外开销，我们在计算图中针对不同操作符组启用或禁用优化并进行性能测试。性能可能受限于硬件资源，或受到额外开销的影响。例如，在内核中通过共享内存进行数据复用，会增加寄存器与共享内存的使用，进而影响整体性能。因此，我们使用内核分析器快速评估不同优化策略的性能表现，以确定最优方案。

与某些需要数百轮评估的编译器不同，我们仅对少量启用不同优化策略的内核进行评估即可。
image-20250713100713.webp

评估（实验）

  • NVIDIA 40GB A100
    • CUDA 12.0
    • cuDNN v8.7.0
  • AMD Radeon MI100
    • ROCm 5.4.3
DNN Workloads
  • BERT、NeRF、Swin-Transformer、ViT、Conformer、NAFNet、BSRN、MMoE、MetaHeac、SparseMLP、GPT-3、LLaMA、Switch-Transformer
基线
  • ONNX Runtime v1.14.0
  • PyTorch v1.12
  • PyTorch XLA v2.2
  • TensorRT v8.5.3
  • TVM v0.12（使用 Ansor 进行 kernel tuning）
  • Welder
  • BladeDISC v0.4.0（实现了 AStitch）
  • MIGraphX v2.4（AMD 平台）
  • CUTLASS 3.1
  • xFormers v0.0.29
实验部分
  • 端到端推理性能
  • 真实大模型的推理评估
  • 性能提升（局部性、并行性）
  • ROCm 的评估
  • 编译时间

讨论

BlockDepend 主要面向主流深度神经网络中具有规则访问模式的操作符，此类操作符在不同输入下的数据访问行为保持一致，且在推理过程中占据主要执行时间，是本研究优化的重点对象。相比之下，那些常出现在特定领域应用中的稀疏或随机访问模式的操作符并不在本方法的覆盖范围之内。

BlockDepend 通过预剖析手段捕获内核的块级依赖关系。由于所针对操作符的访问模式稳定，这种在编译阶段完成的优化在执行过程中仍能保持其有效性。我们的实验采用随机输入数据进行验证，结果显示该方法具有良好的鲁棒性。

上述所提及的执行效率问题同样存在于训练过程中，因此本研究提出的优化方法并不限于推理任务，在训练框架中同样具有优化内核性能和提升 GPU 资源利用率的潜力。尽管本文的优化实践聚焦于 GPU，但相关策略亦可迁移至其他加速器上，通过分析任务间的细粒度依赖关系及优化空间，实现类似的性能提升。未来工作将进一步探索对其他异构加速器的支持。

相关工作

近年来，操作符融合技术在提升 DNN 执行效率方面受到广泛关注。一些研究工作基于元素级依赖对操作符进行分类。例如，DNNFusion 与 AStitch 依据依赖特征定义融合规则，Deepcuts 则结合硬件参数构建融合策略。然而，这类方法依赖人工设计规则，存在遗漏优化机会的风险。

另一些方法则利用操作符级别的依赖关系开展优化。例如，MetaFlow 与 TASO 通过图级变换融合操作符，发掘潜在并行性；TensorRT 与 ONNX Runtime 集成专家制定的融合规则以适配常见模型结构；Rammer 聚焦于并行操作符的优化，但缺乏处理复杂依赖关系所需的分析与优化能力；Welder 借助 rTile 完成操作符融合与重用数据的高效 tile 实现。

相比之下，BlockDepend 更进一步，深入挖掘不同块之间的依赖关系及相应的优化机会，提出了包括动态块调度在内的四类优化策略，以提升资源利用率。通过在更低层级开展分析，BlockDepend 可提取关键的块级依赖信息及其他细节数据（如每个执行单元的活跃块数、缓存分区归属等），从而实现跨操作符的深度优化。

此外，部分研究关注子图级的全局优化，例如 XLA 与 MLIR 通过图级分析识别融合机会，但其优化方式主要依赖人工设计的启发式规则，难以适应更复杂的场景。Souffle 虽对张量表达式进行数据依赖分析，但由于缺乏块级依赖信息，其优化能力仍然受限。

元信息概括

  1. 领域：深度学习编译与推理优化，特别聚焦于 GPU 上的跨操作符优化与内核调度效率提升。
  2. 创新点：
    • 提出 BlockDepend 框架，首次引入 块级依赖分析（Block-level Dependency Analysis）来优化不同内核之间的数据依赖关系。
    • 定义并利用四种典型的块级依赖模式（如一对一、多对多、部分独立、完全独立）以发现融合与并行化机会。
    • 实现跨操作符共享内存复用、内核划分重调度、并行融合、权重预取等新颖的优化策略。
    • 引入对缓存分区、执行单元活跃块数等底层硬件特征的感知，提升 GPU 资源利用率。
  3. 动机：现有编译器（如 TVM、TensorRT）主要基于操作符或元素级依赖进行融合，无法捕捉复杂的跨块依赖结构，从而错失优化机会，并导致执行资源浪费。BlockDepend 旨在弥补这一空白，实现更细粒度、更高效的推理优化。
  4. 方法简述：
    1. 依赖分析：
      • 使用预剖析（profiling）技术，分析生产者与消费者内核之间的块级依赖。
    2. 优化策略：
      • 共享内存复用（适用于一对一依赖）
      • 内核划分与重调度（适用于多对多、部分独立）
      • 并行内核融合（适用于完全独立）
      • 权重数据预取与 L2 缓存优化（跨内核数据访问优化）
    3. 代码生成与执行：
      • 基于依赖图构造新融合内核，重写 blockIdx / threadIdx，统一管理共享内存，自动插入同步与预取逻辑。
  5. 工具+平台：
    • 编译工具：基于开源编译器 Roller 与 NNFusion 进行二次开发
    • 输入模型：ONNX 格式
    • 目标平台：NVIDIA GPU（兼容 A100/H100，支持分区 L2 缓存）
  6. 局限性：
    • 目前仅支持访问模式稳定、规则的数据并行操作符，不适用于稀疏或非结构化访问场景。
    • 仅在 GPU 上实现，尚未适配其他异构加速器（如 TPU、ASIC 等）。
    • 在资源有限或共享内存压力较大的情况下，部分优化策略（如共享内存复用）可能带来负面性能影响。