​作者​：昇腾实战派
​知识地图​：https://blog.csdn.net/Lumos_Lovegood/article/details/161601003

0. 背景与动机

问题描述

在使用 torch.compile 时，Inductor 后端需要编译 C++ kernel，这一过程需要获取设备相关的编译选项，包括 include 路径、链接库和宏定义等。当前，这部分逻辑由 torch/_inductor/cpp_builder.py 中的 get_cpp_torch_device_options() 函数提供。

然而，该函数目前仅硬编码支持三种设备：

设备	判定条件	提供内容
CUDA	device_type == "cuda"	USE_CUDA / USE_ROCM、libcuda.so 等
XPU	device_type == "xpu"	USE_XPU、ze_loader、sycl 等
MPS	device_type == "mps"	USE_MPS

对于像 NPU（Ascend）这样的 PrivateUse1 设备，并不在支持范围内。当调用 get_cpp_torch_device_options("npu") 时，三个 if 分支全部跳过，返回空的设备编译选项，导致 C++ kernel 编译失败。

当前临时方案

目前，torch_npu 通过 monkey-patch 在运行时替换整个函数：

# torch_npu/_inductor/cpp_builder.py
def patch_get_cpp_torch_device_options():
    torch._inductor.cpp_builder.get_cpp_torch_device_options = get_cpp_torch_device_options

这种方式存在以下问题：

• 非正规 API，多个后端同时 patch 会互相覆盖
• 上游函数签名变化时会导致静默崩溃
• 无法被上游社区接受

为什么不能直接将 NPU 硬编码进上游

NPU 依赖的外部资源均不在 PyTorch 仓库内：

NPU 所需资源	维护方	PyTorch 仓库中是否存在
/usr/local/Ascend/include	CANN 工具包（华为）	无
libascendcl.so	CANN 工具包（华为闭源）	无
libtorch_npu.so	torch_npu 项目	无
Ascend 910B 测试硬件	华为	PyTorch CI 中不存在

若硬编码 elif device_type == "npu"，PyTorch CI 会因找不到这些路径和库而构建失败。

1. 方案设计

核心思路

上游提供注册 API，NPU 在 torch_npu 初始化时登记回调。上游不感知 NPU 的具体细节。

                    PyTorch 上游                          torch_npu
                   ┌──────────────────────┐         ┌──────────────────┐
                   │ _registry = {}        │         │ def npu_options():│
                   │                      │  登记    │   return [        │
                   │ register_cpp_device  │←────────│     "USE_NPU",    │
                   │   _options("npu",fn) │         │     "/Ascend/...",│
                   │                      │         │     "ascendcl"    │
                   │ get_cpp_torch_device │         │   ]               │
                   │   _options("npu")    │         └──────────────────┘
                   │   ├─ cuda? no        │
                   │   ├─ xpu?  no        │
                   │   ├─ mps?  no        │
                   │   └─ 查注册表 → 命中  │
                   │       → 调用 npu回调  │
                   │       → 返回 NPU选项  │
                   └──────────────────────┘

与已有注册模式的对照

PyTorch Inductor 已有多个设备注册机制，本实践补齐最后一块缺失的拼图：

注册机制	文件	模式
register_backend_for_device	codegen/common.py	注册表 dict
register_interface_for_device	_dynamo/device_interface.py	注册表 dict
register_device_op_overrides	codegen/common.py	注册表 dict
register_cpp_device_options（本次新增）	cpp_builder.py	注册表 dict

2. 改动内容

文件：torch/_inductor/cpp_builder.py（+38 行，-1 行）

2.1 import 行

修改前：

from collections.abc import Sequence

修改后：

from collections.abc import Callable, Sequence

目的：注册表类型标注 dict[str, Callable] 需要 Callable。

2.2 新增注册表与注册函数（插在 get_cpp_torch_device_options 定义之前）

修改前（_transform_rocm_paths 之后直接是函数定义）：

def _transform_rocm_paths(lpaths: list[str]) -> None:
    if config.is_fbcode():
        sdk_lib = build_paths.sdk_lib
        if os.path.isdir(sdk_lib) and sdk_lib not in lpaths:
            lpaths.append(sdk_lib)

