CANN ops-cv Fast Kernel Launch 实战：基于 PyTorch Extension 与 Ascend C 的自定义 NPU 算子开发指南

CANN ops-cv Fast Kernel Launch 实战：基于 PyTorch Extension 与 Ascend C 的自定义 NPU 算子开发指南

【免费下载链接】ops-cv本项目是CANN提供的图像处理、目标检测相关的算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-cv

导读

本文围绕 CANN ops-cv 仓库中的 Fast Kernel Launch 示例 展开，系统讲解如何在 CANN 生态下，用Ascend C（AI Core 编程语言）+ PyTorch Extension（C++ 扩展）开发自定义 NPU 算子，并像普通 PyTorch 算子一样被torch.ops直接调用。读完本文，你将掌握从环境准备、Wheel 构建安装、Python 端调用，到"单文件落地一个算子"（Schema 注册、Meta 推导、Ascend C Kernel、NPU 调用）的完整开发链路，并了解仓库中 add、upsample_nearest3d 两个算子示例的源码级实现细节。

一、示例定位与核心优势

Fast Kernel Launch 是 ops-cv 仓库中面向"快速算子开发"场景的完整工程示例，位于 examples/fast_kernel_launch_example。它的目标非常明确：用最少的工作量，让开发者把 Ascend C 算子接入 PyTorch 生态，从而直接利用torch.nn、torch.ops以及 NPU 上的自动设备管理能力。

相比传统算子交付流程（算子实现、框架适配层、注册机制相互分离，往往需要多个交付件），该示例强调两点核心优势：

• 单交付件：一个 C++ 文件即可同时完成算子开发与 PyTorch 框架适配，无需拆分多个模块分别交付；
• 高效调用：使用类似 CUDA 的<<<>>>语法直接启动核函数，调用流程简单高效，无需经过 aclnn 接口封装层。

从仓库实际代码看，csrc/add/ascend910b/add.cpp 一个文件内就包含了算子 Schema 注册、Meta 函数、Ascend C Kernel 与 NPU 调用四大部分，正是"单交付件"设计的最佳体现。

二、环境部署与前置条件

开始前需要完成基础环境搭建，具体要求如下：

TorchNPU 是 PyTorch 在昇腾 NPU 上的适配层，本示例编译时依赖它的头文件与库（torch_npu/csrc/core/npu/NPUStream.h、torch_npu/csrc/framework/OpCommand.h），运行时依赖它完成 NPU 设备管理与算子分派。

三、安装步骤：构建并安装 Wheel 包

3.1 安装依赖

进入示例目录并安装 Python 依赖：

cd examples/fast_kernel_launch_example python3 -m pip install -r requirements.txt

requirements.txt 内容如下：

--extra-index-url https://download.pytorch.org/whl/cpu build pyyaml numpy<2 pytest

其中build用于构建 Wheel，pyyaml为构建脚本依赖，numpy<2保证与 PyTorch 兼容，pytest用于运行算子测试。

3.2 设置编译款型并构建 Wheel

# NPU_SOC_VERSION 设置编译款型： # Atlas A2 系列产品使用 "ascend910b"（默认） # Atlas A3 系列产品使用 "ascend910_93" # Ascend 950PR / Ascend 950DT 产品使用 "ascend950" export NPU_SOC_VERSION=ascend910b # -n: non-isolated build（使用当前已存在的环境，不创建隔离构建环境） python3 -m build --wheel -n

关于款型参数的细节，可以从 setup.py 中得到印证：CMakeBuildCommand会读取环境变量NPU_SOC_VERSION，若未设置则回退到NPU_ARCH（此时会打印NPU_ARCH is deprecated, please use NPU_SOC_VERSION.的弃用提示），最终默认值为ascend910b，并通过-DNPU_SOC_VERSION=传入 CMake。

构建完成后，产物位于当前目录的dist文件夹下，产物命名为：

ascend_ops-1.0.0-${python_version}-abi3-${arch}.whl

• ${python_version}：当前环境的 Python 版本标签，例如 Python 3.8.3 对应cp38；
• ${arch}：CPU 架构。

产物中的abi3标签来自 setup.py 中的ABI3Wheel类：它强制将 Wheel 标记为cp38/abi3，使同一个 Wheel 可跨多个 Python 3.8+ 版本使用。

