背景概述

PyTorch 作为主流的深度学习框架，其算子代码生成（Codegen）机制是支撑多后端扩展的核心基础设施。原生 PyTorch 通过 torchgen 统一生成各后端的算子分发与注册代码，而 NPU（Neural Processing Unit）适配层则需要在继承这一机制的基础上，针对 NPU 的硬件特性和软件栈进行深度定制。

本文从架构设计、数据模型、配置规范、代码生成逻辑四个维度，系统解析 torch_npu 对原生 Codegen 机制的适配与扩展，重点剖析 NPU 在算子分发策略、版本化管理及结构化代码生成方面的创新设计，为理解 NPU 算子适配体系提供技术参考。

一、架构差异

维度原生 PyTorch (torchgen)torch_npu
代码位置单一目录 torchgen/两个目录：op-plugin/codegen + torchnpugen
职责划分统一生成所有后端代码torchnpugen：后端分发注册；op-plugin/codegen：算子接口层
与 torchgen 关系自身就是框架继承并猴子补丁 torchgen（apply_torchgen_patch()）
DispatchKeyCPU 等标准 keyPrivateUse1 重命名为 NPU

架构图对比

原生 PyTorch：

torchgen/
├── model.py              # NativeFunction 等核心数据模型
├── dest/                 # 代码生成目标
│   └── register_dispatch_key.py  # 分发注册代码生成
├── api/                  # C++/Native/Dispatcher 签名生成
└── gen.py                # 主入口

torch_npu：

torchnpugen/              # 后端分发注册（扩展 torchgen）
├── gen_backend_stubs.py  # 主入口，apply_torchgen_patch()
├── custom_functions.py   # 自定义算子注册
├── gen_functionalization_type.py
├── gen_npu_c_shim.py
└── utils.py              # gen_unstructured 等替换逻辑

op-plugin/codegen/        # 算子接口层（独立于 torchgen）
├── gen.py                # 双路径分发代码生成
├── model.py              # NPU 扩展的 NativeFunction
├── struct/               # 结构化 aclnn kernel 生成
│   ├── struct_codegen.py
│   └── model.py
├── gen_op_plugin_functions.py  # 版本化 YAML 预处理
└── gen_derivatives.py    # 版本化 derivatives 过滤

二、数据模型差异

原生 NativeFunction

文件：torchgen/model.py

class NativeFunction:
    func: FunctionSchema
    use_const_ref_for_mutable_tensors: bool
    device_guard: bool
    device_check: DeviceCheckType
    python_module: str | None
    category_override: str | None
    variants: set[Variant]
    manual_kernel_registration: bool
    manual_cpp_binding: bool
    loc: Location
    autogen: list[OperatorName]
    ufunc_inner_loop: dict[UfuncKey, UfuncInnerLoop]
    structured: bool
    structured_delegate: Optional[OperatorName]
    # ... 等完整字段

NPU 扩展 NativeFunction

文件：op-plugin/codegen/model.py

class NativeFunction:
    func: FunctionSchema
    impl_name: Optional[str]              # ← NPU 新增：实现函数名
    impl_ns: Optional[str]                # ← NPU 新增：实现命名空间 (acl_op/op_api)
    sparse: Optional[str]                 # ← NPU 新增：稀疏支持版本范围
    internal_format_opapi: Optional[str]  # ← NPU 新增：内部格式 opapi 版本
    use_const_ref_for_mutable_tensors: bool
    structured: bool
    structured_delegate: Optional[OperatorName]
    # 以下原生字段被注释掉：
    # device_guard, device_check, python_module, category_override,
    # variants, manual_kernel_registration, manual_cpp_binding,
    # loc, autogen, ufunc_inner_loop

关键区别

字段原生NPU说明
impl_ns无acl_op / op_api标记双路径实现，是 NPU 双路径分发的核心
impl_name无可选覆盖默认的 kernel 名
sparse无版本范围字符串控制稀疏算子的版本支持
internal_format_opapi无版本范围字符串控制内部格式（5D等）的 opapi 支持
device_guard有注释掉NPU 在 codegen 层面自行处理
variants有注释掉NPU 不需要 function/method 区分
ufunc_inner_loop有注释掉NPU 不使用 ufunc 机制

三、YAML 配置差异

原生 native_functions.yaml

- func: add.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!)
  structured: True
  structured_inherits: TensorIteratorBase
  dispatch:
    CPU, MPS: add_out
    SparseCPU: add_out_sparse

NPU op_plugin_functions.yaml

- func: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  acl_op: all_version             # ← NPU 新增：acl_op 实现版本范围
  op_api: all_version             # ← NPU 新增：op_api 实现版本范围

