Fast Kernel Launch：基于 Ascend C 与 PyTorch Extension 的 NPU 自定义算子极速开发指南

• 人工智能
• 算子库
• 深度学习
• CANN
• Ascend

导读

在 CANN（Compute Architecture for Neural Networks）生态中，为 NPU 编写自定义算子的传统路径往往需要分别维护算子原型定义、Tiling 策略、Kernel 实现与框架适配层，工程链路长、交付成本高。本篇文章聚焦于 ops-nn 仓库中 examples/fast_kernel_launch_example 这一实战示例，系统讲解如何仅用一个 C++ 文件完成 AI Core 算子的开发、编译与 PyTorch 框架适配，并借助<<<>>>核函数启动语法实现简单高效的调用。读完本文，你将掌握从环境搭建、Wheel 构建安装、Python 侧调用，到新增算子全流程的完整实战能力，并理解其底层源码级原理（Schema 注册、Meta 函数、Ascend C Kernel、torch.ops分发机制）。

什么是 Fast Kernel Launch

Fast Kernel Launch 是 ops-nn 仓库提供的一种单交付件（Single Artifact）算子开发范式。它的核心思想是：把传统算子开发中分散在不同目录、不同文件中的逻辑——算子 Schema、Meta 函数（InferShape/InferDtype）、Ascend C Kernel、NPU 调用与注册——全部收敛到一个 C++ 文件中，再通过 PyTorch C++ Extension 机制打包成一个 Python 可导入的 Wheel 包。

其核心优势可以概括为两点：

• 单交付件：一个文件完成算子开发和 PyTorch 框架适配，无需额外维护算子原型（.ini/.json）与 Tiling 分离工程；
• 高效调用：使用 CUDA 风格的<<<>>>语法启动核函数，流程简单高效，PyTorch 侧以torch.ops.<library_name>.<operator_name>方式直接调用。

从目录结构看，该示例仓库的布局如下（examples/fast_kernel_launch_example）：

fast_kernel_launch_example/ ├── ascend_ops/ # Python 包：__init__.py、ops.py、trans_conv3d_format.py ├── cmake/ # 构建模块：ascend.cmake、torch.cmake、torch_npu.cmake、func.cmake 等 ├── csrc/ # C++ 源码：add/、conv3d_custom/、extension.cpp │ └── add/ascend910b/ # 算子名/soc名 双层目录结构 ├── tests/ # pytest 用例：tests/add/test_add.py 等 ├── build_and_test.sh # 一键构建+测试脚本 ├── requirements.txt └── setup.py # 自定义 bdist_wheel / cmake_build / clean 命令

值得留意的是，示例除了最简的add算子外，还包含一个更复杂的 3D 卷积自定义算子（csrc/conv3d_custom），分别提供了ascend910b与ascend950两个 SoC 的实现，说明这套框架同样能承载工程级算子。

环境部署

前置要求

按照 examples/fast_kernel_launch_example/README.md 的说明，开发前需要准备：

此外，构建过程依赖昇腾 CANN Toolkit 的安装路径。从 cmake/ascend.cmake 的实现可以看到，构建脚本会按如下优先级定位 Toolkit：

1. 环境变量ASCEND_HOME_PATH（优先级最高）；
2. root 用户的默认路径/usr/local/Ascend/ascend-toolkit/latest或/usr/local/Ascend/latest；
3. 非 root 用户的$HOME/Ascend/ascend-toolkit/latest或$HOME/Ascend/latest。

若以上路径均找不到，CMake 会直接报错提示设置ASCEND_HOME_PATH。找到 Toolkit 后，构建脚本会使用其中的Bisheng 编译器（bisheng）作为 C/C++ 编译器与链接器，并自动引入ascendc（compiler/ascendc/include/...）、tikcpp、op_common、pkg_inc等头文件目录。

关键 SoC 型号与编译参数

由于 NPU Kernel 面向具体芯片架构编写，示例在csrc下使用算子名/SoC 名的双层目录组织源码，并支持通过环境变量控制编译目标：

以add算子的 SoC 目录为例，其 CMakeLists.txt 仅有一行：

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