def get_cpp_torch_device_options(

修改后：

def _transform_rocm_paths(lpaths: list[str]) -> None:
    if config.is_fbcode():
        sdk_lib = build_paths.sdk_lib
        if os.path.isdir(sdk_lib) and sdk_lib not in lpaths:
            lpaths.append(sdk_lib)

# Registry for out-of-tree device backends to provide C++ compile options.
# Maps device_type (e.g. "npu") to a callback with the same signature as
# get_cpp_torch_device_options.
_cpp_device_options_registry: dict[str, Callable] = {}

def register_cpp_device_options(
    device_type: str,
    callback: Callable,
) -> None:
    """Register C++ build options for a device type.

Allows out-of-tree device backends (e.g., PrivateUse1 devices like NPU)
    to inject their own compile options into Inductor's C++ compilation
    pipeline without monkey-patching.

Args:
        device_type: The device type string (e.g. "npu").
        callback: A callable with the same signature as
            :func:`get_cpp_torch_device_options`:
            ``(device_type, aot_mode, compile_only) ->
            tuple[definitions, include_dirs, cflags, ldflags,
                  libraries_dirs, libraries, passthrough_args]``
    """
    if device_type in _cpp_device_options_registry:
        warnings.warn(
            f"Device options for '{device_type}' is already registered, overwriting."
        )
    _cpp_device_options_registry[device_type] = callback

def get_cpp_torch_device_options(

目的：提供模块级注册表 _cpp_device_options_registry 与公开 API register_cpp_device_options()。重复注册时发出警告并覆盖，与 PyTorch 其他注册机制一致。

2.3 函数内新增查表逻辑（插在 mps 分支之后、fbcode 分支之前）

修改前：

    if device_type == "mps":
        definitions.append(" USE_MPS")

if config.is_fbcode():

修改后：

    if device_type == "mps":
        definitions.append(" USE_MPS")

# Check registered device backends (e.g., PrivateUse1 devices like NPU)
    if device_type in _cpp_device_options_registry:
        return _cpp_device_options_registry[device_type](
            device_type, aot_mode, compile_only
        )

if config.is_fbcode():

目的：cuda/xpu/mps 均不命中时查注册表。命中则直接委托给注册的回调并返回，跳过后续公共路径。未命中则继续原逻辑（返回空设备选项）。

文件：torch_npu/_inductor/cpp_builder.py（下游适配）

修改前：

def patch_get_cpp_torch_device_options():
    torch._inductor.cpp_builder.get_cpp_torch_device_options = get_cpp_torch_device_options

修改后：

def patch_get_cpp_torch_device_options():
    """Register NPU device options to PyTorch Inductor C++ builder.

Uses the official registration API when available (PyTorch >= 2.7),
    falls back to monkey-patching for older versions.
    """
    if hasattr(torch._inductor.cpp_builder, 'register_cpp_device_options'):
        # New path: proper API registration
        torch._inductor.cpp_builder.register_cpp_device_options(
            "npu", get_cpp_torch_device_options
        )
    else:
        # Legacy path: monkey-patch for older PyTorch versions
        torch._inductor.cpp_builder.get_cpp_torch_device_options = \
            get_cpp_torch_device_options

目的：优先使用上游注册 API；未合入此 PR 的旧版 PyTorch 回退 monkey-patch。get_cpp_torch_device_options 回调本身无需改动，签名已与上游一致。

3. 调用链路

完整时序

1. torch_npu 加载
   → __init__.py → patch_torch_for_aoti()
   → patch_get_cpp_torch_device_options()
   → register_cpp_device_options("npu", npu_get_options)
   → _registry = {"npu": npu_get_options}            ← 登记完成

2. 用户代码
   @torch.compile(backend="inductor")
   def model(x): ...

3. Inductor 生成 C++ kernel → 需要编译
   → CppBuilder.build()
   → CppTorchDeviceOptions(device_type="npu")
   → get_cpp_torch_device_options("npu")
     ├─ device_type == "cuda"?  → no
     ├─ device_type == "xpu"?   → no
     ├─ device_type == "mps"?   → no
     ├─ "npu" in _registry?     → yes
     └─ return _registry["npu"]("npu", aot_mode, compile_only)
        → 返回 (["USE_NPU"], ["/usr/local/Ascend/include", ...],
                 [], [], ["/usr/lib64", ...], ["ascendcl", "torch_npu"],
                 ["-DBUILD_LIBTORCH=ON"])

4. g++ 命令组装
   g++ kernel.cpp
     -DUSE_NPU
     -I/usr/local/Ascend/include
     -L/usr/lib64
     -lascendcl -ltorch_npu
     -DBUILD_LIBTORCH=ON
     -o kernel.so

关键不变量

• CUDA/XPU/MPS 的处理逻辑一行不改
• 未注册的设备类型行为不变（返回空设备选项）
• torch_npu 在旧版 PyTorch 上兼容（回退 monkey-patch）

4. PR 信息

项目	内容
PR 编号	#190614
仓库	pytorch/pytorch
分支	lyx0517:inductor-npu-device-options → pytorch:main
改动量	1 file, +38 / -1
RFC	https://gitcode.com/suhaibo/pytorch/wiki

5. 兼容性与风险

项	策略
旧版 PyTorch 无注册 API	torch_npu 保留 monkey-patch fallback
两个后端注册同名设备	warn + overwrite（与 register_backend_for_device 一致）
注册回调抛异常	异常直接传播，不吞
已有 cuda/xpu/mps	逻辑一字不改

6. 验证

在 CANN 9.0.1 + PyTorch 2.13.0（安装版，piped）容器内完成功能验证：

#	测试项	结果
1	mock 设备注册 + 查询	✅
2	重复注册触发 warning	✅
3	未知设备不污染 CUDA 定义	✅
4	未知设备不污染 XPU 定义	✅
5	未知设备不污染 MPS 定义	✅
6	未知设备不泄露 mock 定义	✅
7	NPU 回调注册后 USE_NPU 注入	✅
8	NPU include 路径正确	✅
9	NPU 链接库正确（ascendcl, torch_npu）	✅
10	NPU passthrough 参数正确	✅