TorchAir 的 inductor-npu-ext 精度对比与基准测试工具：从 benchmark.py 生成到 ASC kernel Profiling 的完整实践

导读

本文深入解析 TorchAir 仓库中 inductor-npu-ext（AscendC 后端）自带的精度对比与基准测试工具 precision_check.py。该工具为每个 AscendC 融合 kernel 生成独立的 benchmark.py，将 FX 片段、输入元信息和 ASC artifacts 固化到调试目录，用于单 kernel 级的精度复现、性能对比与 Profiling 定位。读完本文，你将掌握 --check / --bench / --profile / --e2e 四种运行模式的语义边界、benchmark 脚本的生成原理，以及如何用最小示例快速定位融合 kernel 的精度、性能与 Profiling 问题。

1 工具定位与整体设计

1.1 工具定位

precision_check.py（位于 python/inductor_npu_ext/common/precision_check.py）为每个 AscendC 融合 kernel 生成独立的 benchmark.py，将 FX 片段、输入元信息和 ASC artifacts 固化到调试目录中，用于复现单个 kernel 的精度、性能和 Profiling 问题。

工具包含三部分：

1. Codegen 阶段：从 scheduler/FX 节点收集生成脚本所需的信息。
2. 精度对比阶段：恢复 eager_fn，比较 eager 与重新编译后的 FX 片段。
3. ASC Profiling 阶段：编译并调用当前 kernel artifacts，采集 NPU 设备数据。

从源码结构看（generate_benchmark_script() 内部先调用 _build_benchmark_spec() 收集信息，再由 _write_benchmark_script() 统一写文件），Codegen 阶段与脚本写入阶段被刻意解耦：所有对 Inductor 图内部对象的访问都收敛在 _build_benchmark_spec() 内，writer 只读取 BenchmarkSpec，这保证了生成逻辑的稳定与可测试性。

1.2 运行模式

模式	执行对象	主要输出	说明
--check	eager_fn vs torch.compile(eager_fn)	CANN 精度指标、退出码	不直接验证当前目录的 ASC artifact
--bench	上述两条路径	设备耗时、Ratio、Avg	比较重新编译片段的性能
--profile	async_compile_ascendc(artifacts)	profiling/	验证和 Profiling 实际 ASC kernel，默认模式
--e2e	与三种模式组合	使用内联的 asc_graph.py 源码	源码在生成 benchmark.py 时内联

有完整 FX 精度上下文时，脚本支持前三种模式；FX 片段无法恢复时自动生成 profile-only 脚本，只保留 --profile。

1.3 总体流程

NpuKernel.benchmark(nodes, artifacts, file_path)
                  │
                  ▼
       generate_benchmark_script()
                  │
       ┌──────────┴──────────┐
       ▼                     ▼
  提取 FX 片段成功       提取失败/无法序列化
       │                     │
       ▼                     ▼
 写入 eager_fn、输入 meta   warning + profile-only
       │                     │
       └──────────┬──────────┘
                  ▼
           运行 benchmark.py
        ┌─────────┼─────────┐
        ▼         ▼         ▼
     --check   --bench   --profile

2 benchmark.py 生成机制

2.1 入口和调用链

入口位于 python/inductor_npu_ext/npu.py 的 NpuKernel.benchmark()：

def benchmark(self, nodes, file_path=None):
    from .common.precision_check import generate_benchmark_script

    return generate_benchmark_script(
        nodes=nodes,
        kernel_name=self.kernel_name,
        artifacts=self._artifacts,
        args=self.args,
        file_path=file_path,
        is_static=self._is_static,
        prof_size_vars=getattr(self.fused_graph, "prof_size_vars", []),
        cpp_wrapper=bool(V.graph.cpp_wrapper),
    )

generate_benchmark_script() 返回值为是否生成了 precision context（spec.precision is not None），调用方据此可判断当前 kernel 是否具备精度对比能力。prof_size_vars 会在 _build_benchmark_spec() 中被追加进 call_args，使 benchmark 调用点与 C++ wrapper 的签名对齐。

2.2 生成规格：PrecisionContext 与 BenchmarkSpec

PrecisionContext 和 BenchmarkSpec 将图访问与脚本写入分离：

数据结构	主要字段	用途
PrecisionContext	eager_fn_code、fx_input_meta、output_names	描述可执行的 FX 参考函数、输入 shape/stride/dtype 和输出顺序
BenchmarkSpec	kernel 名、artifacts、ASC 调用参数、符号 hint、profile buffers、precision	writer 只读该规格，不再访问 Inductor 图内部对象