dav-2201是 ascend910b 芯片对应的编译参数，更多 SoC 的取值可参考仓库内的 NpuArch 说明和使用指导（cann/ops-math 的 NpuArch 文档）。构建系统通过 cmake/func.cmake 中的recursive_add_subdirectory()宏，自动遍历csrc下所有子目录，并只把包含${NPU_SOC_VERSION}/CMakeLists.txt的算子纳入编译——这就是"同仓库多 SoC 并存"的工程组织方式。

安装步骤

进入示例目录后，按以下三步完成构建与安装：

1. 安装依赖

python3 -m pip install -r requirements.txt

requirements.txt 中除了 build、pyyaml、numpy、pytest 外，还通过--extra-index-url https://download.pytorch.org/whl/cpu指向 PyTorch CPU 版 wheel 源，便于在无 GPU 的构建机上获取 torch。

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

构建完成后，产物位于当前目录的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(bdist_wheel)类——它强制 wheel 使用 Stable ABI（abi3），使同一个包可以兼容多个 Python 版本（>=3.8），这也是该示例"一个包支持多 Python 版本"的关键设计。

3. 安装 Wheel 包

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

4.（可选）清理编译缓存：再次构建前，建议先执行python setup.py clean。该命令由 setup.py 中的CleanCommand实现，会删除build、dist、ascend_ops.egg-info目录及所有.pyc/.pyo文件。

如果希望一步到位，示例还提供了 build_and_test.sh 脚本：它会自动安装依赖、clean、构建 wheel、安装包，然后根据NPU_SOC_VERSION选择运行tests/conv3d_custom/ascend950或tests/conv3d_custom/ascend910b下的 pytest 用例。

构建系统内部是如何工作的

了解底层构建流程有助于排查问题。从 setup.py 可以看到：

• CMakeBuildCommand通过import torch/import torch_npu动态获取Torch_DIR与TORCH_NPU_PATH，读取环境变量NPU_SOC_VERSION（默认ascend910b），随后执行：

cmake -S . -B build -DCMAKE_BUILD_TYPE=Release \ -DTorch_DIR=... -DTORCH_NPU_PATH=... -DNPU_SOC_VERSION=... cmake --build build --parallel <cpu_count>

• 顶层 CMakeLists.txt 定义EXTENSION_MODULE_NAME（默认ascend_ops），并链接torch_npu、ascendcl、platform、register、tiling_api、runtime等库；同时通过-DEXTENSION_MODULE_NAME=ascend_ops将该宏传递给所有算子源码——这正是 C++ 代码中TORCH_LIBRARY_FRAGMENT(EXTENSION_MODULE_NAME, m)能正确生成库名的原因。
• csrc/extension.cpp 定义了一个名为_C的 Python C 扩展模块初始化函数PyInit__C。当 Pythonimport _C时，共享库被加载，其中所有TORCH_LIBRARY静态初始化器自动执行，从而把自定义算子注册进 PyTorch 系统。
• 各算子编译产物（*_obj对象库）通过OBJECTS_LIST汇总后，与extension.cpp一起链接为_C.abi3.so，并被拷贝到ascend_ops/Python 包目录下（参见顶层 CMakeLists 的add_custom_commandPOST_BUILD 步骤）。

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

安装完成后，NPU 算子对用户完全透明。以add算子为例，调用方式与原生 PyTorch 操作几乎一致：

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 npu_result = torch.ops.ascend_ops.add(x, y) # PyTorch Custom Operator Dispatch 机制: torch.ops.<library_name>.<operator_name> # 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. import ascend_ops触发注册：导入包时，ascend_ops/init.py 会先from . import _C（加载共享库，触发 C++ 侧全部TORCH_LIBRARY注册），再导入ops模块。若_C未正确安装，会抛出带提示的ImportError。
2. torch.ops.ascend_ops.add的分发路径：torch.ops是 PyTorch 的自定义算子命名空间。C++ 侧通过TORCH_LIBRARY_FRAGMENT(EXTENSION_MODULE_NAME, m)注册名为ascend_ops的库及算子 Schema，再通过TORCH_LIBRARY_IMPL(..., Meta, m)和TORCH_LIBRARY_IMPL(..., PrivateUse1, m)分别注册 Meta 实现与 NPU（PrivateUse1设备，即 Ascend NPU 的私有设备类型）实现。当输入张量位于 NPU 设备时，Dispatch 会自动落到 NPU 实现。
3. .npu()与torch_npu：张量通过torch_npu的npu()方法迁移到 NPU 设备，这是示例代码中import torch_npu的用途。