- func: abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)
  acl_op: all_version
  op_api: all_version
  gen_opapi:                      # ← NPU 新增：结构化 aclnn 代码生成配置
    out:
      size: self                  # 输出 size = self.sizes()
      dtype: self                 # 输出 dtype = self.scalar_type()
      name: self                  # name propagation 源
    exec: aclnnAbs, self, out     # aclnn 调用映射

关键区别

配置项原生NPU说明
dispatchDispatchKey → kernel 名无原生按后端分发，NPU 按路径分发
acl_op无版本范围ACL Device 侧实现，版本范围如 all_version / [v2.1, newest]
op_api无版本范围aclnn Host 侧实现
gen_opapi无结构体描述 aclnn 调用参数映射和输出张量信息
structured_inherit无函数名op_api 中 inplace 委托给 out 实现
sparse无版本范围稀疏算子版本控制
internal_format_opapi无版本范围内部格式 opapi 版本控制

版本范围格式

acl_op: all_version          # 所有版本支持
acl_op: [v2.1, newest]       # v2.1 及以上支持
op_api: v2.1                 # 仅 v2.1 支持

四、代码生成逻辑差异

4.1 Unstructured Kernel 生成

原生生成代码

文件：torchgen/dest/register_dispatch_key.py

生成逻辑简单直接：设备检查 → 设备守卫 → 调用 kernel。

namespace {
Tensor wrapper_CPU_add(const Tensor& self, const Tensor& other,
                        const Scalar& alpha) {
  // No device check
  const OptionalDeviceGuard device_guard(device_of(self));
  return at::cpu::add(self, other, alpha);
}
}
NPU 生成代码

文件：torchnpugen/utils.py

生成逻辑复杂：设备检查 → 内存重叠检查 → Profiler 守卫 → OpHook → 双路径分发。

namespace {
Tensor wrapper_NPU_add(const Tensor& self, const Tensor& other,
                        const Scalar& alpha) {
  // 1. Device check
  const OptionalDeviceGuard device_guard(device_of(self));
  // 2. UnsafeTensorCheck（内存重叠检查）
  at_npu::native::OpPreparation::check_tensor_mem_overlap(...);
  // 3. NPURecordFunctionGuard（Profiler 守卫）
  at_npu::native::NPURecordFunctionGuard guard;
  // 4. OpHook 前置检查
  if (C10_UNLIKELY(at_npu::native::env::CheckOpHookEnable())) {
      at_npu::native::OpHook::GetInstance().PreHook("add.Tensor", ...);
  }
  // 5. 双路径分发：op_api vs acl_op
  if ((ForceAclnn || at_npu::native::env::CheckJitDisable())
      && FormatHelper::IsOpInputBaseFormat(self)) {
      return at_npu::native::NPUNativeOpApiFunctions::add(self, other, alpha);
  } else {
      if (c10_npu::IsAclnnOnly()) {
          TORCH_CHECK(false, "Unsupported operator execution path");
      }
      return op_plugin::add(self, other, alpha);
  }
}
}
NPU 新增逻辑详解
特性代码位置说明
双路径分发utils.pyCheckJitDisable() + ForceAclnn 决定走 op_api 还是 acl_op
FormatHelper::IsOpInputBaseFormatutils.py检查输入张量是否为基础格式
IsAclnnOnlyutils.py识别设备能力差异，选择相应执行路径
OpHookutils.py算子前后钩子，用于调试/profiler
UnsafeTensorCheckutils.py内存重叠检查
NPURecordFunctionGuardutils.pyNPU 特有的 profiler 守卫
兼容实现开关utils.py环境变量控制是否使用兼容实现

4.2 Structured Kernel 生成

原生 Structured Kernel

文件：torchgen/dest/register_dispatch_key.py

struct structured_add_native_functional final : public at::native::structured_add {
  void set_output_strided(...) override { /* ... */ }
  void set_output_raw_strided(...) override { /* ... */ }
  const Tensor& maybe_get_output(int64_t output_idx) override {
    return outputs_[output_idx];
  }
  std::array outputs_;
  c10::native::OptionalNativeGuard guard_;   // 原生设备守卫
};

Tensor wrapper_add_Tensor(const Tensor& self, const Tensor& other,
                           const Scalar& alpha) {
  OptionalDeviceGuard device_guard(device_of(self));
  structured_add_native_functional op;
  op.meta(self, other, alpha);              // shape/dtype 推导
  op.impl(self, other, alpha, op.outputs_[0]); // 计算
  return std::move(op.outputs_[0]);
}
NPU Structured Kernel

文件：torchnpugen/gen_backend_stubs.py