3.3 安装 Wheel 包

python3 -m pip install dist/*.whl --force-reinstall --no-deps

--force-reinstall确保覆盖旧版本，--no-deps跳过依赖安装（依赖已在前文安装）。

3.4 清理编译缓存（可选）

再次构建前建议执行：

python setup.py clean

该命令由 setup.py 中的CleanCommand实现，会删除build、dist、ascend_ops.egg-info目录及*.pyc、*.pyo缓存文件，避免增量构建时产生陈旧产物。

仓库还提供了一键脚本 build_and_test.sh，内部依次执行"安装依赖 → 清理 → 构建 Wheel → 安装 → 运行pytest tests/* -v"，可直接复现完整流程。

四、快速开始：像普通 PyTorch 算子一样调用

安装完成后，即可在 Python 中以普通 PyTorch 算子方式使用 NPU 算子。以 add 算子为例：

import torch import torch_npu import ascend_ops # 构建出的 python 包 # Initialize data on NPU x = torch.randn(10, 32, dtype=torch.float32).npu() y = torch.randn(10, 32, dtype=torch.float32).npu() # Call the custom NPU operator # PyTorch Custom Operator Dispatch 机制: torch.ops.<library_name>.<operator_name> npu_result = torch.ops.ascend_ops.add(x, y) # Verify against CPU ATen implementation cpu_x = x.cpu() cpu_y = y.cpu() cpu_result = cpu_x + cpu_y assert torch.allclose(cpu_result, npu_result.cpu(), rtol=1e-6) print("Verification successful!")

调用链背后包含三层机制：

1. Python 包导入：ascend_ops/init.py 中执行from . import _C，触发加载编译产物_C.abi3.so；
2. 静态初始化：csrc/extension.cpp 定义PyInit__C创建一个空模块，其注释明确指出：Python 导入该.so的目的，就是触发其中TORCH_LIBRARY静态初始化器执行算子注册；
3. 算子分派：注册到torch.ops.ascend_ops命名空间下的add，在输入为 NPU（PrivateUse1）张量时被分派到add_npu实现。

五、开发指南：新增一个算子（以 add 为例）

5.1 目录与构建配置

实现一个新算子只需一个 C++ 实现文件，按如下结构组织：

csrc/ ├── add/ # 以算子名建立文件夹 │ └── ascend910b/ # 以目标 SoC 名建立子文件夹 │ ├── CMakeLists.txt # 编译参数 │ └── add.cpp # 算子完整实现（建议以算子名为文件名） └── extension.cpp # PyInit__C 空模块

SoC 目录下的 CMakeLists.txt 内容极为精简：

add_sources("--npu-arch=dav-2201")

其中dav-2201是 ascend910b 芯片对应的编译参数。从仓库现有实现看，另一个算子示例 csrc/upsample_nearest3d/ascend910b/upsample_nearest3d_torch.cpp 采用同样的目录结构与构建方式，验证了该模式的通用性。csrc/CMakeLists.txt 通过recursive_add_subdirectory()自动递归收集各算子目录的源文件，累加到OBJECTS_LIST，最终由顶层 CMakeLists.txt 生成_C.abi3.so并拷贝到ascend_ops/包目录下。

5.2 单文件四要素：一个算子需要什么

一个完整的算子实现文件（csrc/add/ascend910b/add.cpp）包含四个模块：

1. 算子 Schema 注册：告诉 PyTorch 框架存在该算子及其签名；
2. Meta Function 实现与注册：InferShape + InferDtype，推导输出形状，不做实际计算；
3. 算子 Kernel 实现：使用 Ascend C API 编写面向特定 SoC 的核函数；
4. 算子 NPU 调用实现与注册：完成 Tiling 计算、Stream 获取与 Kernel 启动。

下面按这四部分逐一展开。

① Schema 注册：声明算子签名

TORCH_LIBRARY_FRAGMENT(EXTENSION_MODULE_NAME, m) { m.def("add(Tensor x, Tensor y) -> Tensor"); }

EXTENSION_MODULE_NAME是构建期宏，由顶层 CMakeLists.txt 通过-DEXTENSION_MODULE_NAME=ascend_ops注入。TORCH_LIBRARY_FRAGMENT以片段形式把算子加入ascend_ops库，这样torch.ops.ascend_ops.add才能被 Python 侧发现。

② Meta Function：形状推导

torch::Tensor add_meta(const torch::Tensor &x, const torch::Tensor &y) { TORCH_CHECK(x.sizes() == y.sizes(), "The shapes of x and y must be the same."); auto z = torch::empty_like(x); return z; } TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, Meta, m) { m.impl("add", add_meta); }

Meta 函数在 CPU 上执行，负责在真正计算前确定输出张量的形状、数据类型与所需空间。注册到Meta分派键后，后续可支撑torch.compile、AutoGrad、AclGraph 等图加速能力——框架无需真实执行计算即可完成图级推导。

③ Ascend C Kernel：核函数实现

template <typename T> __global__ __aicore__ void add_kernel(GM_ADDR x, GM_ADDR y, GM_ADDR z, int64_t totalLength, int64_t blockLength, uint32_t tileSize) { // kernel implementation }

仓库中 csrc/add/ascend910b/add.cpp 给出了完整实现，其核心设计可归纳为三点：

• 流水线（Pipeline）：TPipe+TQue<QuePosition::VECIN/VECOUT, PIPELINE_DEPTH>构建三阶段流水线，每个 tile 依次执行CopyIn（GM→Local）→ Compute（AscendC::Add 向量加）→ CopyOut（Local→GM），PIPELINE_DEPTH = 2让数据传输与计算重叠；
• 双缓冲/多缓冲：BUFFER_NUM = 3，队列缓冲区数量与流水线深度共同决定并发度；
• 数据分块：DataCopyPad+DataCopyExtParams（blockCount、blockLen、srcStride、dstStride）完成带边界处理的数据搬运，完整 tile 与尾部残块（tailTileElementNum）分别处理，保证任意形状输入均正确。

④ NPU 调用与注册：启动 Kernel

torch::Tensor add_npu(const torch::Tensor &x, const torch::Tensor &y) { const c10::OptionalDeviceGuard guard(x.device()); // 记录并在作用域后恢复设备上下文 auto z = add_meta(x, y); // 获取输出张量 auto stream = c10_npu::getCurrentNPUStream().stream(false); // 当前 NPU 流 int64_t totalLength, numBlocks, blockLength, tileSize; totalLength = x.numel(); std::tie(numBlocks, blockLength, tileSize) = calc_tiling_params(totalLength); auto x_ptr = (GM_ADDR)x.data_ptr(); auto y_ptr = (GM_ADDR)y.data_ptr(); auto z_ptr = (GM_ADDR)z.data_ptr(); auto acl_call = [=]() -> int { AT_DISPATCH_SWITCH(x.scalar_type(), "add_npu", AT_DISPATCH_CASE(torch::kFloat32, [&] { using scalar_t = float; add_kernel<scalar_t><<<numBlocks, nullptr, stream>>>(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize); }) // ... kFloat16 / kInt32 同理 ); return 0; }; at_npu::native::OpCommand::RunOpApi("Add", acl_call); // 保证与 TorchNPU 调用 aclnn 接口时序一致 return z; } TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, PrivateUse1, m) { m.impl("add", add_npu); }

该函数承担了 README 中提到的三件事：① 计算输出 Tensor（直接调用add_meta）；② 计算 Tiling（调用calc_tiling_params）；③ 启动 NPU Kernel（<<<numBlocks, nullptr, stream>>>语法）。

calc_tiling_params的源码实现（见 csrc/add/ascend910b/add.cpp）值得单独说明，它通过platform_ascendc::PlatformAscendCManager::GetInstance()查询硬件能力：

• MIN_ELEMS_PER_CORE = 1024：每个 AI Core 最少处理元素数，防止任务过小导致并行效率低；
• numBlocks = min(coreNum, (totalLength + MIN_ELEMS_PER_CORE - 1) / MIN_ELEMS_PER_CORE)：实际使用的 AI Core 数不超过数据量所需；
• blockLength = ceil(totalLength / numBlocks)：每个 AI Core 处理元素数；
• tileSize = ubSize / PIPELINE_DEPTH / BUFFER_NUM：由 UB（Unified Buffer）大小、流水线深度和缓冲数量共同决定每次搬运的数据块大小；
• 数据量为 0 等边界情况下，TORCH_CHECK(coreNum > 0)保证核数参数合法。

关键点在于OpCommand::RunOpApi("Add", acl_call)：Kernel 启动被包装进 lambda 后统一经at_npu::native::OpCommand执行，保证算子调用时序与 TorchNPU 调用 aclnn 接口的时序一致，避免流同步与内存生命周期问题。

最后注册到PrivateUse1分派键——这是 PyTorch 为昇腾 NPU 等私有设备预留的分派键，框架据此在输入张量位于 NPU 设备时自动分派到该实现。

5.3 测试验证

算子开发完成后，可参照 tests/add/test_add.py 用 pytest 进行验证，该测试文件包含两类用例：

1. 接口存在性测试：断言torch.ops.ascend_ops命名空间中存在add，用于防护因 Schema 与 C++ 注册签名不匹配（参数名、类型、重载）导致算子未导出到 Python 的常见问题；
2. 功能正确性测试：通过@pytest.mark.parametrize组合 19 种形状（从(1,)到(1000, 1000)、(8, 3, 128, 128)等）与 3 种 dtype（float32 / float16 / int32），将 NPU 结果与 CPU 上a + b的结果对比：浮点用torch.allclose(rtol=1e-4, atol=1e-4)，整型用torch.equal精确比对，并用@pytest.mark.skipif(not torch.npu.is_available(), ...)在无 NPU 环境自动跳过。

运行方式（需 NPU 环境）：

pytest tests/add -v

六、多算子扩展与 Python 封装

除 add 外，示例还包含第二个算子upsample_nearest3d（实现见 csrc/upsample_nearest3d/ascend910b/upsample_nearest3d_torch.cpp，测试见 tests/upsample_nearest3d/test_upsamplenearest3d.py）。它演示了带 List 参数（size: List[int]）的算子如何注册与调用，为开发输入含列表/标量参数的算子提供了参考模板。

如果需要为算子提供更友好的 Python 函数接口，可以在 ascend_ops/ops.py 中做薄封装：

def upsample_nearest3d(x: Tensor, size: List[int]) -> Tensor: """Performs upsample_nearest3d(x) in an efficient fused kernel""" return torch.ops.ascend_ops.upsample_nearest3d(x, size)

即：Python 层仅做类型标注与参数透传，实际执行仍由 C++ 注册的torch.ops.ascend_ops.*算子完成，Python 侧保持"零计算逻辑"。

七、构建系统内部机制速览

理解构建链路有助于排查编译问题，关键机制如下：

• SoC 选择：顶层 CMakeLists.txt 读取NPU_SOC_VERSION（默认ascend910b），各算子的add_sources("--npu-arch=...")参数需与所选 SoC 匹配（如 ascend910b 对应dav-2201）；
• 扩展名控制：目标库_C设置PREFIX ""、SUFFIX ".abi3.so"、Py_LIMITED_API=0x03080000，与 setup.py 的abi3标签呼应；
• 链接库：torch_npu、ascendcl、platform、register、tiling_api、runtime等 NPU 侧库；
• 产物拷贝：构建后通过add_custom_command(POST_BUILD)将_C.abi3.so拷贝到ascend_ops/包目录，保证from . import _C能加载到；
• 编译加速：CMakeBuildCommand使用os.cpu_count()作为并行编译任务数。

八、小结

Fast Kernel Launch 示例为 CANN ops-cv 用户提供了一条最短路径：一个文件、四条注册（Schema / Meta / Kernel / NPU 调用）、一条torch.ops.ascend_ops.*调用链，即可完成自定义 NPU 算子的开发、构建与验证。其核心价值在于把 Ascend C 的算子表达力与 PyTorch 的生态无缝衔接——单交付件降低了交付复杂度，<<<>>>语法保持了调用直观性，而OpCommand::RunOpApi则确保自定义算子与官方 aclnn 路径在时序语义上保持一致。

后续开发新算子时，可直接以 csrc/add/ascend910b/add.cpp 为蓝本：先复制目录骨架与 CMakeLists，再依次替换 Schema、Meta、Kernel 与 NPU 调用四个部分，最后参照 tests/add/test_add.py 补充参数化测试即可。

【免费下载链接】ops-cv本项目是CANN提供的图像处理、目标检测相关的算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-cv