对应的单元测试位于 tests/add/test_add.py，其中包含两类用例：

• test_add_interface_exist：仅断言torch.ops.ascend_ops.add在命名空间中可见，用于守护"Schema 与 C++ 注册签名不一致导致算子未导出"这类常见故障；
• test_add_operator：对 19 种形状（从(1,)到(1000, 1000)、含多维(1,3,32,32)等）与 3 种 dtype（float32、float16、int32）做参数化测试，与 CPU 结果比对（int32 用torch.equal精确比较，浮点用torch.allclose(rtol=1e-4, atol=1e-4)）。

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

新增算子的整体流程非常轻量：只需提供一个 C++ 实现文件。下面以add算子为例逐步拆解。

第一步：建立目录结构

在csrc目录下使用算子名（如add）建立文件夹，再在其中使用目标 SoC 名建立子文件夹（Atlas A2 系列用ascend910b，Atlas A3 系列用ascend910_93，Ascend 950PR/950DT 用ascend950）：

csrc/ └── add/ └── ascend910b/ ├── CMakeLists.txt └── add.cpp

recursive_add_subdirectory()宏会自动发现该目录并纳入构建。

第二步：编写 SoC 目录下的 CMakeLists.txt

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

add_sources是 cmake/func.cmake 中定义的宏，其行为包括：

• 以当前目录的父目录名作为算子名（OP_NAME），生成${OP_NAME}_obj对象库目标；
• 用 Bisheng 编译器编译目录下所有.cpp（也可通过第二个参数传入自定义源文件列表）；
• 把编译参数拼成"--npu-arch=dav-2201 -xasc "，其中-xasc表示按 Ascend C 语言模式编译；
• 将对象目标追加到全局OBJECTS_LIST，最终由顶层 CMakeLists 统一链接。

第三步：编写单文件算子（add.cpp）

文件 csrc/add/ascend910b/add.cpp 完整包含了开发一个 AI Core 算子所需的全部四个模块：

• 算子 Schema 注册
• 算子 Meta Function 实现与注册
• 算子 Kernel 实现（Ascend C）
• 算子 NPU 调用实现与注册

3.1 头文件与命名空间

#include <ATen/Operators.h> #include <torch/all.h> #include <torch/library.h> #include "torch_npu/csrc/core/npu/NPUStream.h" #include "torch_npu/csrc/framework/OpCommand.h" #include "kernel_operator.h" #include "platform/platform_ascendc.h" #include <type_traits> namespace ascend_ops { // 当前项目为一个命名空间 namespace Add { // 建议每个算子有独立 namespace，防止全局变量污染

torch_npu的两个头文件分别用于获取当前 NPU 流（c10_npu::getCurrentNPUStream）和封装 aclnn 调用时序（at_npu::native::OpCommand）。

3.2 算子 Schema 注册

// Register the operator's schema TORCH_LIBRARY_FRAGMENT(EXTENSION_MODULE_NAME, m) { m.def("add(Tensor x, Tensor y) -> Tensor"); }

EXTENSION_MODULE_NAME在构建时被替换为ascend_ops，因此这条语句向 PyTorch 框架声明了ascend_ops::add这一算子的函数签名，框架据此知道存在这样一个算子。注意 schema 中参数名/类型必须与 C++ 注册签名严格一致，否则算子会被隐藏、无法从 Python 侧torch.ops访问——这正是测试用例test_add_interface_exist要守护的故障模式。

3.3 Meta 函数：InferShape + InferDtype

// Meta function implementation of Add 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; } // Register the Meta implementation TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, Meta, m) { m.impl("add", add_meta); }

Meta 函数只做形状与类型推导、不做实际计算：add的输出形状与输入一致，因此直接torch::empty_like申请空张量即可。它注册在Meta分发键下，框架在执行算子计算前即可获知输出所需空间，后续可支持torch.compile/ AutoGrad / AclGraph 等图加速能力。

