TorchAir inductor-npu-ext 静态 Shape 融合算子 TopN 调优与 PGO 选优完整指南-育师

人工智能
模型编译
Ascend

【免费下载链接】torchair

TorchAir 支持用户基于PyTorch框架和torch_npu插件在昇腾NPU上使用图模式进行推理。

项目地址：https://gitcode.com/Ascend/torchair

导读

本文介绍 TorchAir 仓库中experimental/_inductor_npu_ext（AscendC 后端）在静态 shape 场景下的融合算子性能调优能力：编译期由 Autofuser（AF）生成 TopN 候选 tiling 解并逐一编译为静态 kernel，运行时（或 CppWrapper 的代码生成期）通过torch_npu.profiler对候选做 PGO（Profile-Guided Optimization）实测选优，最终把最优解固化进缓存，后续调用零 profiling、零运行时 tiling 计算开销。读完本文，你将掌握该能力的三环节原理（TopN 候选生成 → PGO 实测选优 → 缓存）、触发条件、TORCHINDUCTOR_ASCENDC_AUTOTUNE_TOPN等配置用法、CppWrapper 集成方式、设备控核联动以及调试与排障手段。

说明：文中涉及的代码文件路径均相对experimental/_inductor_npu_ext/，文章末尾给出了从仓库根目录出发的完整相对路径索引。

1 能力定位与术语

1.1 两条编译路径

AscendC 后端的编译单元是融合 kernel（FusedASCGraph）。按 shape 是否静态，融合 kernel 有两条编译路径：

路径	适用场景	tiling 时机	说明
jit 路径	动态 shape（kernel 参数含`SizeArg`）	运行时 host 侧执行`AutofuseTiling`	每次下发时计算 tiling（`TASK_QUEUE_ENABLE=2`时在队列任务内计算），灵活但有 host 开销
静态 topn 路径	静态 shape	编译期确定	tiling 固化进静态 kernel，消除运行时 tiling 计算开销，并支持调优选优

静态 shape 调优能力构建在 topn 路径之上，由三个环节组成：

TopN 候选生成（编译期）：调用 Autofuser 生成的GenerateTopnSolutions，产出最多 10 组候选 tiling 解（tiling_repr/workspace/block_dim），每个候选编译为独立的静态kernel.so；
PGO 实测选优：PythonWrapper 在首次调用时、CppWrapper 在代码生成期用真实输入在设备上运行全部候选，通过torch_npu.profiler采集真实执行耗时，选出最优候选；
缓存（后续调用/运行）：选优结果落盘optimize_tiling，后续调用直接加载最优静态 kernel，不再重复 profiling。

在 README.md 的特性规划中，“静态化编译支持”与“PGO 性能寻优”即对应本能力，“基于硬件性能建模的模板&Tiling寻优”对应 TopN 候选生成环节。README 同时给出启用方式：已随 torch_npu 打包为torch_npu/_inductor/ascendc，通过torch.compile(options={"npu_backend": "ascendc"})启用，也可在import torch前设置环境变量TORCHINDUCTOR_NPU_BACKEND="ascendc"。

1.2 术语说明

术语	含义
静态 shape	融合 kernel 的参数中不含动态尺寸符号（`SizeArg`），即`_is_static == True`
TopN	编译期让 AF 生成候选 tiling 解的机制；扩展固定请求最多 10 组，实际保留数量由`TORCHINDUCTOR_ASCENDC_AUTOTUNE_TOPN`控制
PGO	Profile-Guided Optimization，通过 profiler 采集候选 kernel 的真实执行性能来指导选优
autotune	候选实测选优过程：对候选逐一 profiling、比较并缓存最优解；PythonWrapper 在首次调用时执行，CppWrapper 在代码生成期执行
kernel_key	候选的唯一标识，`md5(tiling_repr\|workspace\|block_dim)`，wrapper 据此索引 tiling 信息与静态 kernel
optimize_tiling	选优结果缓存文件，内容为`tiling_repr\|workspace\|block_dim`

2 触发条件

