## 1. 背景

### 1.1 业务背景

PyTorch 原生 `torch.utils.cpp_extension` 模块为多种原生源语言提供了统一的编译扩展能力：开发者声明 `Extension`、注册 `BuildExtension`，即可自动获得 ninja 增量编译（lazy）、多核并行编译（parallel）、ABI 探测、依赖追踪等能力。

昇腾 NPU 适配工程共有 9 个示例工程，其中 4 个涉及 Ascend C 内核源码（`.asc` 文件），必须使用昇腾专用编译器 `bisheng` 编译。由于原生 `cpp_extension.py` 没有 `.asc` 和 `bisheng` 的支持插槽，这 4 个工程被迫自定义 `build_ext`，绕过原生构建链路，采用单命令 bisheng 一体编译+链接。

### 1.2 问题触发

在相关适配工程的评审中，有开发者反馈：现有构建方式在 Python 侧单独注入，没有纳入原生 cmake→ninja 的 lazy / parallel 构建体系，也无法把 Ascend 内核二进制产物直接和 Python 扩展链接到同一个动态库中，希望获得与 PyTorch 原生编译一致的体验。

该反馈包含三个递进的技术判断，本质指向同一个根因：**原生 `BuildExtension` 不认识 `.asc` 源和 `bisheng` 编译器**。

### 1.3 名词术语

| 术语 | 定义 |

|------|------|

| `BuildExtension` | PyTorch 原生类（`torch.utils.cpp_extension.BuildExtension`），继承 setuptools `build_ext`，是原生扩展的构建驱动 |

| `_write_ninja_file` | 原生函数，负责生成 `build.ninja` 文件，内含 compile/link 规则和依赖追踪配置 |

| `_is_cuda_file` / `_is_sycl_file` | 原生函数，按文件后缀识别对应源语言，路由到对应编译规则 |

| `bisheng` | 昇腾专用编译器，支持 `-x asc` 编译 Ascend C 源码、`--npu-arch` 指定芯片架构 |

| `.asc` | Ascend C 内核源文件后缀，编译产物为 AI Core 二进制（aibin） |

| `aibin` | AI Core 二进制产物，由 bisheng 编译 `.asc` 时附带生成，需合入最终 `.so` |

| lazy build | 增量编译，只重编变更的源文件，依赖 `-MMD -MF $out.d + deps=gcc` |

| parallel build | 并行编译，每个源文件一条 ninja build 规则，ninja 按 CPU 核数并行 |

| `NpuExtension` | NPU 适配层提供的 `CppExtension` 封装，自动注入 NPU 头文件/库路径 |

| `--npu-arch` | bisheng 的芯片架构参数，按芯片型号自动映射（支持 910 系列/950 系列芯片） |

---

## 2. 现状分析

### 2.1 原生 PyTorch 的编译能力架构

原生 `BuildExtension` 的 lazy/parallel 能力通过以下机制实现：

```

setup(ext_modules=[ext], cmdclass={"build_ext": BuildExtension})

    │

    ▼

BuildExtension.build_extensions()

    ├─ ① _check_abi()                    → 自动探测 C++ ABI

    ├─ ② src_extensions += [.cu/.sycl]   → 注册合法源后缀

    ├─ ③ compiler.compile = ninja 版     → 替换编译方法（关键！）

    └─ ④ super().build_extensions()      → 走 distutils 默认流程

            └─ compiler.compile(...)

                    └─ unix_wrap_ninja_compile()

                            ├─ _write_ninja_file()     → 生成 build.ninja

                            │   ├─ rule compile:  $cxx -MMD -MF $out.d ... -c $in -o $out

                            │   │   depfile = $out.d / deps = gcc          ← lazy

                            │   ├─ rule cuda_compile: $nvcc ... -c $in -o $out

                            │   ├─ rule sycl_compile: $sycl ... -c $in -o $out

                            │   ├─ rule link: $cxx $in $ldflags -o $out

                            │   └─ build x.o: compile x.cpp                 ← parallel（每源一条）

                            └─ subprocess.run(["ninja"])

```

**机制要点**（具体实现细节可参考 PyTorch 官方源码，此处仅描述原理）：

- lazy 能力依赖 `depfile` + `deps = gcc`：编译器输出 `.d` 依赖文件，ninja 据此实现只重编变更源文件。

- parallel 能力依赖「每源一条 build 规则」：ninja 按依赖图并行调度。

- 编译器的替换点在于 `compiler.compile` 与内部 `_write_ninja_file` 生成的 ninja 规则。