3.4 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 }

__global__ __aicore__是 Ascend C 核函数的标准声明方式，GM_ADDR表示全局内存（Global Memory）地址。示例仓库中的完整实现（add.cpp）展示了 AI Core 编程的三个经典阶段，值得逐段研读：

• 流水线与队列：使用AscendC::TPipe管理流水线，TQue<QuePosition::VECIN, 2>/TQue<QuePosition::VECOUT, 2>声明输入/输出队列（深度 2，对应双缓冲）；pipe.InitBuffer(..., PIPELINE_DEPTH, tileSize)为队列分配本地内存（UB）。
• 分块（Block）划分：通过AscendC::GetBlockIdx()让每个 AI Core 定位到自己负责的数据段：xGm.SetGlobalBuffer((__gm__ T*)x + blockLength * GetBlockIdx())；并用currentBlockLength处理数据量不能整除 Core 数时的尾部情况。
• Tile 内三段流水：每个 tile 依次执行CopyIn（AscendC::DataCopyPad从 GM 拷入本地）、Compute（AscendC::Add(zLocal, xLocal, yLocal, n)向量加法）、CopyOut（DataCopyPad写回 GM），并通过AllocTensor/EnQue/DeQue/FreeTensor管理队列缓冲区，实现数据传输与计算的重叠。代码对"完整 tile"与"最后一个不完整 tile"（tailTileElementNum）分别处理。

3.5 Tiling 计算（calc_tiling_params）

Kernel 之外，add.cpp 还实现了calc_tiling_params，它决定了算子如何分块并行：