脚本头部写入 kernel 名、节点摘要、ASC artifacts 字典和运行模式说明；输入 buffer 由 rand_strided 按保存的 layout 构造（来自 torch._dynamo.testing）。artifacts 字典包含 name、cpp_wrapper、cpp_tiling_wrapper、tiling_def、host_impl、device_impl 等字段，其中 vector_core_num、ai_core_num 来自设备核数限制信息，保证缓存键隔离。

2.3 环境变量快照：可复现性的关键细节

头部还会快照生成进程的 TORCHINDUCTOR_* 环境变量（_benchmark_env_lines() 实现），在 import torch 之前恢复。原因在于：这些变量在 import 期被 torch._inductor.config 和 inductor_npu_ext.config 读取，若不带回，--check 会按默认配置重新编译，无法复现整网运行时开启的融合（例如 TORCHINDUCTOR_ASCENDC_DEBUG=gather）。

两个变量被显式排除：

• TORCHINDUCTOR_NPU_BACKEND：已由脚本头部固定写为 ascendc（os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'ascendc'），不参与快照；
• TORCHINDUCTOR_CACHE_DIR：由 Inductor 自己回写（cache_dir() 每次运行都会解析并写入；force_disable_caches 下 with_fresh_cache_if_config() 会把它指向一次性的 mkdtemp 目录），不承载 codegen 决策。

快照以 _bench_env 字典加一段恢复循环的形式内联进脚本，而不是调用本包的 helper——因为本包在模块级就 import torch，而变量必须在 import torch 之前就位。恢复循环保持"命令行优先"的语义（等价于 setdefault），便于做对照实验；但对取值不同的变量会向 stderr 打印 WARNING 并给出生成期原值，避免"以为在复现、其实配置已被改掉"静默发生。值相同或变量未设置时不产生任何输出。

2.4 FX 片段提取

_extract_aten_fragment() 的处理步骤：

1. 对 scheduler nodes 聚合 origin（aggregate_origins），只保留 FX call_function 节点并按拓扑序排序。
2. 从 scheduler 的 read_writes 收集外部输入；只有所有依赖已经可用的节点才加入片段。
3. 优先通过消费 OUTPUT 的 scheduler buffer 确定真实输出（_collect_external_output_fx_names），并保持输出顺序。
4. 若无法找到外部输出，使用写入节点的叶子节点作为回退；仍无法确定时返回 None。
5. 将不在片段内的节点记录为 external inputs，交给 _extract_fx_input_meta() 和 _gen_eager_fn_code()。

这种方式只恢复当前融合 kernel 对应的 FX 子图，不重新执行整网模型。_buffer_origin_name() 负责把 Inductor IR buffer 映射回其 FX origin 节点名（分别处理 TensorBox/StorageBox、ComputedBuffer 等分支）。

2.5 输入和参数序列化

_extract_fx_input_meta() 从 node.meta['val'] 读取：

• shape：静态整数或通过 _size_hint() 得到的动态 shape hint；
• stride：保留非连续输入和广播 stride（如转置输入的 stride、广播维度的 stride 0）；
• dtype：写入 rand_strided；
• device：生成脚本统一使用当前 NPU 设备。

_serialize_fx_arg() 支持 FX Node、列表/元组、torch.device、NaN/Inf 和 Tensor 常量。多机生成的 npu:<rank> 会归一化为 npu，避免在另一张卡上运行时报设备不存在。具体地，三类参数会被归一化：torch.device 对象（如 device(type='npu', index=3)）、字符串 device（'npu:3'、'npu:0'）、以及 repr 内嵌 device='npu:3' 的 Tensor 常量。

关于动态 shape 的 _size_hint()，测试 test_dynamic_shape_no_recompile_with_debug 记录了一个重要回归点：早期实现用 int(SymInt) 取值，会经 guard_int -> evaluate_expr -> _set_replacement 把具体 shape 值写入 shape_env.replacements，导致每个新 shape 都触发 guard miss 而逐 shape 重编；修复后 _extract_fx_input_meta 改用 dim.node.expr + _size_hint（只读 hint），不修改 shape_env，dynamic=True 下整网只编译一张图。

节点 target 必须能序列化为 torch.* 调用；参数无法序列化、缺少 Tensor meta 或输出不完整时不生成半成品 eager_fn。测试 test_no_input_fragment_device_serialization 还覆盖了无外部输入节点的 FX 片段（如常量折叠后的 aten.full.default），修复后 NPU device 归一化为 torch.device('npu')，--check 正常通过。

