用途:把已有可编译运行的 Ascend C <<<>>> 直调工程改造成 PyTorch 扩展,使 Python 侧可通过 torch.ops.npu.xxx() 调用 NPU 算子。前置条件为已有直调工程、已安装 torch 与 torch_npu、CANN 环境已配置(ASCEND_HOME_PATH 已设置)。
主要能力:改造后产出 op_kernel/(xxx_tiling.h、xxx_kernel.asc)、op_host/(xxx.asc、data_utils.h)、op_extension/(xxx_torch.cpp、register.cpp、ops.h)与 CMakeLists.txt(双 target:可执行文件 + libxxx_ops.so)。调用链为 torch.ops.npu.xxx → PyTorch dispatch(PrivateUse1)→ ascend_kernel::xxx_torch() → ComputeTiling、aclrtMemcpy 搬 Tiling、kernel 函数调用 → NPU 执行。
使用方式:按 Step 0 拆分 kernel 与 host(仅单文件 .asc 需要),Step 1 建 ops.h,Step 2 写 xxx_torch.cpp(acl/acl.h 放在 torch 头文件前,参数校验用 is_privateuseone()),Step 3 在 register.cpp 中做 TORCH_LIBRARY_FRAGMENT 签名、TORCH_LIBRARY_IMPL PrivateUse1 绑定与 Meta 后端注册,Step 4 更新 CMakeLists.txt(CMAKE_CXX_STANDARD 17,npu-arch 按芯片:A2/A3=dav-2201,A5=dav-3510),Step 5 执行 cmake -S . -B build && cmake --build build -j4,Step 6 用 torch.ops.load_library("build/libxxx_ops.so") 加载并测试。
限制与注意:xxx_torch.cpp 由 C++ 编译器编译,不能用 <<<>>>,须手动管理 stream 同步,推荐 stream(true) 清 queue,禁止 stream(false) 直接调用、lambda 内传 NPUStream + OpCommand、zeros_like 创建输出;输出用 at::empty_like。Meta 后端为 torch.compile/fx 必需。已注册算子天然支持 aclgraph。多算子可在一个 .so 注册,也可加载多个 .so。具体模板、反模式与排错细节见资料所列 references 与官方示例路径。