std::tuple<int64_t, int64_t, int64_t> calc_tiling_params(int64_t totalLength) { constexpr static int64_t MIN_ELEMS_PER_CORE = 1024; // 每个 AI Core 最少处理的元素数 constexpr static int64_t PIPELINE_DEPTH = 2; // 流水线深度 constexpr static int64_t BUFFER_NUM = 3; // 缓冲区数量（双缓冲） // 通过平台接口获取 UB 大小与 AI Core 数量 auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance(); uint64_t ubSize; ascendcPlatform->GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); int64_t coreNum = ascendcPlatform->GetCoreNumAiv(); ... int64_t numBlocks = std::min(coreNum, (totalLength + MIN_ELEMS_PER_CORE - 1) / MIN_ELEMS_PER_CORE); int64_t blockLength = (totalLength + numBlocks - 1) / numBlocks; int64_t tileSize = ubSize / PIPELINE_DEPTH / BUFFER_NUM; return std::make_tuple(numBlocks, blockLength, tileSize); }

它返回三元组(numBlocks, blockLength, tileSize)，分别代表：实际使用的 AI Core 数量、每个 Core 处理的元素数、每次 tile 处理的数据块大小。其中tileSize依据 UB 容量、流水线深度与缓冲区数量推导，保证队列缓冲区放得下且流水线可重叠。

3.6 NPU 调用接口（add_npu）

torch::Tensor add_npu(const torch::Tensor &x, const torch::Tensor &y) { // OptionalDeviceGuard 确保后续操作在正确的设备上下文执行，作用域结束后自动恢复 const c10::OptionalDeviceGuard guard(x.device()); auto z = add_meta(x, y); // 复用 Meta 函数确定输出 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); // 计算 Tiling 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", // 根据不同的数据类型，调用不同的 NPU Kernel 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); }) AT_DISPATCH_CASE(torch::kFloat16, [&] { using scalar_t = half; add_kernel<scalar_t><<<numBlocks, nullptr, stream>>>(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize); }) AT_DISPATCH_CASE(torch::kInt32, [&] { using scalar_t = int32_t; add_kernel<scalar_t><<<numBlocks, nullptr, stream>>>(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize); }) ); return 0; }; // 需使用 RunOpApi/RunOpApiV2 接口调用，保证时序与 TorchNPU 调用 aclnn 接口一致 at_npu::native::OpCommand::RunOpApi("Add", acl_call); return z; }

这段代码是"Host 侧调用"的完整范式，要点包括：

1. 计算输出的 Shape/Dtype：直接调用前面注册过的add_meta，一处实现、多处复用；
2. 计算 Tiling：依据x.numel()得到numBlocks/blockLength/tileSize；
3. 调用 NPU Kernel：add_kernel<scalar_t><<<numBlocks, nullptr, stream>>>即文档强调的<<<>>>语法——第一参数是 Grid（使用的 AI Core 数量），第三参数是 NPU 流；外层用AT_DISPATCH_SWITCH按 dtype 分发到 float32 / float16 / int32 三个特化；
4. 保证执行时序：整个 Kernel 启动闭包通过at_npu::native::OpCommand::RunOpApi("Add", acl_call)提交，与 TorchNPU 调用 aclnn 接口保持一致的入队/同步语义。

3.7 注册 NPU 实现

// Register the NPU implementation TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, PrivateUse1, m) { m.impl("add", add_npu); }

PrivateUse1是 PyTorch 为第三方私有设备（此处为 Ascend NPU）保留的分发键。注册后，当算子的输入张量位于 NPU 设备时，PyTorch 自动 Dispatch 到add_npu，对用户完全透明。

第四步：重新构建并验证

1. 参考上文"安装步骤"重新构建 Wheel 包并安装；
2. 基于 pytest 测试算子 API，直接参考 tests/add/test_add.py 的写法：先做接口存在性断言，再用多形状多 dtype 的参数化用例与 CPU 结果对拍。

工程延伸：从单算子到更复杂的算子

Fast Kernel Launch 框架并不局限于逐元素算子。示例仓库中的 3D 卷积自定义算子展示了更复杂的工程实践（csrc/conv3d_custom）：

• 多 SoC 并存：ascend910b/与ascend950/各有一套实现（Kernel 与 Torch 适配文件分离，如conv3d_v2_kernel.cpp/h、conv3d_v2_torch.cpp），由构建系统按NPU_SOC_VERSION自动选择；
• Python 层封装与格式转换：ascend_ops/ops.py 中的conv3d_custom根据torch_npu.npu.get_device_name()判断设备：Ascend950 直接使用 NCDHW 格式调用conv3d_v2_custom；Ascend910B 则需要借助 trans_conv3d_format.py 完成 NCDHW ↔ NDC1HWC0 / FRACTAL_Z_3D 的格式转换（C0=16），并支持 HF32 混合精度开关；
• 测试分层：tests/conv3d_custom 下按 SoC 组织 pytest 用例，ascend950 侧还配合example_case.csv提供案例化测试数据。

常见问题排查

结合源码可以给出以下排查建议：

• Cannot import _C：说明共享库未正确安装或路径异常。重新执行pip install dist/*.whl --force-reinstall --no-deps，并确认ascend_ops/包内存在_C.abi3.so（构建时由 POST_BUILD 拷贝）。
• torch.ops.ascend_ops.add不存在：多为 Schema 声明与 C++ 注册签名不一致（参数名、类型或重载不匹配），参考test_add_interface_exist的守护思路排查注册段。
• CMake 报找不到 Ascend Toolkit：确认ASCEND_HOME_PATH已设置，或 Toolkit 位于默认路径。
• 换芯片型号后 Kernel 行为异常：确认NPU_SOC_VERSION与--npu-arch与目标芯片匹配，并切换到对应 SoC 目录下的实现。
• 重新构建出现旧产物：先执行python setup.py clean清理build、dist与 egg-info。

总结

Fast Kernel Launch 为 NPU 自定义算子开发提供了一条"极简路径"：一个 C++ 文件承载 Schema、Meta、Kernel、Host 调用四件套，配合 PyTorch C++ Extension 与<<<>>>语法，即可将 Ascend C 算子无缝接入 PyTorch 生态，以torch.ops.ascend_ops.add的形式被上层直接调用。其构建体系（SoC 目录约定 +add_sources宏 + abi3 Wheel）在保证单交付件便捷性的同时，也兼顾了多芯片、多 Python 版本与工程级算子的可扩展性，是理解 CANN 算子开发与 PyTorch 适配衔接的优质入门范本。

• 人工智能
• 算子库
• 深度学习
• CANN
• Ascend