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 CAI Core 编程语言 PyTorch ExtensionC 扩展开发自定义 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 调用四大部分正是单交付件设计的最佳体现。二、环境部署与前置条件开始前需要完成基础环境搭建具体要求如下依赖项版本/说明CANN 基础环境参考 docs/zh/install/quick_install.md 完成部署gcc9.4.0 及以上python3.8 及以上PyTorchtorch2.6.0TorchNPU与 torch 版本对应的 torch_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.txtrequirements.txt 内容如下--extra-index-url https://download.pytorch.org/whl/cpu build pyyaml numpy2 pytest其中build用于构建 Wheelpyyaml为构建脚本依赖numpy2保证与 PyTorch 兼容pytest用于运行算子测试。3.2 设置编译款型并构建 Wheel# NPU_SOC_VERSION 设置编译款型 # Atlas A2 系列产品使用 ascend910b默认 # Atlas A3 系列产品使用 ascend910_93 # Ascend 950PR / Ascend 950DT 产品使用 ascend950 export NPU_SOC_VERSIONascend910b # -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, dtypetorch.float32).npu() y torch.randn(10, 32, dtypetorch.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(), rtol1e-6) print(Verification successful!)调用链背后包含三层机制Python 包导入ascend_ops/init.py 中执行from . import _C触发加载编译产物_C.abi3.so静态初始化csrc/extension.cpp 定义PyInit__C创建一个空模块其注释明确指出Python 导入该.so的目的就是触发其中TORCH_LIBRARY静态初始化器执行算子注册算子分派注册到torch.ops.ascend_ops命名空间下的add在输入为 NPUPrivateUse1张量时被分派到add_npu实现。五、开发指南新增一个算子以 add 为例5.1 目录与构建配置实现一个新算子只需一个 C 实现文件按如下结构组织csrc/ ├── add/ # 以算子名建立文件夹 │ └── ascend910b/ # 以目标 SoC 名建立子文件夹 │ ├── CMakeLists.txt # 编译参数 │ └── add.cpp # 算子完整实现建议以算子名为文件名 └── extension.cpp # PyInit__C 空模块SoC 目录下的 CMakeLists.txt 内容极为精简add_sources(--npu-archdav-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包含四个模块算子 Schema 注册告诉 PyTorch 框架存在该算子及其签名Meta Function 实现与注册InferShape InferDtype推导输出形状不做实际计算算子 Kernel 实现使用 Ascend C API 编写面向特定 SoC 的核函数算子 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_NAMEascend_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 给出了完整实现其核心设计可归纳为三点流水线PipelineTPipeTQueQuePosition::VECIN/VECOUT, PIPELINE_DEPTH构建三阶段流水线每个 tile 依次执行CopyInGM→Local→ ComputeAscendC::Add 向量加→ CopyOutLocal→GMPIPELINE_DEPTH 2让数据传输与计算重叠双缓冲/多缓冲BUFFER_NUM 3队列缓冲区数量与流水线深度共同决定并发度数据分块DataCopyPadDataCopyExtParamsblockCount、blockLen、srcStride、dstStride完成带边界处理的数据搬运完整 tile 与尾部残块tailTileElementNum分别处理保证任意形状输入均正确。④ NPU 调用与注册启动 Kerneltorch::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_kernelscalar_tnumBlocks, 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 KernelnumBlocks, 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由 UBUnified 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 进行验证该测试文件包含两类用例接口存在性测试断言torch.ops.ascend_ops命名空间中存在add用于防护因 Schema 与 C 注册签名不匹配参数名、类型、重载导致算子未导出到 Python 的常见问题功能正确性测试通过pytest.mark.parametrize组合 19 种形状从(1,)到(1000, 1000)、(8, 3, 128, 128)等与 3 种 dtypefloat32 / float16 / int32将 NPU 结果与 CPU 上a b的结果对比浮点用torch.allclose(rtol1e-4, atol1e-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_API0x03080000与 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