struct structured_add_npu_functional final : public at_npu::native::structured_add_npu {
  void set_output_strided(...) override { /* ... */ }
  void set_output_raw_strided(...) override { /* ... */ }
  const Tensor& maybe_get_output(int64_t output_idx) override {
    return outputs_[output_idx];
  }
  std::array outputs_;
  c10_npu::OptionalNPUGuard guard_;        // NPU 设备守卫
};

Tensor wrapper_NPU_add_Tensor(const Tensor& self, const Tensor& other,
                               const Scalar& alpha) {
  // 1. Device check
  const OptionalDeviceGuard device_guard(device_of(self));
  // 2. UnsafeTensorCheck
  at_npu::native::OpPreparation::check_tensor_mem_overlap(...);
  // 3. NPURecordFunctionGuard
  at_npu::native::NPURecordFunctionGuard guard;
  // 4. OpHook
  const bool op_hook_enabled = C10_UNLIKELY(at_npu::native::env::CheckOpHookEnable());
  if (op_hook_enabled) {
      at_npu::native::OpHook::GetInstance().PreHook("add.Tensor", self, other, alpha);
  }
  // 5. Structured meta + impl
  structured_add_npu_functional op;
  op.meta(self, other, alpha);
  op.impl(self, other, alpha, op.outputs_[0]);
  // 6. OpHook 后置
  // 7. 返回
  return std::move(op.outputs_[0]);
}
Structured Kernel 差异对比
差异点原生NPU
设备守卫原生 Guard / OptionalDeviceGuardc10_npu::OptionalNPUGuard
OpHook无有（PreHook/PostHook）
UnsafeTensorCheck无有
NPURecordFunctionGuard无有
inplace 独立 impl无NPU_STRUCTURED_INPLACE_OPS 中的算子有独立 impl
structured 声明生成标准声明跳过外部 structured 声明（compute_native_function_declaration 补丁）

4.3 op-plugin/codegen 独有的生成逻辑

模块生成内容原生是否有
gen.py双路径分发代码（op_api/acl_op 动态选择）无
struct/struct_codegen.py结构化 aclnn kernel（EXEC_NPU_CMD）无
gen_op_plugin_functions.py版本化 YAML 预处理无
gen_derivatives.py版本化 derivatives 过滤无
gen_backend_stubs.pyop_plugin 函数声明/定义无

4.4 torchnpugen 独有的生成逻辑

模块生成内容原生是否有
custom_functions.py自定义算子注册 + Python 绑定无
gen_functionalization_type.pyFunctionalization 代码torchgen 有但 NPU 扩展了
gen_npu_c_shim.pyAOTInductor C Shim无
autograd/gen_autograd.pyAutograd 代码torchgen 有但 NPU 扩展了

五、生成文件对比

原生 torchgen 生成torchnpugen 对应op-plugin/codegen 对应
RegisterCPU.cpp / RegisterNative.cppRegisterNPU.cpp—
CPUFunctions.h / NativeFunctions.hNPUNativeFunctions.h—
NativeFunctions.hNPUStructuredNativeFunctions.h/cpp—
—CustomRegisterSchema.cpp—
—CustomFunctions.h/cpp—
—custom_ops.py—
——OpInterface.cpp（双路径分发）
——StructKernelNpuOpApi.cpp（aclnn kernel）
——Interface.h（函数声明）

六、猴子补丁机制

torchnpugen 通过 apply_torchgen_patch() 在运行时替换 torchgen 的关键方法，而非 fork 整个仓库：

# torchnpugen/gen_backend_stubs.py
def apply_torchgen_patch():
    # 替换 unstructured 生成逻辑
    dest.RegisterDispatchKey.gen_unstructured = gen_unstructured
    # 替换 device check 逻辑
    dest.RegisterDispatchKey.gen_device_check = gen_device_check
    # 替换 structured class 生成
    StructuredRegisterDispatchKey.gen_class = gen_class
    # 替换 structured set_output_body
    StructuredRegisterDispatchKey.gen_class_set_output_body = gen_class_set_output_body
    # 替换 structured gen_one
    StructuredRegisterDispatchKey.gen_one = gen_one
    # 替换 native function 声明
    dest.compute_native_function_declaration = compute_native_function_declaration
    # 新增默认参数类型
    JIT_TO_CPP_DEFAULT["contiguous_format"] = "c10::MemoryFormat::Contiguous"
    # 替换参数生成
    dispatcher.arguments = native.arguments

被替换的方法及其作用：

原生方法NPU 替换替换原因
gen_unstructuredtorchnpugen/utils.py:gen_unstructured注入双路径分发、OpHook、UnsafeTensorCheck 等
gen_device_checktorchnpugen/utils.py:gen_device_check增加 NPU 设备检查逻辑
gen_class闭包中的 gen_class替换原生 Guard 为 OptionalNPUGuard
gen_class_set_output_body闭包中的 gen_class_set_output_body增加 NPU 设备守卫逻辑
gen_onegen_npu_structured_one注入 OpHook、UnsafeTensorCheck、RecordFunctionGuard
compute_native_function_declaration闭包中的版本跳过外部 structured 声明（由 NPU 自行生成）