topn 路径与 autotune 的使能条件是逐融合 kernel 判定的，判定函数为 compiler/_compiler.py 中的decide_use_topn_tiling（源码注释明确其为 jit/topn 两条分发路径的“单一判定来源” single source of truth）：

条件	判定依据	说明
静态 shape	`is_static == True`	常规融合 kernel：参数中无`SizeArg`（npu.py 中`_is_static = not any(isinstance(v, SizeArg) for v in self._fused_graph.args)`）；独立转置 kernel：无动态符号
CANN 版本	`>= 9.1.0`	`GenerateTopnSolutions`接口要求，对应`is_gte_cann_version("9.1.0")`检查
host 代码能力	host_impl 含`GenerateTopnSolutions`与`GetTilingDataRepr`	由 AF 生成，代码侧据此探测（`_can_use_topn_tiling`）
非 matmul template	`has_matmul_template != True`	matmul template 走专用路径

以上全部满足即走 topn 编译路径。是否进一步触发 autotune 取决于 topn 配置（源码中_compiler.py第 846~848 行对候选做tiling_infos[:autotune_topn]裁剪）：

autotune_topn > 1：请求生成多个候选；实际候选 >1 时返回NpuInductorAutoTuneKernel，PythonWrapper 在首次调用时、CppWrapper 在代码生成期触发 PGO autotune（若 AF 实际只给出 1 组解，则退化为下述行为）；
autotune_topn == 1（默认）：候选裁剪为 1 个，直接把该候选写入optimize_tiling缓存并返回普通NpuInductorKernel——有静态化收益，但没有PGO 实测选优；
动态 shape kernel 一律走 jit 路径，不参与本能力。同一张图中静态/动态 kernel 混合时，各自走对应路径，互不影响。

注意：topn 配置只影响候选裁剪数量与是否 autotune，不影响 topn 路径本身的判定。topn 路径内部生成候选时固定请求AUTOTUNE_TOPN=10个解，再裁剪到用户配置值（见 3.2 节）。

3 编译期：TopN 候选生成

3.1 编译产物

与 jit 路径（kernel.so+wrapper.so两个产物）相比，topn 路径的产物结构为：

<缓存根目录>/<融合kernel>/<artifact_hash>/ ├── asc_tiling.py # tiling 编译脚本（tiling_def + host_impl） ├── tiling.so # tiling 编译产物，导出 AutofuseTiling / GenerateTopnSolutions / GetTilingDataRepr ├── tiling_wrapper.so # TopN 获取 wrapper，导出 init_tiling / get_topn / set_device_limit ├── wrapper.so # 运行时 wrapper，导出 init / wrapper / release / set_device_limit ├── static_kernel/ │ └── <kernel_key>/ │ ├── asc_kernel.py # 该候选的静态 kernel 编译脚本 │ └── kernel.so # 每个候选一个静态 kernel ├── optimize_tiling # 选优结果缓存；候选数为 1 或 PGO 成功后才生成 └── compile.lock # 编译/选优文件锁

调试目录：TORCH_COMPILE_DEBUG=1时另有asc_kernel_<kernel_key>.py等带 hash 的脚本。

缓存根目录由 config.py 的_get_asc_cache_dir()决定：复用原生 Inductor 的TORCHINDUCTOR_CACHE_DIR（默认/tmp/torchinductor_{用户名}），AscendC 产物落在其ascendc子目录；_get_lib_dir（compiler/init.py）以排除name/is_static/has_matmul_template后的源码 artifact 哈希为缓存 key。

3.2 编译流程

compile_ascendc（compiler/_compiler.py）中 topn 分支按依赖串行执行：

