---
tags: [CANN]
doc_type: knowledge-note
返回目录: "[[计算机系/昇腾与CANN/大模型与算子专题学习导航|大模型与算子专题学习导航]]"
上一篇: "[[计算机系/昇腾与CANN/Ascend C融合算子与核间协同|Ascend C融合算子与核间协同]]"
下一篇: "[[计算机系/昇腾与CANN/Ascend C算子入图与图执行|Ascend C算子入图与图执行]]"
后续学习: ["[[计算机系/昇腾与CANN/Ascend C算子入图与图执行|Ascend C算子入图与图执行]]"]
---

# Ascend C接入PyTorch单算子

Ascend C 核函数规定 NPU 如何计算；PyTorch 接入层把 Tensor、设备、执行流、内存和函数调用连接起来。一个算子可以看作带输入输出的函数，但进入框架后还需要明确类型、形状、设备分发和求导语义。

本篇只处理单算子调用。编译基础见 [CMake 入门](CMake与Ascend%20C编译入门.md)，图模式扩展见 [算子入图与图执行](Ascend%20C算子入图与图执行.md)。

## 1. 两条接入路线

| 路线 | Python 调用 | 注册内容 | 适用场景 |
| --- | --- | --- | --- |
| PyTorch 算子注册 | `torch.ops.ascendc_ops.ascendc_add(x,y)` | Schema、设备实现等 | 与 dispatcher 和图编译配合 |
| pybind11 | `ascendc_ops.ascendc_add(x,y)` | Python 模块与 C++ 函数绑定 | 快速暴露 C++ 接口 |

`torch.library` 是 Python 侧算子注册 API 家族；C++ 可用 `TORCH_LIBRARY/TORCH_LIBRARY_IMPL` 进入同一套调度体系。

单独使用 pybind11 不会自动提供 Schema 或图追踪语义，可能导致图中断。它并非永远不能与 `torch.compile` 共存：可以再包成正式自定义算子，并补充相应图支持。反过来，完成 Schema 注册也不等于所有编译后端立即支持该算子。

## 2. Host 适配层的职责

```text
Python Tensor
→ C++ 包装函数
→ 检查设备、dtype、shape、stride
→ 申请输出 Tensor
→ 获取当前 NPU stream
→ 启动 Ascend C kernel
→ 返回输出 Tensor
```

输出可用 `at::empty(x.sizes(), x.options())` 创建，使设备和类型跟随输入。它只分配空间，不初始化有效值；Kernel 必须覆盖所有输出元素。

当前流的典型获取方式为：

```cpp
auto stream = c10_npu::getCurrentNPUStream().stream(false);
```

使用当前流可以维持与前后 PyTorch 操作的执行顺序。包装函数返回时，NPU 工作可能尚未完成；不能用 Host 函数返回时刻作为 Kernel 结束时刻。多设备情况下还需设置与输入设备匹配的设备上下文。

适配层至少要处理：

- x、y 是否都在同一个 NPU 上。
- dtype 是否与 Kernel 中的指针解释一致。
- shape 是否符合实现；简单等长 Add 不会自动获得广播能力。
- Tensor 是否连续；非连续输入需要 stride 支持、连续化处理或明确拒绝。
- `numel()` 是否超出 Kernel 长度参数范围，空输入和尾块如何处理。

把数据指针转换成 `uint8_t*` 只是以字节指针传参，不能改变实际元素类型。Kernel 如果把它解释为 float，而输入为 half，会产生错误访问。

## 3. 注册 Schema 与 NPU 实现

![PyTorch 正式注册流程](.assets/AscendC进阶/torch-13.png)

假设 `ascendc_ops::ascendc_add` 已负责校验、分配输出和启动核函数：