2.6 profile-only 降级

以下情况会记录 warning 并降级：

• 没有可用的 FX origin 或 call_function 节点；
• 无法确定 graph outputs；
• 输入缺少 meta['val']、动态值无法得到 hint；
• 节点参数或 target 不支持序列化。

降级后仍写入完整 ASC artifacts 和 profile 入口，因此可以用 python3 benchmark.py --profile 定位真实 kernel 问题。测试 test_no_half_baked_eager_fn 专门看护"半坏脚本"：无论片段提取成功或失败，生成的 benchmark.py 必须语法正确（可 ast.parse），有 eager_fn 则必有 --check，无 eager_fn 则必须是纯 profile-only 脚本。

3 精度检查设计

3.1 --check 的验证边界

生成脚本中的两个执行函数为：

def _run_eager(inputs):
    return eager_fn(inputs)

_compiled_fn = torch.compile(
    eager_fn,
    backend="inductor",
    options={"npu_backend": "ascendc"},
    dynamic=not is_static,
)

def _run_compiled(inputs):
    with inductor_config.patch({"cpp_wrapper": <生成时的 cpp_wrapper 配置>}):
        result = _compiled_fn(inputs)
    return result

--check 比较的是恢复的 eager 片段与重新编译后的片段。它不读取当前目录 ASC kernel 的 device/host 实现；要验证该 artifacts 必须使用 --profile。这一点由测试 test_check_semantics_not_asc_artifact 直接验证：把 asc_graph.py 中的 ops.add 篡改为 ops.sub 后，--check 依然通过，因为 --check 根本不读 asc_graph.py 的 kernel 实现。

3.2 Tensor 比较指标

指标实现位于 compare_tensors()（precision_check.py），参考 CANN msaccucmp.py 精度标准。源码顶部的阈值常量即为判定基准：

指标	实现/含义	通过条件
CosineSim	sum((x/‖x‖) * (y/‖y‖))	> 0.99
RelEuclid	sqrt(sum((x-y)^2) / sum(y^2))	< 0.05
KLDiv	对两端分别 min-max 归一化后计算 sum(p*log(p/q))	< 0.005
RMSE/MAE	观察误差是否集中	> 1.1 只告警，不影响前三项判定

其中 x 为 expected、y 为 actual；浮点计算先转 float64 并展平。三个主指标必须同时满足，精度检查才通过。

源码中还有几个值得注意的细节：

• _EPSILON = torch.finfo(torch.float32).eps，用于范数、分母接近 0 时的保护分支；
• _align_nonfinite_positions() 实现 equal_nan 语义：当两侧在完全一致的位置携带 NaN/Inf（且 Inf 符号一致）时，这些位置从指标计算中剔除，比较仍可通过，并在结果行打印 [NOTE: ...]；若 NaN/Inf 位置或符号不一致，则直接判定 FAILED；
• RMSE/MAE 比值仅作为告警（[WARN: RMSE/MAE=..., deviation not concentrated]），不参与通过/失败判定，用于提示误差是否集中在少数离群点上。

3.3 输出和 dtype 处理

• 空 Tensor 直接通过（打印 precision check passed: empty tensor）；
• 量化 Tensor 先 dequantize()；
• 非浮点且非复数 Tensor 逐元素精确比较（int、bool 等必须完全一致，测试覆盖了 int32 dtype）；
• 非 Tensor 输出打印 skipped 并跳过；Tensor 输出按 output_names 顺序逐个比较（compare_outputs()）；
• 检查异常会被 run_check() 捕获，打印失败原因并以非零退出码结束（sys.exit(0 if _ok else 1)）。

4 Profiling 和性能对比

4.1 --profile：实际 ASC kernel

--profile 分支执行：

1. async_compile_ascendc(None, artifacts) 编译当前 ASC artifacts；
2. 根据 TensorArg layout 创建随机 strided NPU Tensor（rand_strided，按生成的 profile_buffers 规格）；
3. 调用 kernel，并使用 torch_npu.profiler.profile 采集 11 次循环（schedule(wait=0, warmup=0, active=10, repeat=1, skip_first=1)）；
4. 跳过首个 step，保留 10 个 active step 到 ./profiling。

脚本内联的 experimental_config 做了精简采集配置：export_type=[Text, Db]、profiler_level=Level2、msprof_tx=False、aic_metrics=AiCoreNone、l2_cache=False、op_attr=False、record_op_args=False 等，尽可能减少 Profiling 本身的开销与产物体积。

