SuperAPI AI Directory
Partner with us中文AI newsFavorites
Skillsexternal-cannbot-ops-torch-ascendc-op-extension

external-cannbot-ops-torch-ascendc-op-extension

Coding & development

View repository

Updated 10-15 17:19

About this skill

Ascend C Operator TORCH_LIBRARY Integration Guide

Prerequisites

  • An existing compilable and runnable Ascend C <<<>>> direct-call project
  • Environment already has: torch, torch_npu installed
  • CANN environment already configured (ASCEND_HOME_PATH already set)

Target Output

Regardless of the source project's directory structure, the following files are required after modification:

├── op_kernel/
│   ├── xxx_tiling.h         (Tiling 结构体 + 常量，kernel 和 host 共用)
│   └── xxx_kernel.asc       (纯 kernel 代码，引用 xxx_tiling.h)
├── op_host/
│   ├── xxx.asc              (host+main 可执行文件入口，引用 xxx_kernel.asc)
│   └── data_utils.h         (文件读写工具，位置以 include 路径可达为准)
├── op_extension/
│   ├── xxx_torch.cpp        (PyTorch 接入层：torch::Tensor 接口 + Tiling 计算 + kernel launch)
│   ├── register.cpp         (TORCH_LIBRARY 注册，含 Meta backend)
│   └── ops.h                (函数声明)
└── CMakeLists.txt           (双 target：可执行文件 + libxxx_ops.so)

Existing files are kept as-is; missing ones are added step by step according to the Steps.

Architecture Principles