```cpp
TORCH_LIBRARY(ascendc_ops, m)
{
    m.def("ascendc_add(Tensor x, Tensor y) -> Tensor");
}

TORCH_LIBRARY_IMPL(ascendc_ops, PrivateUse1, m)
{
    m.impl("ascendc_add", TORCH_FN(ascendc_ops::ascendc_add));
}
```

第一段定义接口契约；第二段为设备分发键绑定实现。torch_npu 使用 `PrivateUse1` 接入 NPU 后端。注册 NPU 实现不会同时产生 CPU 实现。

Schema 还涉及别名和原地修改语义。这里声明返回新 Tensor，因此实现不应偷偷修改 x 或 y。若要实现原地算子，应使用正确的别名/修改声明。

## 4. 动态库构建和 Python 验证

构建目标从独立程序的 `add_executable` 改为共享库：

```cmake
add_library(ascendc_ops SHARED add_custom.asc)
target_include_directories(ascendc_ops PRIVATE
    ${TORCH_INCLUDE_DIRS}
    ${TORCH_NPU_INCLUDE_DIRS}
)
target_link_libraries(ascendc_ops PRIVATE
    ${TORCH_LIBRARIES}
    torch_npu
    Python3::Python
)
```

此前还需完成 ASC、PyTorch、Python 与 torch_npu 的发现和路径配置，设置适配目标芯片的编译参数。头文件目录与库搜索目录是两回事；仅把目录保存到变量中，不会让链接器自动使用它。PyTorch、torch_npu、CANN 与编译 ABI 也需要兼容。

在 Linux NPU 环境中加载 `.so` 后，静态注册代码会把算子加入 dispatcher：

```python
import torch
import torch_npu

torch.ops.load_library("./libascendc_ops.so")
x = torch.rand((8, 2048), dtype=torch.float16)
y = torch.rand_like(x)
actual = torch.ops.ascendc_ops.ascendc_add(x.npu(), y.npu()).cpu()
expected = x + y
torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3)
```

这里假定 Kernel 支持 float16、连续输入和该形状。`rtol/atol` 应随算子数值特性制定，不是所有算子的统一合格线。测试还应覆盖负值、零、不同长度、尾块和不支持的输入是否被正确拒绝。

单个示例中固定 8 个 block，不代表所有输入或芯片都应固定 8 个核。Tiling 应从输入规模与硬件资源推导。

## 5. pybind11 接入

![pybind11 绑定流程](.assets/AscendC进阶/torch-19.png)

同一个 C++ 包装函数也可以绑定成 Python 扩展：

```cpp
PYBIND11_MODULE(ascendc_ops, m)
{
    m.def("ascendc_add", &ascendc_ops::ascendc_add);
}
```

```python
import torch
import torch_npu
import ascendc_ops

x = torch.rand((8, 2048), device="npu", dtype=torch.float16)
y = torch.rand_like(x)
z = ascendc_ops.ascendc_add(x, y)
```

扩展文件名与模块初始化名要匹配，模块所在目录必须可被 Python 导入。其加载方式与 `torch.ops.load_library` 不同：前者加载 Python 扩展模块，后者加载带有算子注册逻辑的动态库。

## 6. 从能调用到可训练、可入图

前向计算正确，并不自动支持反向传播。训练用途要提供或注册正确的 autograd 规则，并检查梯度；不要把 `requires_grad=True` 当成已经实现自定义反向。

进入图编译还需要 Meta/Fake 信息，使编译器在不读取真实数据时推导输出描述；TorchAir 路径还需要相应转换与 GE 执行实现。由此可把接入工作分为三层：**算子计算正确 → 框架调用正确 → 图编译及训练语义正确**。

实现参考：[torch.library 路径](https://gitcode.com/cann/asc-devkit/tree/master/examples/01_simd_cpp_api/02_features/00_framework/00_pytorch/torch_library)、[pybind11 路径](https://gitcode.com/cann/asc-devkit/tree/master/examples/01_simd_cpp_api/02_features/00_framework/00_pytorch/pybind)。