该模式使用 Text/Db profiler 输出和 AscendC 实际调用路径，适合验证 tiling、host wrapper 和 device kernel；配合 --e2e 可以验证修改后的 ASC 源码（见 4.3，修改需写入 benchmark.py 内联副本）。

4.2 --bench：eager 与 compiled 片段

run_bench() 对 _run_eager 和 _run_compiled 分别调用 run_profiling()：

_run_eager     → bench_eager_prof     → kernel_details.csv
_run_compiled  → bench_compiled_prof  → kernel_details.csv
                                      │
                                      ▼
                        按 Step Id 累加 Duration(us)
                                      │
                                      ▼
                         Graph Mode / Eager / Ratio

run_profiling() 使用 torch_npu.profiler.profile（CPU+NPU activity、schedule(wait=0, warmup=0, active=10, repeat=1)）采集 11 次迭代，on_trace_ready 为 tensorboard_trace_handler；结束后 torch.npu.synchronize() 并 sleep 等待落盘。

parse_device_times() 读取最新的 *_ascend_pt/ASCEND_PROFILER_OUTPUT/kernel_details.csv，要求存在 Step Id 和 Duration(us) 列，按 Step Id 累加设备耗时。print_bench_result() 最多输出 10 个 step，并输出平均值；Ratio 定义为 Eager / Graph Mode，大于 1 表示 compiled 路径更快。

注意：这里的 Graph Mode 是重新编译后的 eager_fn，不是 --profile 分支中的当前 ASC artifacts。

4.3 --e2e：内联 AscendC 源码

生成 benchmark.py 时读取同目录的 asc_graph.py（_write_artifacts_loading() 在 if args.e2e: 分支内将文件内容 splice 进脚本），将源码内联到脚本。运行时传入 --e2e 后，用内联的 tiling_def、host_impl 和 device_impl 覆盖 artifacts 字段。

因此，生成脚本后即使删除或修改旁边的 asc_graph.py，--profile --e2e 仍使用生成时保存的内容。测试 test_e2e_missing_asc_graph 与 test_e2e_corrupted_asc_graph 分别验证了删除与篡改 asc_graph.py 后 --e2e 依然正常。--check --e2e 和 --bench --e2e 可以使用同一命令行，但它们仍不直接验证原 ASC artifact。

5 限制和问题定位

5.1 输入与 shape 限制

• 动态 shape 只使用当前编译样例的具体 hint，不能代表所有 shape。
• benchmark 输入保留 shape/stride/dtype，但随机数据不等价于完整模型数据分布。
• profile-only 说明 FX 参考片段不可恢复，不表示 ASC kernel 本身一定失败。

5.2 常见产物和错误

现象	可能原因	建议
benchmark.py 没有 --check	FX 片段提取失败，生成 profile-only	查看 codegen warning，执行 --profile
--check 返回 1	指标未达阈值或 eager/compiled 抛异常	查看各输出指标和异常文本
--bench 找不到 CSV	CANN profiler 目录或字段格式不匹配	检查 *_ascend_pt、Step Id、Duration(us)
--profile 编译失败	ASC artifacts、CANN 或 wrapper 环境问题	直接执行 asc_graph.py/asc_kernel_*.py 复现

开启 TORCH_COMPILE_DEBUG=1 可查看 torch_compile_debug 下的 benchmark.py、asc_graph.py、ir_pre_fusion.txt 和 output_code.py。使用手册 docs/manuals.md 中同样强调：首次迭代数据可能包含编译开销，建议关注 Avg 行的 Ratio 值。

6 最小示例

6.1 生成 benchmark.py

以一个最小 pointwise 函数为例：

import torch


@torch.compile(options={"npu_backend": "ascendc"})
def fn(x, y):
    return (x + y).abs()


x = torch.randn(8, 16, device="npu")
y = torch.randn(8, 16, device="npu")
fn(x, y)

打开 TORCH_COMPILE_DEBUG=1 运行后，在 torch_compile_debug/.../torchinductor/ 目录下可以看到：

model__0_inference_0.0/
├── autofused_*/
│   ├── benchmark.py          # 本工具生成的独立复现入口
│   ├── asc_graph.py
│   ├── asc_kernel_*.py
│   └── inductor_wrapper.cpp 等
├── output_code.py
└── ir_pre_fusion.txt

其中 benchmark.py、asc_graph.py 位于 autofused_*/ 内，output_code.py、ir_pre_fusion.txt 位于上一级 model__*_inference_*.*/ 目录。多次执行相同脚本可能因缓存跳过编译，可配合 TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 强制每次重新编译（测试基类 _BenchmarkPrecisionBase.setUp 同样如此设置）。