编译 tiling.so：由asc_tiling.py（setup_asc_host_compile_command）产出，包含 tiling 定义与GenerateTopnSolutions实现；
编译 tiling_wrapper.so：codegen_cpp_tiling_wrapper（codegen/_asc_codegen.py）生成 C++ 源码，dlopen tiling.so，向 Python 暴露init_tiling/get_topn/set_device_limit接口；
生成候选：get_topn通过 ctypes 调用get_topn(configs, 10, ...)，内部执行GenerateTopnSolutions，返回最多 10 组{tiling_repr, workspace, block_dim}；候选数量由 AF 决定，可能少于请求值；随后按autotune_topn裁剪；
生成 tiling_map 并编译 wrapper.so：每个候选计算kernel_key = md5(tiling_repr|workspace|block_dim)（对应get_kernel_hash），全部候选的kernel_key → TilingInfo映射注入 wrapper 源码（std::map<std::string, TilingInfo> tiling_map，见_compiler.py第 849~872 行）后编译；
逐候选编译静态 kernel：每个候选以固定 tiling_repr 走setup_asc_kernel_compile_command编译出static_kernel/<kernel_key>/kernel.so（对应调测产物asc_kernel_<hash>.py）；
返回值：候选数 >1 返回NpuInductorAutoTuneKernel（携带全部候选的KernelInfo列表，见_autotune.py）；候选数 ==1 直接写optimize_tiling并返回普通NpuInductorKernel（_compiler.py第 906~909 行）。

wrapper 运行时按kernel_key查tiling_map获取 tiling_data / workspace / block_dim 后直接下发静态 kernel，无需再执行 tiling 计算。

3.3 与异步编译的关系

async_compile（compiler/init.py）按路径分发：

jit 路径：wrapper / tiling obj / kernel obj 三路编译任务并行提交进程池，全部完成后回调链接 kernel.so；
topn 路径：存在 host→select→device 的依赖链，整体作为一个任务（compile_ascendc）提交进程池串行编译；
同步模式（TORCHINDUCTOR_COMPILE_THREADS=1或 CPU 调试）直接在当前进程执行compile_ascendc，两条路径逻辑一致（decide_use_topn_tiling是两条分发路径的单一判定来源）。

4 PGO 实测选优

4.1 PythonWrapper 触发时机

PythonWrapper 模式下，NpuInductorAutoTuneKernel首次被调用时（__call__ → load(args)）执行autotune()（compiler/_autotune.py）。全过程持有compile.lock文件锁，多进程/多线程并发场景下只有一个执行者。

4.2 选优流程

缓存检查：若optimize_tiling已存在，直接加载最优解（load_optimize_file对文件内容做 md5 得到 kernel_key），init 后立即返回，跳过后续所有步骤；
加载候选：dl.init逐个加载全部候选静态 kernel；
PGO 实测：
- 使用torch_npu.profiler，ProfilerLevel.Level1+AiCMetrics.PipeUtilization（见_autotune.py中_ExperimentalConfig构造）；
- schedule 为wait=0, warmup=1, active=1, repeat=1, skip_first=1，共执行 4 轮，每轮所有候选轮流下发并同步 stream，其中 1 轮为有效采集窗口；
- 输入为 PythonWrapper 首次调用时传入的实际参数（tensor 转 device 指针，标量按 ABI 转换）；CppWrapper 的前置 Python 编译阶段则使用 materialize 后的real_inputs；
- 源码级细节：为避免原地（in-place）kernel 在重复计时中写穿真实输入缓冲，benchmark 会先用clone_preserve_strides克隆输入、每轮恢复留底副本，再经_to_ctypes_arg转 ctypes 参数（与 compiler/_kernel.py 的下发 ABI 一致）；
解析耗时：profiling 结果落盘./profile_result/<md5>/，解析其中kernel_details.csv，按 kernel 名autofused前缀过滤，读取Duration(us)；耗时条数必须等于候选数，否则报错；
选优与缓存：取耗时最小的候选为最优（min(timings.items(), key=...)）；调用dl.release传入最优 kernel_key（wrapper 据此卸载其余候选，仅保留最优静态 kernel）；save_optimize_file写入tiling_repr|workspace|block_dim，随后run()一直以最优kernel_key下发（见_autotune.py第 167~174 行）。

日志会输出每个候选的耗时与最优 index（inductor_asclogger，info 级别，Autotune completed ...），可用于确认选优结果。

4.3 流程图

