前言在深度学习推理加速领域计算图的优化程度直接决定了模型在硬件上的执行效率。昇腾NPU上的CANN 作为昇腾 NPU 的核心软件栈其图编译引擎 GE 承担着从框架算子到硬件指令的复杂转换任务。在实际部署过程中我们发现未经优化的计算图往往存在大量细粒度算子这些算子之间的边界成为性能瓶颈的主要来源。graph-autofusion 正是为解决这一问题而设计的自动融合引擎它在 GE 的图编译阶段自动识别可融合的算子组合将其合并为单一的融合算子从而显著减少 Kernel 启动开销和 Global Memory 的访问次数。本文将深入剖析 graph-autofusion 的架构设计、融合策略以及在昇腾 NPU 上的实践经验。graph-autofusion 的核心价值在于其自动化的融合能力。传统的算子融合往往需要手工编写融合算子的实现代码这不仅工作量巨大而且难以覆盖所有可能的融合场景。graph-autofusion 通过模式匹配和数据流分析能够自动识别计算图中的融合机会并在保证语义正确的前提下执行融合变换。这种自动化机制使得模型开发者无需关注底层优化细节同时又能获得接近手工优化的性能表现。对于昇腾 NPU 这种具有特殊硬件架构的加速器而言自动融合引擎的价值更加凸显因为它能够充分挖掘硬件的并行计算能力和存储层次结构。计算图优化的挑战与融合的必要性现代深度学习模型通常由数十甚至数百个算子组成这些算子通过数据依赖关系连接形成计算图。在传统的执行模式下每个算子对应一个独立的 Kernel运行时需要依次启动这些 Kernel 并在它们之间传递中间结果。这种执行方式存在两个主要问题。第一Kernel 启动本身存在固定开销当算子粒度较细时这部分开销在总执行时间中的占比会变得非常显著。第二算子之间的数据传递需要经过 Global Memory而 Global Memory 的带宽远低于片上存储频繁的读写操作会严重制约整体性能。融合优化正是针对这两个问题提出的解决方案。通过将多个相邻算子合并为一个融合算子可以减少 Kernel 的启动次数从而降低启动开销的累积效应。同时融合后的算子可以将中间结果保存在片上存储或寄存器中避免不必要的 Global Memory 访问。对于昇腾 NPU 而言这种优化尤为重要因为 NPU 的计算单元具有高度的并行性只有充分的数据复用才能发挥其峰值性能。graph-autofusion 作为 CANN 图编译器的重要组成部分承担着自动发现和实施融合优化的核心职责。从硬件架构的角度来看昇腾 NPU 采用了立方体矩阵计算单元和向量计算单元的异构设计。立方体单元专门用于矩阵乘法等密集计算而向量单元则处理激活函数、归一化等轻量级运算。graph-autofusion 的融合策略充分考虑了这种硬件特性它会将适合立方体单元的算子组合与适合向量单元的算子组合分别进行优化。这种硬件感知的融合策略能够最大化硬件利用率同时避免因不合理的融合导致的资源冲突。graph-autofusion 的架构设计graph-autofusion 的架构可以划分为三个主要层次图分析层、模式匹配层和融合执行层。图分析层负责对输入的计算图进行静态分析构建数据流图和控制流图识别算子之间的依赖关系和执行顺序。模式匹配层维护着一组预定义的融合模式通过子图同构算法在计算图中搜索匹配的算子组合。融合执行层则负责实际的融合变换包括算子的合并、中间张量的消除以及融合算子的代码生成。图分析层是整个融合引擎的基础。在这一层graph-autofusion 首先对计算图进行拓扑排序确定算子的执行顺序。然后它会对每个算子进行数据流分析计算其输入输出张量的形状、数据类型和内存布局。这些信息对于后续的模式匹配至关重要因为融合的可行性不仅取决于算子的类型还取决于张量的具体属性。例如两个连续的 Elementwise 算子只有在输出张量的形状完全一致时才能融合否则会产生语义错误。# WHY: 展示图分析层如何提取算子属性用于融合判断 def analyze_operator_properties(graph, op_name): op graph.get_operator(op_name) properties { type: op.type, inputs: [(t.shape, t.dtype, t.layout) for t in op.inputs], outputs: [(t.shape, t.dtype, t.layout) for t in op.outputs], compute_cost: estimate_flops(op), memory_footprint: estimate_memory(op) } return properties模式匹配层是 graph-autofusion 的核心组件。它维护着一个融合模式库其中包含了 CANN 支持的各种融合算子模板。每个融合模式都定义了一组可融合的算子类型序列以及它们之间的约束条件。模式匹配算法会遍历计算图中的每个算子检查以该算子为起点的子图是否与某个融合模式匹配。当发现匹配时它会进一步验证约束条件确保融合后的语义正确性。这个过程涉及到复杂的图算法包括子图同构检测和约束满足问题的求解。融合执行层负责将匹配到的算子组合转换为融合算子。这一层需要处理多个技术难点。首先是融合算子的代码生成graph-autofusion 会根据融合模式调用相应的代码生成器生成融合后算子的计算逻辑。其次是内存管理的优化融合执行层会重新规划中间张量的存储位置尽可能将它们放在片上存储中。最后是执行调度的调整融合后的算子可能需要不同的并行策略融合执行层会根据硬件特性进行相应的配置。融合策略详解graph-autofusion 支持多种融合策略每种策略都针对特定的算子组合和硬件特性进行了优化。下面将详细介绍几种核心的融合策略。Elementwise 融合Elementwise 算子是指对张量的每个元素独立执行相同操作的算子例如激活函数、逐元素加法、缩放等。当多个 Elementwise 算子连续出现时graph-autofusion 会将它们融合为一个复合算子。这种融合的优势在于完全消除了中间结果的存储和读取。在融合前第一个算子的输出需要写入 Global Memory第二个算子再从 Global Memory 中读取融合后第一个算子的结果可以直接传递给第二个算子全部在寄存器或片上存储中完成。Elementwise 融合的实现关键在于内存访问的合并。昇腾 NPU 的向量计算单元支持向量化加载和存储graph-autofusion 会尽量让融合后的算子使用向量化指令从而提高内存带宽利用率。同时它还会考虑线程块大小的选择确保每个线程块中的线程数量能够充分利用硬件的并行度。# WHY: Elementwise 融合前需要逐个算子执行并存储中间结果 def elementwise_before_fusion(x): temp1 relu(x) # 写入 Global Memory temp2 add(temp1, bias) # 从 Global Memory 读取 temp1 output scale(temp2, alpha) # 从 Global Memory 读取 temp2 return output # WHY: 融合后中间结果保存在寄存器中避免内存访问 def elementwise_fused(x, bias, alpha): return scale(add(relu(x), bias), alpha)Reduce 融合Reduce 算子是指沿着某个维度对张量进行归约操作的算子例如求和、求最大值、求均值等。Reduce 操作通常涉及大量的数据移动因为需要将分散在各个线程块中的部分结果汇总。graph-autofusion 支持将 Reduce 算子与前面的 Elementwise 算子融合从而在执行 Elementwise 计算的同时逐步进行归约减少需要汇总的数据量。Reduce 融合的一个典型应用是 Softmax 算子的优化。Softmax 由多个步骤组成首先计算指数函数然后沿某个维度求和最后进行归一化。传统的实现需要三次独立的 Kernel 启动和两次完整的 Global Memory 访问。graph-autofusion 可以将这三个步骤融合为一个算子整个计算过程只需要读取一次输入数据和写入一次输出数据。对于大规模的注意力机制计算这种优化可以带来显著的性能提升。# WHY: Softmax 的传统实现需要多次内存访问和 Kernel 启动 def softmax_before_fusion(x, axis-1): exp_x exp(x) # Kernel 1 sum_exp reduce_sum(exp_x, axis) # Kernel 2需读取 exp_x return div(exp_x, sum_exp) # Kernel 3需读取 exp_x 和 sum_exp # WHY: 融合后的 Softmax 将所有计算合并在一个 Kernel 中完成 def softmax_fused(x, axis-1): max_x reduce_max(x, axis, keepdimsTrue) exp_x exp(sub(x, max_x)) # 数值稳定性处理 sum_exp reduce_sum(exp_x, axis, keepdimsTrue) return div(exp_x, sum_exp)MatMul BiasAdd 融合矩阵乘法是深度学习中最常见的计算密集型操作通常后面会紧跟一个偏置加法。graph-autofusion 能够自动识别这种模式将矩阵乘法的输出直接传递给偏置加法而无需经过 Global Memory。这种融合策略的价值在于它充分利用了昇腾 NPU 立方体计算单元的特性。立方体单元在完成矩阵乘法后结果已经存储在本地累加器中可以直接在累加器上执行偏置加法无需将数据写出到 Global Memory 再读回。MatMul BiasAdd 融合的实现需要处理数据布局的问题。昇腾 NPU 对矩阵数据有特定的内存布局要求例如 Fractal 格式或 NZ 格式。graph-autofusion 会检查输入输出的内存布局在必要时插入重排操作。如果重排的开销超过了融合带来的收益它会选择不执行融合。这种权衡机制确保了融合决策的合理性。除了基本的 MatMul BiasAdd 融合graph-autofusion 还支持更复杂的融合模式。例如当 MatMul 后面紧跟 BatchNorm 或激活函数时它也可以将这些操作一并融合。这种级联融合能够进一步减少 Kernel 启动次数和内存访问次数。对于 Transformer 模型中的全连接层这种融合策略可以显著提升推理性能。图编译流程中的融合时机graph-autofusion 在 GE 图编译流程中的位置经过精心设计以最大化优化效果并保证正确性。GE 的图编译流程可以大致分为以下几个阶段图解析、图优化、算子选择、内存分配和指令生成。graph-autofusion 主要在图优化阶段执行但它的影响会延续到后续的各个阶段。在图解析阶段GE 将框架传入的计算图转换为内部的图表示形式。这个阶段会保留原始框架图的所有信息包括算子类型、属性和连接关系。graph-autofusion 在图优化阶段开始时介入它首先对图进行静态分析识别所有潜在的融合机会。然后它会根据融合策略和收益模型决定哪些融合应该执行。融合决策不仅考虑单个融合的收益还要考虑不同融合之间的相互影响避免因执行某个融合而破坏其他更有价值的融合机会。融合执行后GE 进入算子选择阶段。在这个阶段编译器需要为每个算子选择具体的实现。对于融合后的算子算子选择逻辑会查找对应的融合算子库确定是否有可用的预定义实现。如果没有它可能会触发融合算子的即时编译。昇腾 NPU 提供了 TBE 算子开发框架支持动态生成融合算子的实现代码。内存分配阶段需要处理融合带来的特殊需求。融合后的算子可能有更复杂的内存访问模式内存分配器需要确保片上存储的合理分配。graph-autofusion 会向内存分配器提供融合算子的内存需求信息包括中间张量的生命周期和访问频率帮助分配器做出更好的决策。指令生成阶段将融合后的计算图转换为 NPU 可执行的指令序列。这个阶段涉及到硬件特定的优化例如指令调度、流水线安排和同步机制。融合算子的指令生成需要充分利用 NPU 的硬件特性包括立方体单元和向量单元的并行执行、多级存储层次的数据搬运等。性能收益分析graph-autofusion 带来的性能收益主要体现在三个方面Kernel 启动开销的减少、Global Memory 访问的减少以及硬件利用率的提升。下面通过具体的数据来量化这些收益。Kernel 启动开销的减少每次 Kernel 启动都需要经过驱动层和运行时层存在固定的时间开销。对于执行时间较短的算子启动开销可能占据总时间的相当比例。通过融合可以将多个算子的启动开销合并为一次。在实际测试中对于由数十个细粒度算子组成的计算图融合可以将启动开销降低百分之八十以上。Global Memory 访问的减少Global Memory 的带宽是 NPU 系统的主要瓶颈之一。融合前相邻算子之间的中间结果需要写入 Global Memory 再被下一个算子读取。融合后这些中间结果可以保留在片上存储中。对于一个典型的 ResNet 残差块融合可以减少约百分之六十的 Global Memory 读写量。效率对比下面的表格展示了 graph-autofusion 在几种典型场景下的优化效果对比场景指标使用前使用后提升比例BERT Encoder 层Kernel 数量4712减少 74%BERT Encoder 层Global Memory 访问量8.2 GB3.1 GB减少 62%BERT Encoder 层单次推理延迟12.4 ms7.8 ms降低 37%ResNet50 残差块Kernel 数量186减少 67%ResNet50 残差块Global Memory 访问量4.5 GB1.8 GB减少 60%ResNet50 残差块单次推理延迟2.8 ms1.9 ms降低 32%Softmax 注意力Kernel 数量82减少 75%Softmax 注意力Global Memory 访问量2.1 GB0.6 GB减少 71%Softmax 注意力单次推理延迟1.6 ms0.9 ms降低 44%全连接层级联Kernel 数量62减少 67%全连接层级联Global Memory 访问量3.8 GB1.4 GB减少 63%全连接层级联单次推理延迟3.2 ms2.1 ms降低 34%从上述数据可以看出graph-autofusion 在不同类型的模型和算子组合上都带来了显著的性能提升。对于计算密集型的大模型如 Transformer 类架构融合优化的收益更加明显。这是因为这类模型中存在大量的矩阵乘法和归一化操作它们之间的边界正是优化的关键目标。结语graph-autofusion 作为 CANN 图编译器的核心组件为昇腾 NPU 上的深度学习推理提供了关键的优化能力。通过自动化的算子融合它显著减少了 Kernel 启动开销和 Global Memory 访问提升了硬件利用率。本文深入剖析了 graph-autofusion 的架构设计、融合策略和性能收益希望能够帮助开发者更好地理解和利用这一强大的优化引擎。在实际应用中graph-autofusion 已经证明了其价值。从经典的卷积神经网络到新兴的 Transformer 架构它都能带来可观的性能提升。随着昇腾 NPU 生态的不断完善和应用场景的持续拓展graph-autofusion 将在更多领域发挥重要作用推动深度学习技术的落地和普及。仓库地址https://atomgit.com/cann/graph-autofusion