七、核心设计差异总结

1. 双路径分发机制（最大创新）

每个算子可以有 op_api（Host 侧 aclnn）和 acl_op（Device 侧 ACL）两种实现，运行时根据以下条件动态选择：

CheckJitDisable()  →  true  →  op_api（aclnn Host 调用）
ForceAclnn          →  true  →  op_api
IsOpInputBaseFormat →  true  →  op_api
否则                       →  acl_op（ACL Device 调用）

2. 版本化算子管理

YAML 中每个算子都标注版本范围，codegen 时根据当前版本过滤，实现算子的渐进式支持：

op_api: all_version      # 所有版本
acl_op: [v2.1, newest]   # v2.1+
op_api: v2.1             # 仅 v2.1

3. 结构化 aclnn 代码生成

通过 gen_opapi 配置，自动生成 EXEC_NPU_CMD(aclnnXxx, ...) 调用代码，无需手写 op_api 实现：

gen_opapi:
  out:
    size: self
    dtype: self
  exec: aclnnAbs, self, out

生成代码：

at::Tensor abs(const at::Tensor& self) {
    DO_COMPATIBILITY(aclnnAbs, acl_op::abs(self));
    auto output_size_0 = self.sizes();
    auto output_dtype_0 = self.scalar_type();
    at::Tensor out = npu_preparation::apply_tensor_without_format(
        output_size_0, self.options().dtype(output_dtype_0));
    EXEC_NPU_CMD(aclnnAbs, self, out);
    at::namedinference::propagate_names(out, self);
    return out;
}

4. 猴子补丁扩展

torchnpugen 不 fork torchgen，而是通过 apply_torchgen_patch() 在运行时替换关键方法，保持与上游同步的同时注入 NPU 逻辑。

5. 精简数据模型

op-plugin 的 NativeFunction 去掉了 variants、ufunc_inner_loop、python_module 等不需要的字段，新增 impl_ns/impl_name/sparse/internal_format_opapi。

6. NPU 特有基础设施

基础设施说明原生是否有
c10_npu::OptionalNPUGuardNPU 设备守卫无（原生为通用设备守卫）
OpHook算子前后钩子无
UnsafeTensorCheck内存重叠检查无
NPURecordFunctionGuardProfiler 守卫无
ForceAclnn强制使用 aclnn无
IsAclnnOnly识别设备能力差异无
FormatHelper::IsOpInputBaseFormat张量格式检查无
DO_COMPATIBILITYaclnn 兼容性回退无
EXEC_NPU_CMDaclnn 调用封装无

八、总结

torch_npu 的 Codegen 体系在继承原生 torchgen 框架的基础上，针对 NPU 硬件特性和软件栈进行了系统性扩展。其核心创新体现在三个方面：一是通过双路径分发机制实现算子在不同执行模式间的灵活切换；二是通过版本化 YAML 配置实现算子的渐进式支持；三是通过猴子补丁机制在保持与上游同步的同时注入 NPU 特有逻辑。这一设计为 NPU 算子的高效适配和持续演进提供了坚实的技术基础。

# pytorch
Logo
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

更多推荐

cover

DeepInfer 完全教程：DeepSeek 内部推理框架原理、架构与部署实战

avatar 鲲鹏昇腾开发者社区

Ascend for PyTorch 中虚拟机 run 的执行过程

TorchDynamo 是 PyTorch 2.x 中负责将 Python 字节码转换为计算图的核心组件。其设计目标是在不修改用户代码的前提下，通过符号执行的方式捕获模型计算图，进而交由后端编译器进行优化。理解 TorchDynamo 的虚拟机（VM）执行机制，对于深入掌握 PyTorch 编译优化原理、排查图捕获问题以及进行二次开发都至关重要。

avatar 鲲鹏昇腾开发者社区

昇腾平台 CV 模型微调：图像交互场景实践

本文系统阐述了昇腾平台在图像交互任务（如VQA、Caption、交互式分割）中的CV模型微调实践，涵盖硬件底座、多模态数据处理、微调策略、分布式训练、性能优化与评估全流程。基于昇腾DaVinci架构与MindSpore生态，实现从数据预处理到模型部署的完整链路，强调冻结主干、离线预处理、激活重计算等关键调优手段，确保高分辨率下显存可控与训推一致。实证表明，结合业务指标（mIoU/VQA-Acc/B

avatar 鲲鹏昇腾开发者社区