6.2 最小精度对比

在 benchmark.py 所在目录执行：

python3 benchmark.py --check

输出示例：

output:abs precision PASSED: CosineSim=... RelEuclid=... KLDiv=... MAE=... RMSE=...
Precision check passed.

该命令比较恢复的 eager_fn 和重新编译后的 FX 片段；如果要验证目录中的实际 ASC artifacts，应使用 --profile。

6.3 最小性能对比

python3 benchmark.py --bench

输出示例：

Benchmark: <kernel_name>
 Step  | Graph Mode (us)    | Eager (us)    |  Ratio
-------+--------------------+--------------------+---------
  1   |        ...         |        ...         |  ... x
  2   |        ...         |        ...         |  ... x
  ...
  10   |        ...         |        ...         |  ... x
-------+--------------------+--------------------+---------
 Avg   |        ...         |        ...         |  ... x

step 序号从 1 开始，最多输出 10 个 step；脚本会分别生成 bench_compiled_prof 和 bench_eager_prof，再从 kernel_details.csv 汇总 Duration(us)。实际运行中可参考手册的示例数值（如 Avg 行 Ratio 2.18x），Ratio 大于 1 表示图模式更快。

6.4 最小 ASC Profiling

python3 benchmark.py --profile

该命令编译并调用当前 ASC artifacts，在当前目录生成 profiling/。需要验证修改后的 ASC 源码时，直接修改 benchmark.py 中 if args.e2e: 块内联的源码，再执行：

python3 benchmark.py --profile --e2e

此时使用的是生成脚本中内联的 ASC 源码；生成脚本后删除或修改旁边的 asc_graph.py 都不影响该命令（不会被读取），因此修改 asc_graph.py 本身不会生效。

6.5 profile-only 示例

当 FX 片段无法提取或输入元信息不完整时，生成的脚本只保留 ASC 运行入口：

python3 benchmark.py --profile

这类脚本没有 --check 和 --bench，但仍可用于确认 artifacts 是否能够编译、加载和执行。该降级只表示参考 FX 片段不可恢复，不等同于 ASC kernel 失败。

6.6 推荐使用顺序

先用 --check 确认 eager/compiled 片段的输出
        │
        ├── 需要比较耗时：--bench
        └── 需要确认当前 ASC kernel：--profile
                                      │
                                      └── 修改内联源码后：--profile --e2e

不要将 --check 的结果直接当作当前 ASC artifacts 的正确性结论，也不要用一个动态 shape 样例替代完整 shape 范围的验证。

7 测试保障与源码佐证

该工具的语义边界在仓库中有成体系的测试守护，位于 tests/test_benchmark_precision_st.py（NPU 端到端用例，@unittest.skipUnless(_has_npu, "需要真实 NPU 设备")）：

• 片段一致性看护：_assert_fragments_consistent 对比 asc_graph.py 的 "Graph fragment:" 段与 benchmark.py 中 eager_fn 的 aten/prims 调用序列必须完全一致；_assert_outputs_consistent 对比两者的 return 变量列表（含顺序）。这保证了从 scheduler 节点反向推导的 FX 片段与真正上板执行的 ASC 片段对应同一段计算逻辑。
• 多输出覆盖：中间值同时作为输出与后续输入、独立多输出、输出顺序与拓扑序不一致、四路 fanout 等场景均有专门用例（test_multi_output_*）。
• 输入构造覆盖：非连续输入（转置）、广播 stride [0, 1]、zero-size 输入、fp16/int32 dtype、torch.ops.prims.convert_element_type 等。
• 动态 shape 回归：dynamic=True + TORCH_COMPILE_DEBUG=1 下只编译一张图，且 benchmark.py 中 _fx_inputs 仍为具体 shape 值（来自 _size_hint）。
• 模式一致性：_assert_has_check_mode / _assert_profile_only 确保三模式与 profile-only 脚本的参数集各自完整、不产生半坏脚本。

结合 docs/manuals.md 第 4 节的使用指引，推荐的问题定位路径为：先跑默认 --profile 确认 kernel 能稳定运行；修改了 AscIR 或 Codegen 后用 --profile --e2e 验证融合 kernel 修改效果；--check 与 --bench 则用于验证融合 kernel 对应单算子片段在图模式下的行为。这一套"生成 → 校验 → 复现 → 修改 → 再验证"的闭环，正是融合 kernel 开发与精度问题定位的高效工作流。