文件依赖:
  op_kernel/xxx_tiling.h              (单一数据源：Tiling 结构体 + 常量)
    ├── op_kernel/xxx_kernel.asc      (#include "xxx_tiling.h")
    │   └── op_host/xxx.asc           (#include "../op_kernel/xxx_kernel.asc"，间接引入)
    └── op_extension/xxx_torch.cpp (#include "../op_kernel/xxx_tiling.h")

Python 调用链:
  torch.ops.npu.xxx(x1, x2)
      │
      ▼ PyTorch dispatch（PrivateUse1 后端）
  ascend_kernel::xxx_torch()         ← op_extension/xxx_torch.cpp 实现
      │
      ├─ ComputeTiling()             ← 计算 Tiling 参数
      ├─ aclrtMemcpy(tiling→device)  ← Tiling 搬到 device
      └─ xxx_kernel(...)             ← 函数调用（等价于 <<<>>>）
            │
            ▼ NPU 执行

Stream Synchronization Modes

Core Background: The NPU's taskqueue is a device-side task queue. "Flush queue" means waiting for existing tasks in the queue to complete before executing the current task; "enqueue" means placing the current task into the queue for sequential execution.

This solution uses a CMake multi-file architecture,xxx_torch.cpp compiled by the C++ compiler (not ASC), thereforecannot use <<<>>> syntax, and can only call the kernel as a normal C function. Under this architecture, stream synchronization must be managed manually.

Recommended mode:stream(true) flush queue

auto aclStream = c10_npu::getCurrentNPUStream().stream(true);
xxx_kernel(blockNum, nullptr, aclStream, ...);

stream(true) flushes the queue before returning aclrtStream (waits for previous tasks to complete), ensuring no out-of-order execution. For output allocation, using at::empty_like（zeros_like triggers NPU operations and may cause out-of-order execution).

Forbidden Modes

Anti-pattern Risk Fix
stream(false) + direct call Out-of-order: queue not flushed, kernel executes before previous operations Use stream(true)
Inside lambda, pass NPUStream + OpCommand Deadlock: queue waits for lambda, lambda waits for empty queue Inside lambda, use aclrtStream
zeros_like to create output Out-of-order: zeros_like is enqueued but kernel is not enqueued Use empty_like

Quick Reference of 8 Launch Modes

# Mode Verdict
1 Pass NPUStream object + <<<>>> Recommended(single-file .asc)
2 stream(true) + <<<>>> / function call Correct（The only applicable one in this solution）
3 stream(false) + OpCommand + <<<>>> Recommended(single-file .asc)
4 stream(false) + direct launch Anti-pattern: out-of-order
5 NPUStream + OpCommand Anti-pattern: deadlock
6 stream(true) + OpCommand Correct(single-file .asc)
7 stream(true) + zeros_like Anti-pattern: out-of-order
8 stream() with no arguments Correct(equivalent to mode 2)

For complete code, principles, and deadlock flow of each mode, see references/anti_patterns.md。

Modification Steps

Step 0: Split kernel and host (only needed for single-file .asc projects)

Skip condition: kernel and host are already in different files.

For single-file .asc projects (kernel + host + main in one file), perform the following operations:

  1. Extract op_kernel/xxx_tiling.h: from the original .asc extract the following content into a new file op_kernel/xxx_tiling.h：

    • Named constants (constexpr, such as TILE_LENGTH、DOUBLE_BUFFER）
    • Tiling data structure (struct XxxTilingData）

    xxx_tiling.h contains only pure C/C++ syntax, without __aicore__、__gm__ and other ASC keywords. All compilation units (kernel, host main, torch host) share this file.

  2. Extract op_kernel/xxx_kernel.asc: cut the following content from the original .asc into a new file op_kernel/xxx_kernel.asc：

    • #include "kernel_operator.h" and the ASC-specific header file
    • #include "xxx_tiling.h"(replacing the original inline constants and structures)
    • Kernel class (class KernelXxx）
    • kernel function entry (extern "C" __global__ __vector__ void xxx_kernel(...)）

    xxx_kernel.asc contains only ASC syntax, without any host code (main, file read/write, iostream, etc.).

  3. Modify the original .asc(moved to the op_host/ directory): delete the already-cut kernel content, and change it to:

    #include "../op_kernel/xxx_kernel.asc"    // 引用拆出的纯 kernel 文件
// ... 以下保留原有的 host 侧代码（main 等），不做任何修改

  4. Ensure the include path is correct：op_host/xxx.asc The auxiliary files that depends on (such as data_utils.h) need to be found within the paths covered by target_include_directories in CMakeLists.txt.

  5. Clean up old files: delete the original .asc files that have been replaced (no longer referenced for compilation after the split), to avoid leaving orphaned files that confuse maintenance.

  6. Verification: compile the executable target to confirm the split has not broken functionality.

Step 1: Create op_extension/ops.h(function declarations)

Skip if it already exists.

Read the template:templates/ops_template.h

In the ascend_kernel namespace, declare the operator function signature, for use by both op_extension/register.cpp and op_extension/xxx_torch.cpp .

Step 2: Create op_extension/xxx_torch.cpp(PyTorch integration layer implementation)

This is the core file. After reading the source project code, modify it based on the template:templates/torch_template.cpp

The template uses #include "../op_kernel/xxx_tiling.h" to include the Tiling structure and constants (single source of truth).ComputeTiling The function is anexample placeholder, and should be replaced with the actual computation logic of the source project.

Key design decisions：

  • #include "acl/acl.h" placed before the torch header filebefore— torch's macros will interfere with ACL header file parsing
  • #include "../op_kernel/xxx_tiling.h" Include the Tiling structure and constants
  • Parameter validation uses x1.is_privateuseone()(C++ API), not is_npu()
  • Framework stream:must use stream(true)(clear queue, prevent out-of-order),forbidden stream(false). See the "Stream synchronization mode" section above for details

Adaptation guide: Tiling computation, kernel entry declarations, kernel invocation methods, etc. need to be extracted from the source code. See references/operation_checklist.md。

Step 3: Create op_extension/register.cpp(operator registration)

Skip if it already exists.

Read the template:templates/register_template.cpp

Must include three parts of registration:

// 1. 算子签名定义
TORCH_LIBRARY_FRAGMENT(npu, m) {
    m.def("xxx(Tensor x1, Tensor x2) -> Tensor");
}

// 2. NPU 后端实现绑定
TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) {
    m.impl("xxx", TORCH_FN(ascend_kernel::xxx_torch));
}

// 3. Meta 后端注册（torch.compile / fx 必需）
at::Tensor xxx_meta(const at::Tensor& x1, const at::Tensor& x2) {
    return at::empty_like(x1);  // 简单算子：直接推导；多输出/shape变化需按实际逻辑
}
TORCH_LIBRARY_IMPL(npu, Meta, m) {
    m.impl("xxx", &xxx_meta);
}

namespace selection: by default use npu（torch.ops.npu.xxx). If you need to avoid conflicts with torch_npu built-in operators, you can switch to ascendc_ops（torch.ops.ascendc_ops.xxx）。

Step 4: Update CMakeLists.txt

Read the template:templates/CMakeLists_template.cmake

Dual-target structure: keep the original executable Target, add add_library(xxx_ops SHARED ...) to compile .so。set(CMAKE_CXX_STANDARD 17) cannot be omitted.npu-arch Modify according to the actual chip: A2/A3=dav-2201, A5=dav-3510.

For how to handle CMakeLists when the kernel file format is different, see references/operation_checklist.md。

NoteIf there is already a PyTorch-related Target in CMakeLists.txt, directly replace it with the new Target 2 configuration; do not keep the old target.

Step 5: Compilation Verification

cmake -S . -B build && cmake --build build -j4

After a successful compilation, the following will be produced:

  • build/<operator_name>_custom — Original executable (direct-call verification)
  • build/libxxx_ops.so — PyTorch extension library

Step 6: Python Testing

import torch
import torch_npu

torch.ops.load_library("build/libxxx_ops.so")

# 根据算子签名调整输入数量、dtype 和验证逻辑
x1 = torch.randn(1024, 4096, dtype=torch.float32).npu()
x2 = torch.randn(1024, 4096, dtype=torch.float32).npu()

y = torch.ops.npu.xxx(x1, x2)    # 单输入算子改为 torch.ops.npu.xxx(x1)

assert y.is_npu                    # 注意：Python 中是 property 不是方法
assert torch.allclose(y.cpu(), x1.cpu() + x2.cpu())  # 验证逻辑按算子语义调整

For multi-operator scenarios, load multiple .so as needed:

torch.ops.load_library("/path/to/libadd1_ops.so")
torch.ops.load_library("/path/to/libmul1_ops.so")

y1 = torch.ops.npu.add1(x1, x2)
y2 = torch.ops.npu.mul1(x1, x2)

Multi-operator extensions

One .so can register multiple operators:

register.cpp:
  TORCH_LIBRARY_FRAGMENT(npu, m) {
      m.def("op_a(Tensor x1, Tensor x2) -> Tensor");
      m.def("op_b(Tensor x1, Tensor x2) -> Tensor");
  }
  TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) {
      m.impl("op_a", TORCH_FN(ascend_kernel::op_a));
      m.impl("op_b", TORCH_FN(ascend_kernel::op_b));
  }
  TORCH_LIBRARY_IMPL(npu, Meta, m) {
      m.impl("op_a", &op_a_meta);
      m.impl("op_b", &op_b_meta);
  }

aclgraph Compatibility

Custom operators that have been registered with TORCH_LIBRARY natively support aclgraph, with no additional modification required. Usage:

# 方式1: NPUGraph
g = torch.npu.NPUGraph()
with torch.npu.graph(g):
    output = model(static_x, static_y)
g.replay()

# 方式2: make_graphed_callables
model = torch.npu.make_graphed_callables(model, (x, y))

# 方式3: npugraph_ex backend (需 torch_npu >= 7.3.0)
compiled = torch.compile(model, backend="npugraph_ex", fullgraph=True)

Prerequisite: The operator has correctly registered the Meta backend (see Step 3), and stream synchronization uses the stream(true) mode.

Pitfall Checklist

When encountering compilation/linking/runtime issues, read references/troubleshooting.md。

Reference Resources

  • Official aclgraph + TORCH_LIBRARY example:torchair/examples/_kernel_extension_aclgraph/torch_library/
  • Official TORCH_LIBRARY example:asc-devkit/examples/01_simd_cpp_api/02_features/02_framework/00_pytorch/torch_library/
  • PyTorch framework adaptation documentation:asc-devkit/docs/guide/编程指南/附录/AI框架算子适配/PyTorch框架.md
  • PyTorch Custom Operators：https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html