**语言识别机制**：原生通过**文件后缀**识别源语言，通过 `_write_ninja_file` 里**硬编码的编译器变量**（`$cxx`/`$nvcc`/`$sycl`）和**硬编码的编译规则**（`compile`/`cuda_compile`/`sycl_compile`）路由编译。当前仅支持 `.cu`/`.cuh`/`.sycl` 等原生源后缀，没有 `.asc`。

### 2.2 昇腾适配工程当前实现

4 个 Ascend C 工程的当前实现对比：

| 工程 | 继承基类 | cmdclass 写法 | 芯片架构 | 编译方式 |

|------|---------|--------------|---------|---------|

| `cpp_extension` | torch `BuildExtension` | `.with_options(use_ninja=True)` | 硬编码 | ninja |

| `acl_dnn` | 自定义 `AscendBuildExtension` | 手写 `cmdclass` | 硬编码 | 单命令 bisheng |

| `kernel_ascend` | 自定义 `AscendBuildExtension` | 手写 `cmdclass` | 硬编码 | 单命令 bisheng |

| `attention` | 自定义 `AscendBuildExtension` | 手写 `cmdclass` | 硬编码 | 单命令 bisheng |

> 注：上表为示意（工程名已泛化），当前实现的核心问题在于：4 个工程中除首个外均绕过了原生构建链路，且芯片架构为硬编码参数，无法自动适配不同型号。

**自定义 `build_ext` 的实现方式**：

```python

class AscendBuildExtension(build_ext):

    def build_extension(self, ext):

        # 单命令 bisheng 编译+链接，无 ninja、无增量、无并行

        compile_cmd = [...]

        self.spawn(compile_cmd)

```

**现状问题总结**：

1. 不享受原生 ninja 的 lazy/parallel 能力，全量编译耗时随工程规模线性增长；

2. aibin 产物无显式依赖声明，无法进入构建依赖图、无法追踪；

3. 芯片架构硬编码，换型号需改代码；

4. 重复造轮子：每个工程各写一套 `compile_cmd`。

---

## 3. 目标与方案设计

### 3.1 设计目标

| 编号 | 目标 | 说明 |

|------|------|------|

| US-1 | 原生支持 `.asc` 源 | `Extension(language="asc")` + `BuildExtension` 直接编译 |

| US-2 | 支持 lazy 增量编译 | 修改单个 `.asc` 只重编该文件 |

| US-3 | 支持 parallel 并行编译 | 多个 `.asc` 并行编译 |

| US-4 | aibin 显式产物声明 | 进入构建依赖图，可追踪 |

| US-5 | 芯片架构自动探测 | 免手写 `--npu-arch` |

| US-6 | 适配工程统一迁移 | 4 个工程全部走原生链路 |

### 3.2 总体思路

在 PyTorch 原生编译链路中补齐 `.asc` 与 `bisheng` 的支持插槽，让 Ascend C 工程与原生工程享受一致的构建体验：

```

Extension(language='asc') + BuildExtension

    │

    ▼

_write_ninja_file()

    ├─ src_extensions += ['.asc']          ← FR-1 源注册

    ├─ config += ['bisheng = ...']         ← FR-2 编译器变量

    ├─ rule asc_compile: $bisheng -x asc  ← FR-3 编译规则

    │     --npu-arch=$npu_arch -MMD -MF ...

    ├─ rule asc_devlink: $bisheng ...      ← FR-5 设备链接规则

    └─ per-source build rules              ← lazy + parallel（复用 ninja 机制）

        ▲

        └─ get_npu_arch() 自动探测         ← FR-4

```

---

## 4. 实现方案

### 4.1 原生层改动总览

| 功能需求 | 目标 | 位置 |

|---------|------|------|

| FR-1 | `.asc` 源注册 | `cpp_extension.py`（原生层） |

| FR-2 | `bisheng` 编译器变量 | `cpp_extension.py` |

| FR-3 | `asc_compile` 编译规则 | `cpp_extension.py` |

| FR-4 | `--npu-arch` 架构自动探测 | NPU 适配层 |

| FR-5 | `asc_devlink` 设备链接规则 | `cpp_extension.py` |

| FR-6 | `AscendExtension` 工厂函数 | NPU 适配层 |

### 4.2 原生层改动

#### FR-1：`.asc` 源注册

```python

# 在 BuildExtension 的源后缀注册处新增

if with_asc:  # 检测到 .asc 源时

    self.compiler.src_extensions += ['.cu', '.cuh', '.sycl', '.asc']

```

#### FR-2：`bisheng` 编译器变量

参照 `_write_ninja_file` 中原生编译器变量的写法：