5 CppWrapper 模式下的集成

CppWrapper（TORCHINDUCTOR_CPP_WRAPPER=1或cpp_wrapper=True）下，autotune 被前移到代码生成期完成（cpp_wrapper.py 中 patch 的codegen_with_cpp_wrapper）：

autotune_topn <= 1时直接走 inductor 原生 cpp_wrapper 代码生成，不做任何额外处理；
autotune_topn > 1时：
- 先以 PythonWrapper 模式编译模块，并用真实输入（含编译期记录的参数）执行module.call(real_inputs)——此过程触发各融合 kernel 的 autotune/profiling；被 mutate 的输入先 clone，避免污染真实输入（源码中对self.mutated_inputs命中的输入执行clone_preserve_strides）；
- 记录未被真实执行覆盖的分支（存在候选但未经过 PGO 选优）到_ascendc_unvisited_topn，代码生成时这些 kernel 回退使用第一个候选（cpp_wrapper.py的_emit_compiled_kernel中对NpuInductorAutoTuneKernel且 wrapper 在 unvisited 集合内的场景取kernel_infos[0]）；
- 再以 cpp_wrapper 模式重新 codegen，选优出的最优候选的lib_kernel与kernel_key直接烘焙进 C++ 图代码，运行期零 Python 开销（C++ 侧通过_kernel_state_runtime_source生成的静态状态对象 dlopen wrapper.so，并在构造时init+release选中最优 kernel）。

因此在autotune_topn > 1的 CppWrapper 模式下，首次代码生成耗时包含一次整图的 Python 模式执行 + PGO 采集。扩展缓存命中时会跳过候选编译和 PGO，但若重新触发代码生成，仍可能执行用于生成 CppWrapper 的 Python 模式compile_to_module()/module.call()。

6 缓存与后续运行

选优缓存：optimize_tiling文件是 autotune 的最终结果。async_compile在编译入口即检查（compiler/__init__.py的async_compile）：文件存在且 wrapper/kernel 有效时，直接返回携带最优kernel_key的普通NpuInductorKernel，完全跳过候选编译与 profiling；
编译缓存：缓存根目录跟随原生 Inductor（TORCHINDUCTOR_CACHE_DIR，默认/tmp/torchinductor_{用户名}），AscendC kernel 落在其ascendc子目录下；跨 OS 共享文件系统场景建议按 OS 隔离，避免文件锁失效导致缓存损坏；
调试注意：TORCHINDUCTOR_FORCE_DISABLE_CACHES=1只控制 Inductor 自身缓存，不会绕过本扩展的ascendc缓存。要强制重新编译，应使用新的扩展缓存目录或清理对应扩展缓存；二次运行若要外挂性能采集工具，请保留缓存（见 8 节约束）。

7 设备资源限制（set_device_limit 联动）

torch.npu.set_device_limit（AIV/AIC 核数限制）会影响 tiling 搜索空间，因此需要保证候选生成与运行时看到同一份资源视图：

父进程在提交编译任务前读取并归一化控核限制（_read_device_limit_info，未显式设置过set_device_limit时视为无限制；归一化时以设备属性中的vector_core_num/cube_core_num为上限做 min 截断）；
控核信息随编译任务序列化传递给编译 worker（worker 不重复读取设备状态），避免进程池环境不一致（async_compile中的worker_limit = _read_device_limit_info()显式随任务下发）；
_inject_device_limit将限制注入 tiling_wrapper.so 与 wrapper.so 的set_device_limit接口，GenerateTopnSolutions与运行时下发使用一致的g_device_limit（ub_size传 0，后端视为使用默认 UB）。

静态 topn 的缓存 key 包含vector_core_num/ai_core_num（_get_lib_dir中is_static时不排除这两个字段），因此不同控核配置的静态缓存相互隔离；而动态路径按 tiling 运行时重算，generic .so 可跨核数共享。这一行为在 tests/smoke/test_device_limit_smoke.py 中有对应测试：test_core_limit_is_part_of_cache_key（静态缓存 key 按核数隔离）、test_cache_reuse_across_core_num（动态缓存跨核数复用）、test_autotune_candidates（autotune 候选的 block_dim 受控核限制）、test_static_topn_block_dim_limited（profiler 验证 Block Num <= 8）等。

