CANN graph-autofusion:算子自动融合框架的设计思路

CANN graph-autofusion:算子自动融合框架的设计思路 文章目录前言先搞懂为什么算子需要融合graph-autofusion 的核心SuperKernel codegen JIT自动发现融合机会JIT 即时生成融合策略计算密集 访存密集算子配对在 CANN 五层架构中的位置与 ge 配合的编译流程与上下游仓库的协作关系上游ops-* 算子库上游ge 图引擎下游runtime 执行层关联pto-isa关联ascend-boost-comm一个融合过程的简化还原融合效果的验证方式写在最后前言“graph-autofusion 是自动调优工具吧”——如果你也这么想那这篇文章就是为你写的。graph-autofusion 不是调优不是调优不是调优。它是昇腾CANN生态中的算子自动融合框架干的是另一件事把多个独立的ops-*算子自动焊接成一个SuperKernel让昇腾NPU少跑几次内存搬运多算几轮实际数据。一个类比贯穿全篇算子是流水线上的工人graph-autofusion 是那个发现你俩的工序紧挨着合并成一个岗位更快的工业工程师。先搞懂为什么算子需要融合NPU 执行一个算子不是算完就完事——算完要把结果从计算单元搬回显存HBM下一个算子再从HBM搬进计算单元。这一来一回搬运耗时可能比计算本身还长。问题还不止于此每次算子切换NPU 的多个 AI Core 之间还需要同步屏障barrier等待所有核完成当前算子后才能启动下一个。数据搬运开销 核间同步开销构成了算子间缝隙的两大成本。# 未融合每个算子独立执行中间结果反复搬运 LayerNorm → 写回HBM → GELU → 读HBM → 写回HBM → 下一步 ↑ 开销1数据搬运 ↑ 开销2核间同步 # 融合后合成一个SuperKernel中间数据留在片上 LayerNorm → GELU → 写回HBM只写一次 ↑ 片上SRAM直传无HBM读写无核间同步不是搬得更快是不搬了。这就是融合的本质。对于 Transformer 类模型中常见的 LayerNorm→GELU、MatMul→BiasAdd→ReLU 这类短算子串联模式算子间搬运和同步的开销占比尤为突出——每个算子可能只需几微秒计算但光搬运和同步就要吃掉同样甚至更多的时间。# 一个简化的开销对比示意伪代码数值仅为说明比例# 未融合3个算子各自独立执行# MatMul: 计算 8us 搬运 5us 同步 3us 16us# BiasAdd: 计算 2us 搬运 5us 同步 3us 10us# ReLU: 计算 1us 搬运 5us 同步 3us 9us# 合计: 35us其中非计算开销占 57%# 融合为SuperKernel后# SuperKernel: 计算 11us 搬运 5us仅首尾各1次 16us# 非计算开销从 20us 降到 5usgraph-autofusion 的核心SuperKernel codegen JITgraph-autofusion 的核心机制是SuperKernel codegen JIT——运行时即时编译自动将多个ops-*算子融合为单一的SuperKernel执行体。自动发现融合机会框架不需要开发者手写融合规则它在图编译阶段遍历计算图自动识别可融合的算子组合。识别依据包括算子是否被标记为可融合Fusable属性、算子间是否存在数据依赖边连接、融合后数据流是否合法维度、类型兼容。JIT 即时生成融合后的SuperKernel代码在运行时生成而非预编译。JIT 的价值在于它能根据实际输入的形状shape和数据类型dtype做特化——比如针对[batch, 4096, 768]的输入生成与该形状完全贴合的计算代码省去通用逻辑中的分支判断和动态分发开销。// SuperKernel codegen JIT 的概念性流程// 1. 接收融合子图描述FusionPlan plananalyze_fusion_opportunity(subgraph);// 2. 根据实际输入形状做特化plan.specialize(input_shape,input_dtype);// 3. JIT 生成 SuperKernel 代码std::string kernel_codecodegen_jit(plan);// 4. 编译为可执行体ExecutableKernel super_kernelcompile(kernel_code);融合策略计算密集 访存密集算子配对graph-autofusion 的融合策略不是能融就融而是有选择地配对。核心思路是计算密集型算子与访存密集型算子搭配融合让 NPU 的计算单元和访存单元交替工作减少空闲等待。举例来说MatMul 是计算密集型算子大量乘加运算而 BiasAdd 和 ReLU 是访存密集型算子少量计算主要在做数据搬运。将它们融合后MatMul 产生的中间结果可以在片上直接被 BiasAdd 和 ReLU 消费访存密集算子的等待数据时间被计算密集算子的产出数据时间覆盖两者形成流水线式的交替执行。# fusion_rules.yaml —— 融合策略配置示例fusion_rules:-name:matmul_bias_relupattern:[MatMul,BiasAdd,ReLU]strategy:compute_bound memory_bound pairingconditions:-matmul_m_dim_ge:128# M维度足够大时融合收益明显-bias_dtype_eq:float16# 类型需兼容-name:layernorm_gelupattern:[LayerNorm,GELU]strategy:sequential_fuseconditions:-norm_hidden_dim_ge:64融合策略还考虑了算子的粒度边界——不是无限串联。如果融合链过长单个 SuperKernel 的寄存器和共享内存占用会超出片上资源上限反而导致性能回退。因此 graph-autofusion 内部有资源预算模型对融合深度做截断。// 融合深度截断的资源预算示意boolcan_fuse(constFusionPlanplan,constHWResourcehw){// 估算融合后 SuperKernel 的片上资源占用size_t reg_usageestimate_reg_usage(plan);size_t ub_usageestimate_ub_usage(plan);// Unified Buffer 占用// 超过硬件资源的 80% 即拒绝继续融合if(reg_usagehw.max_regs*0.8)returnfalse;if(ub_usagehw.max_ub*0.8)returnfalse;returntrue;}关键特征总结有三自动发现融合机会不需要开发者手写融合规则框架在图编译阶段自动识别可融合的算子组合JIT 即时生成融合后的SuperKernel代码在运行时生成而非预编译能根据实际输入形状做特化仅依赖Ascend C runtime轻量依赖不引入额外编译链路这里有个更关键的误区要纠正graph-autofusion 不是把算子粘在一起执行而是从代码层面重新生成了一个统一的计算体。就像两个工位合并后不是两个人挤在一起干活而是重新设计了一个更高效的工序。在 CANN 五层架构中的位置graph-autofusion 位于CANN五层架构的第3层——昇腾计算编译层与图引擎ge配合工作第1层AscendCL语言层 第2层AOL算子库 / AOE调优引擎服务层 第3层Graph Compiler / BiSheng / ATC编译层 ← graph-autofusion 在这里 第4层Runtime / Graph Executor执行层 第5层驱动 / 虚拟化基础层它在编译层的位置决定了它的工作时机——图编译时触发融合执行时已经是融合后的SuperKernel。这不是运行时的延迟优化而是编译期的提前规划。与 ge 配合的编译流程geGraph Engine在图编译阶段执行多轮优化遍passgraph-autofusion 作为其中一个优化 pass 介入。具体流程ge 先完成图的解析和基本优化如死代码消除、常量折叠然后调用 graph-autofusion 的融合 pass 遍历计算图识别可融合的子图并替换为 SuperKernel 节点。替换后的图继续进入后续编译阶段最终生成可执行的任务流。# 图编译触发融合的典型流程# Step 1: ATC 工具将模型转换为离线模型atc--mode5--modelmodel.onnx--outputmodel_fused# Step 2: 内部流程——ge 执行图优化# ge 解析图 → 常量折叠 → graph-autofusion融合pass → 后续编译 → 生成om模型# Step 3: 查看融合结果atc--mode1--ommodel_fused.om# 查看融合后的计算图结构# 通过 AscendCL Python 接口触发融合编译importacl# 初始化acl.init()contextacl.context()model_pathmodel_fused.om# 加载离线模型已包含融合后的SuperKernelmodel_idacl.mdl.load_from_file(model_path)# 执行时SuperKernel 作为单一算子被调度outputacl.mdl.execute(model_id,inputs)graph-autofusion 的融合 pass 是可配置的——用户可以通过环境变量或配置文件控制融合的激进程度比如是否允许跨分支融合、融合深度上限等。# 融合策略环境变量配置exportAUTOFUSION_ENABLE1# 启用自动融合exportAUTOFUSION_MAX_DEPTH4# 最大融合深度exportAUTOFUSION_STRATEGYbalanced# 策略balanced / aggressive / conservative与上下游仓库的协作关系graph-autofusion 不是孤立运作的它有自己的上下游依赖链上游ops-* 算子库ops-nn、ops-math 等算子库提供原始算子实现是 graph-autofusion 的原材料。每个算子在注册时会声明自己的融合属性——是否可融合、融合后需要哪些额外信息如中间张量的形状推导函数。graph-autofusion 正是通过这些属性来识别和筛选融合候选。// ops-* 算子的融合属性注册示意REGISTER_OP(LayerNorm).Input(x,TensorDesc(Shape({-1,-1}),ACL_FLOAT16)).Output(y,TensorDesc(Shape({-1,-1}),ACL_FLOAT16)).Attr(epsilon,AttrValue(1e-5)).Fusable(true)// 标记可融合.FusionOutputShape([](constShapein){// 融合后输出形状推导returnin;// LayerNorm 输出与输入同形状});上游ge 图引擎ge 在图编译阶段识别计算图中的融合机会触发 graph-autofusion。ge 负责在什么位置融合的决策graph-autofusion 负责怎么融合的实现。这种分工让融合逻辑与图优化逻辑解耦——ge 不需要知道 SuperKernel 的代码生成细节graph-autofusion 也不需要关心图的全局优化策略。下游runtime 执行层融合产物SuperKernel交给 runtime 调度执行。对 runtime 来说SuperKernel 与普通算子没有本质区别——它也是一个可调度的执行单元只是内部逻辑更复杂、执行时间更长。runtime 无需感知融合的存在这保证了融合框架与执行框架的解耦。关联pto-isapto-isa 是 PTO 虚拟指令集架构它定义了昇腾 NPU 上算子执行的指令抽象。graph-autofusion 生成的 SuperKernel 最终通过 PTO 指令序列描述其计算逻辑。PTO 作为虚拟指令集架构非底层硬件指令集为 SuperKernel 提供了跨硬件代际的可移植性——同一份融合产物可以在不同型号的昇腾 NPU 上执行。关联ascend-boost-commascend-boost-comm 是算子公共平台负责算子注册、发现和基础能力共享。graph-autofusion 通过 ascend-boost-comm 的注册表发现可融合算子也通过它查询算子的元信息计算类型、访存特征、资源占用估算来辅助融合决策。ops-* 算子库 ──→ graph-autofusion ──→ runtime 执行 ↑ │ ge 触发融合 pto-isa 指令生成 │ ascend-boost-comm算子注册一个融合过程的简化还原假设ge发现一个计算子图MatMul → BiasAdd → ReLUgraph-autofusion 的处理逻辑大致如下# 1. ge 识别可融合子图subgraph[MatMul,BiasAdd,ReLU]# 2. graph-autofusion 检查融合可行性# - MatMul: 计算密集, fusableTrue# - BiasAdd: 访存密集, fusableTrue# - ReLU: 访存密集, fusableTrue# - 资源预算: 预估reg/ub占用在阈值内 → 可以融合# 3. graph-autofusion JIT 生成 SuperKernelsuper_kernelcodegen_jit(subgraph)# 内部将三个算子的计算逻辑合并为单一执行体# 中间结果 bias_add_output 不再写回HBM# 4. 编译为PTO指令序列pto_instructionscompile_to_pto(super_kernel)# 5. 交由 runtime 调度执行runtime.launch(super_kernel,inputs,outputs)融合省的不是计算量是搬运量。三个算子变一个HBM读写从6次降到2次——这才是SuperKernel的核心收益。融合效果的验证方式graph-autofusion 提供了融合结果的 profiling 工具帮助开发者确认融合是否生效以及收益大小# 使用 msprof 工具采集融合前后的性能数据msprof--application./inference_app--output./prof_data# 查看融合结果# prof_data 目录下会生成算子级的时间线# 融合后MatMulBiasAddReLU 出现为单一SuperKernel节点# 未融合三个算子分别出现在时间线上中间有HBM读写间隔# 通过环境变量开启融合日志确认融合决策exportAUTOFUSION_LOG_LEVELDEBUG# 日志输出示例# [AUTOFUSION] Found fusable subgraph: MatMul → BiasAdd → ReLU# [AUTOFUSION] Resource estimate: reg62%, ub45% → within budget# [AUTOFUSION] SuperKernel generated: matmul_bias_relu_fused写在最后graph-autofusion 做的事情很聚焦在编译期自动发现融合机会用JIT生成SuperKernel减少NPU与HBM之间不必要的搬运。它不调优参数不搜索配置空间——它是一个焊接工程师不是调音师。如果你想看看它是怎么做的仓库在这里https://atomgit.com/cann/graph-autofusion下一步建议clone 仓库后关注Autofuse组件和SuperKernel codegen的实现逻辑理解它如何识别融合模式、如何生成融合代码——这才是这个框架真正的技术内核。