```python

# 在 _write_ninja_file 的 config 块中新增

if with_asc or asc_dlink_post_cflags:

    if "PYTORCH_BISHENG" in os.environ:

        bisheng = os.getenv("PYTORCH_BISHENG")

    else:

        bisheng = _find_bisheng()  # 从 PATH 查找 bisheng

    config.append(f'bisheng = {bisheng}')

```

#### FR-3：`asc_compile` 编译规则

参照原生编译规则的定义方式：

```python

if with_asc:

    asc_compile_rule = ['rule asc_compile']

    asc_compile_rule.append(

        '  command = $bisheng -x asc --npu-arch=$npu_arch '

        '-MMD -MF $out.d $asc_cflags -c $in -o $out $asc_post_cflags'

    )

    asc_compile_rule.append('  depfile = $out.d')

    asc_compile_rule.append('  deps = gcc')

```

#### FR-4：`--npu-arch` 架构探测

参照原生编译器探测的模式，在 NPU 适配层提供：

```python

def get_npu_arch() -> str:

    result = subprocess.run(["npu-smi", "info"], capture_output=True, text=True)

    # 解析芯片型号，映射到对应架构参数

    if '950' in chip_name:

        return 'dav-3510'

    elif '910' in chip_name:

        return 'dav-2201'

```

#### FR-5：aibin 显式产物声明（设备链接）

参照原生设备链接规则：

```python

if asc_dlink_post_cflags:

    asc_devlink_out = os.path.join(os.path.dirname(objects[0]), 'asc_dlink.o')

    asc_devlink_rule = ['rule asc_devlink']

    asc_devlink_rule.append(

        '  command = $bisheng --npu-arch=$npu_arch $in -o $out $asc_dlink_post_cflags'

    )

    asc_devlink = [f'build {asc_devlink_out}: asc_devlink {" ".join(objects)}']

    objects += [asc_devlink_out]

```

#### FR-6：`AscendExtension` 工厂函数

参照原生扩展工厂函数，放在 NPU 适配层：

```python

def AscendExtension(name, sources, extra_compile_args=None, ...):

    # 注入 NPU 头文件/库路径 + language='asc'

    return Extension(

        name=name,

        sources=sources,

        language='asc',

        extra_compile_args=extra_compile_args or {},

        ...

    )

```

### 4.3 NPU 适配层改动

| 改动项 | 位置 | 说明 |

|--------|------|------|

| `AscendExtension` 工厂函数 | NPU 适配层 utils/cpp_extension.py | 类比 `NpuExtension`，注入 NPU 路径 + `language='asc'` |

| `get_npu_arch()` | NPU 适配层 utils/cpp_extension.py | `npu-smi info` 探测芯片架构 |

| 注册 `.asc` 到 `src_extensions` | NPU 适配层（`_LazyBuildExt` 或 monkey-patch） | 让 `BuildExtension` 认识 `.asc` |

### 4.4 昇腾适配工程层改动

4 个工程的 setup.py 统一改为：

```python

from torch_npu.utils.cpp_extension import AscendExtension

from torch.utils.cpp_extension import BuildExtension

ext = AscendExtension(

    name="xxx.custom_ops_lib",

    sources=source_files,          # glob("csrc/*.asc")

)

setup(

    ext_modules=[ext],

    cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)},

)

```

删除内容：

- 自定义 `AscendBuildExtension` 类

- 覆写的 `build_extension()` 方法

- 手写的 `compile_cmd` 拼接

- 硬编码的 `--npu-arch` 参数

### 4.5 实现优先级与里程碑

| 阶段 | 目标 | 涉及功能 | 交付物 |

|------|------|---------|--------|

| **阶段一** | 原生支持 `.asc` + `bisheng` | FR-1/2/3/6 | 原生 `cpp_extension.py` 补 `asc_compile` 规则 + `bisheng` 变量；NPU 适配层提供 `AscendExtension` |

| **阶段二** | 适配工程 4 工程迁移 | 全部 | 4 个 setup.py 改用 `AscendExtension + BuildExtension`；test.sh 全 PASS |

| **阶段三** | aibin 显式化 + 架构探测 | FR-4/5 | `asc_devlink` 规则合入原生；`get_npu_arch()` 合入 NPU 适配层 |

---

## 5. 验收标准

### 5.1 功能验收

| 编号 | 验收项 | 验收方法 | 对应需求 |

|------|--------|---------|---------|

| AC-1 | `Extension(language="asc")` + `BuildExtension` 可编译 `.asc` 源 | 新建最小工程编译成功 | US-1, FR-1/2/3/6 |

| AC-2 | 修改单个 `.asc` 后只重编该文件 | 触摸一个 `.asc`，重跑 build，ninja 日志显示只编 1 个 | US-2, FR-3 |