8 配置项与约束

8.1 配置项

配置	默认值	说明
`TORCHINDUCTOR_ASCENDC_AUTOTUNE_TOPN`	1（不使能 PGO）	调优最优解数量，正整数；上限 10，超限告警并最多返回 10 组；实际解数量与融合算子相关，可能小于配置值
`TORCHINDUCTOR_COMPILE_THREADS`	按 CPU 核数	设为 1 关闭异步编译（调试用）；topn 路径在异步模式下整体串行编译
`TORCHINDUCTOR_CACHE_DIR`	`/tmp/torchinductor_{用户名}`	原生 Inductor 缓存根目录，AscendC 编译/选优缓存落其`ascendc`子目录
`TORCHINDUCTOR_CPP_WRAPPER`/`cpp_wrapper=True`	关闭	CppWrapper 模式，autotune 前移至代码生成期

配置解析位于 config.py：TORCHINDUCTOR_ASCENDC_AUTOTUNE_TOPN非整数或小于 1 时分别在启动期抛出TORCHINDUCTOR_ASCENDC_AUTOTUNE_TOPN must be a positive integer / must be >= 1；_compiler.py中若配置值超过AUTOTUNE_TOPN=10会打印告警。使用手册 docs/manuals.md 给出示例：export TORCHINDUCTOR_ASCENDC_AUTOTUNE_TOPN=3使能 PGO 调优。

代码内另有AUTOTUNE_CONFIGS（compiler/_compiler.py，当前为空列表），用于向GenerateTopnSolutions透传 tiling 搜索配置项（如ub_threshold、corenum_threshold、enable_multicore_ub_tradeoff等，源码中留有注释示例），为后续搜索空间扩展预留；tests/smoke/test_device_limit_smoke.py的_worker_autotune展示了实际注入三组配置的做法。

8.2 约束与代价

首次编译耗时显著增加：相比 jit 路径多编译 tiling.so、tiling_wrapper.so，以及 N 个静态 kernel.so；
首次选优自动开启 profiling：PythonWrapper 在首次调用时、CppWrapper 在代码生成期执行 PGO；PGO 与外部性能采集工具互斥，嵌套场景现有 profiling 工具无法兼容；如需采集性能指标，保留编译缓存后在二次运行时开启；
候选数量不保证：实际候选数由 AF 的解空间决定，可能少于配置值；候选数为 1 时无选优过程；
适用范围：仅静态 shape 融合 kernel；动态 shape 走 jit 路径，未进入融合图的算子（aten fallback）不经过本编译路径，均不涉及；
profiling 中间产物：选优过程会在执行目录产生./profile_result/<md5>/临时目录，选优完成后自动清理。

8.3 错误提示

场景	异常/表现	处理建议
topn 配置非法	启动即报`TORCHINDUCTOR_ASCENDC_AUTOTUNE_TOPN must be a positive integer / must be >= 1`	修正为正整数
候选 tiling 生成失败	`NPU kernel {name} get topn failed`	检查 CANN 版本 >= 9.1.0 与 tiling.so 产物
profiling 解析失败	`Autotune {kernel_dir} parse profiling error, timing info count: N`	候选数与耗时条数不符，多为 profiler 未采集到 kernel，检查环境后清缓存重试
静态 kernel init 失败	`NPU kernel {lib_kernel} init failed`	检查缓存目录完整性，必要时清缓存重新编译

9 调试产物与运行观测

开启TORCH_COMPILE_DEBUG=1后，torch_compile_debug/<run_id>/torchinductor/model__*/autofused_*/下与 topn 调优直接相关的产物：

