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

摘要：
深度神经网络（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）的依赖分析。
该新颖分析方式带来两项关键优势：

  • 首先，块粒度分析可将复杂的多对多元素级依赖关系简化为不同内核块之间的一对一关系。这种简化允许在融合内核中利用更快速的共享内存（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