Ascend for PyTorch 训练营学习笔记第四期续

API代码生成 Ascend for PyTorch 训练营(https://www.chaspark.com/#/s/AscendforPyTorch2026?multi=zh) 第四期

Dispatcher 分发器

• 作用：算子调度的"路由器"。根据输入张量的设备类型、是否需要梯度等信息，动态匹配最合适的后端实现。
• 机制：维护一个按优先级排序的 DispatchKey 列表，依次查表、命中即执行。
• DispatchKey：
  • CPU / CUDA / PrivateUse1（第三方硬件，昇腾 NPU 挂在这里）
  • Meta（形状推断）
  • Autograd（自动微分）
• 一次调用的查表链路（以 at::add 为例）：
  1. Functionalize → torch.compile 功能化层
  2. PythonDispatcher → Python 钩子
  3. ADInplaceOrView → 追踪 inplace/view 操作
  4. AutogradCPU → 构建反向计算图
  5. BackendSelect → 根据输入设备选后端
  6. CPU / CUDA / NPU → 执行具体 Kernel

代码生成 5 步走：

1. 主仓入口（generate_code.sh）
2. 生成 OpPlugin 配置（op-plugin/gencode.sh）
3. 核心后端 stub 生成（torchnpugen.gen_backend_stubs）
4. 自动微分（autograd）生成
5. 测试与辅助功能生成

核心后端 stub 生成: “NPU 算子如何被 PyTorch 调度”

• 分发键（DispatchKey）：NPU 算子统一通过 PrivateUse1这个 dispatch key 分发到昇腾后端。
• 算子注册：用 TORCH_LIBRARY_IMPL 宏注册算子
  • 原生 API（如 add.Tensor）→ 注册到 aten 命名空间
  • 自定义算子（如 npu_linear）→ 注册到 npu 命名空间
• 生成产物：wrapper 包装函数 + 静态注册代码块，最终产出 RegisterNPU.cpp 等文件。
• Python 侧调用：注册后可直接通过 torch.ops.aten.xxx 或 torch.ops.npu.xxx 调用。

torchnpugen 通过解析算子定义，自动生成 wrapper 和注册代码，把昇腾 NPU 算子挂到 PyTorch 的 PrivateUse1 调度体系上，实现"写一次配置，自动生成适配层"。

自动微分（Autograd）代码生成：“NPU 算子自动求导”

核心机制

• 输入：TorchNPU 自定义 YAML + OpPlugin 算子 YAML
• 生成模块：torchnpugen.autograd.gen_autograd
• 依据：derivatives.yaml（定义算子的导数规则）

生成内容

代码示例

// 反向算子注册到 AutogradPrivateUse1
TORCH_LIBRARY_IMPL(aten, AutogradPrivateUse1, m) {
    m.impl("npu_linear", TORCH_FN(VariableType::npu_linear));
}

// Functionalize：消除 in-place，转为 out-of-place
TORCH_LIBRARY_IMPL(npu, Functionalize, m) {
    m.impl("npu_silu_", TORCH_FN(functionalization::npu_silu_));
}

复用 PyTorch 上游 Codegen 能力（原生算子的反向绑定直接复用），只针对 NPU 差异生成"胶水代码"，实现轻量化适配、随社区版本平滑演进。

通过解析导数配置，自动生成反向传播绑定和图模式所需的 functionalize 转换代码，让 NPU 算子完整支持自动求导和编译优化。