产物	说明
`asc_tiling.py`	tiling 实现 + tiling.so 编译脚本，可直接执行复现 tiling 编译
`asc_kernel_<hash>.py`	每个候选 tiling 的静态 kernel 编译脚本（hash 即 kernel_key）
`inductor_tiling_wrapper.cpp`	TopN 获取 wrapper 源码（get_topn/GenerateTopnSolutions 调用侧）
`inductor_wrapper.cpp`	运行时 wrapper 源码（tiling_map 注入后）
`benchmark.py`	融合 kernel 性能测试脚本，可单独复现与比对候选性能

运行期观测：autotune 选优结果与各候选耗时通过inductor_asc日志（info 级别）输出（_autotune.py中Autotune completed %s, best kernel index: %d及逐候选Timing = xx.xx us）；TORCHINDUCTOR_ASCENDC_DEBUG、layout_check等通用调测开关见使用手册 docs/manuals.md。

相关测试用例（可作为行为参考）：

tests/test_cpu_scalar_tensor_scalardata_st.py：test_add_cpu_scalar_tensor_static_kernel_path（topn=1 时静态 kernel 通过非空 kernel_key 下发）、test_add_cpu_scalar_tensor_topn_tiling_path（topn=2 的多候选编译执行）；
tests/smoke/cpp_wrapper_test.py：CppWrapper 下 topn 多候选编译、执行、输入突变（x.add_(1)原地算子场景）与分支切换（torch.cond未访问分支回退首候选）；
tests/smoke/test_device_limit_smoke.py：设备核数限制与候选 block_dim 的联动、静态/动态缓存隔离。

10 快速上手示例

在满足静态 shape、CANN >= 9.1.0 的前提下，可按下面方式使能 PGO 选优：

# 使能 PGO：请求最多 3 组候选 tiling 解（配置超 10 会告警并截断） export TORCHINDUCTOR_ASCENDC_AUTOTUNE_TOPN=3 # （可选）调试：同步编译 + 保存编译中间产物 export TORCHINDUCTOR_COMPILE_THREADS=1 export TORCH_COMPILE_DEBUG=1

import torch import torch_npu @torch.compile(options={"npu_backend": "ascendc"}) def test_add_sum(x, y): return torch.add(x, y).sum() x = torch.randn(32, 1024).npu() y = torch.randn(1, 1024).npu() out = test_add_sum(x, y) # 首次调用触发 PGO：编译候选 + 实测选优 + 落盘 optimize_tiling golden = (x + y).sum() assert torch.max(torch.abs(out - golden)) < 1e-3, torch.max(torch.abs(out - golden))

首次执行后，检查缓存目录（默认/tmp/torchinductor_{用户名}/ascendc/）下各融合 kernel 目录：存在optimize_tiling表示选优已落盘；static_kernel/<kernel_key>/kernel.so的个数即候选数。二次运行命中缓存后，不再编译候选、不再开启 profiling，直接以最优kernel_key下发静态 kernel。

附：源码与文档路径索引（相对仓库根目录）

设计文档：experimental/_inductor_npu_ext/docs/design/static_shape_autotune.md
使用手册：experimental/_inductor_npu_ext/docs/manuals.md
项目说明：experimental/_inductor_npu_ext/README.md
编译/调优核心实现：experimental/_inductor_npu_ext/python/inductor_npu_ext/compiler/_compiler.py、experimental/_inductor_npu_ext/python/inductor_npu_ext/compiler/_autotune.py、experimental/_inductor_npu_ext/python/inductor_npu_ext/compiler/_kernel.py
异步编译分发：experimental/_inductor_npu_ext/python/inductor_npu_ext/compiler/init.py
配置解析：experimental/_inductor_npu_ext/python/inductor_npu_ext/config.py
CppWrapper 集成：experimental/_inductor_npu_ext/python/inductor_npu_ext/cpp_wrapper.py
静态判定：experimental/_inductor_npu_ext/python/inductor_npu_ext/npu.py
测试用例：experimental/_inductor_npu_ext/tests/test_cpu_scalar_tensor_scalardata_st.py、experimental/_inductor_npu_ext/tests/smoke/cpp_wrapper_test.py、experimental/_inductor_npu_ext/tests/smoke/test_device_limit_smoke.py