| AC-3 | 多个 `.asc` 并行编译 | ninja 日志显示多线程编译 | US-3, FR-3 |

| AC-4 | aibin 出现在构建依赖图中 | `ninja -t deps` 可查到 aibin 依赖 | US-4, FR-5 |

| AC-5 | 芯片架构自动探测 | 不写 `--npu-arch`，构建成功且芯片正确 | US-5, FR-4 |

| AC-6 | 4 个工程迁移后 test.sh 全 PASS | `bash test.sh` 全 PASS | US-6 |

| AC-7 | 原 C++ 工程编译不受影响 | base/full/pybind 等工程编译行为不变 | NFR-1 |

### 5.2 非功能验收

| 编号 | 验收项 |

|------|--------|

| NAC-1 | 适配工程中不再有自定义 `AscendBuildExtension` 类 |

| NAC-2 | 适配工程中不再有手写 `compile_cmd` 拼接 |

| NAC-3 | 原有 C++ 工程编译路径无回归 |

---

## 6. 风险与约束

| 编号 | 风险/约束 | 影响 | 缓解措施 |

|------|----------|------|---------|

| R-1 | bisheng 对 depfile（`-MMD -MF`）的兼容性待验证 | FR-3 增量编译能力依赖此兼容性 | 阶段一优先验证 bisheng 的 depfile 支持 |

| R-2 | bisheng 的 device link 模式待验证 | FR-5 aibin 显式化依赖此能力 | 阶段三前与编译工具链维护方确认 |

| R-3 | 原生上游接受度 | 改动能否合入 PyTorch 主干 | 以 guard 方式实现，非 Ascend 用户不受影响 |

| R-4 | bisheng 不在 PATH | 构建失败 | 支持通过 `PYTORCH_BISHENG` 环境变量指定 bisheng 路径 |

| R-5 | 新芯片型号未覆盖 | `get_npu_arch()` 报错 | 报错信息引导用户手动指定 `--npu-arch` |

---

## 7. 依赖关系

```

bisheng 支持 -MMD -MF（depfile）

    │

    ▼

FR-3 asc_compile 规则

    │

    ├── FR-1 _is_asc_file  ──┐

    ├── FR-2 bisheng 变量   ──┤

    │                        ├── FR-6 AscendExtension ──→ 适配工程迁移

    │                        │

    └── bisheng 支持 dlink ──┤

                             │

                        FR-5 asc_devlink

                             │

                        FR-4 get_npu_arch

```

**关键前置依赖**：bisheng 对 `-MMD -MF` 和 device link 的支持需优先与编译工具链维护方确认。

---

## 8. 附录

### 附录 A：原生扩展编译规则（参照模板）

以下规则结构为新增 `asc_compile` / `asc_devlink` 规则的参照：

```python

# compile 规则（C++ 源）

compile_rule = ['rule compile']

compile_rule.append('  command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags')

compile_rule.append('  depfile = $out.d')

compile_rule.append('  deps = gcc')

# 每源一条 build 规则（parallel 的关键）

for source_file, object_file in zip(sources, objects):

    build.append(f'build {object_file}: {rule} {source_file}')

# 设备链接规则（aibin 显式化的参照）

if dlink_post_cflags:

    dlink_out = os.path.join(os.path.dirname(objects[0]), 'dlink.o')

    dlink_rule = ['rule dlink']

    dlink_rule.append('  command = $compiler $in -o $out $dlink_post_cflags')

    dlink = [f'build {dlink_out}: dlink {" ".join(objects)}']

    objects += [dlink_out]

```

### 附录 B：当前 4 个 Ascend C 工程的编译命令（待替换）

```python

# 4 个工程统一使用此命令（单命令 bisheng，无 ninja）

compile_cmd = [

    "bisheng", "-x", "asc", "--npu-arch=dav-2201",

    "-shared", "-fPIC", "-std=c++17",

    f"-D_GLIBCXX_USE_CXX11_ABI={abi_value}",

    "-ltorch_npu", "-ltorch", "-lc10",

    *ext.sources, "-o", ext_fullpath,

]

self.spawn(compile_cmd)

```

### 附录 C：需求背景说明

评审反馈的核心诉求归纳如下：

| 诉求 | 实际含义 |

|---------|---------|

| 构建方式未纳入原生标准链路 | 绕过原生标准构建体系 |

| 未享受 lazy / parallel 能力 | 单命令 bisheng，无 ninja 增量/并行 |

| 无法将 aibin 链接到同一动态库 | aibin 靠 bisheng 隐式兜底，非显式构建产物 |