SuperAPI AI Directory
Partner中文AI newsFavorites
Skillsexternal-gitcode-ascend-ascendc-operator-code-gen

external-gitcode-ascend-ascendc-operator-code-gen

Agents & models

View repository

Updated 10-15 16:58

About this skill

AscendC Operator Code Generation and Framework Adaptation

Generate op_host + op_kernel code based on the design document, register to the PyTorch framework, and call the ascendc-operator-compile-debug skill to complete compilation, installation, and precision testing.

Prerequisites: Design document ascend-kernel/csrc/ops/<op_name>/design.md is ready, and the directory structure has been created.

Workflow Overview

读取设计文档 → 加载 reference → 选择模板 → 生成 op_host + op_kernel
    → 框架适配 (ops.h + register.cpp + csrc/CMakeLists.txt)
    → 调用 ascendc-operator-compile-debug skill (编译 + 安装 + 测试)

Phase 1: Load Reference Documents

MANDATORY — READ BEFORE CODING: Read references/GUIDE.md, and load the corresponding reference file according to the operator type in the design document.Never skip this step.

Phase 2: Read the Design Document

From ascend-kernel/csrc/ops/<op_name>/design.md extract:

Extraction Item Design Document Section Purpose
Function signature + supported data types Prototype design op_host function prototype, kernel template parameters
Operator type Tiling split → Step 1 Select template (elementwise / row)
UB allocation table UB space allocation Derive bufferCoefficient, InitBuffer size
Computation pseudocode Kernel implementation Compute function logic

Phase 3: Select Template and Generate Code

Select the corresponding template according to the operator type, copy it to the project directory, and then modify it.

Template Selection

Operator type op_host template op_kernel template
Elementwise (ReLU, GELU, Add...) templates/elementwise_op_host.cpp templates/elementwise_op_kernel.cpp
Row processing (LayerNorm, Softmax...) templates/row_op_host.cpp templates/row_op_kernel.cpp

MANDATORY: Read the complete content of the corresponding template file, copy it to the target path, and then modify it.

Generation Steps

  1. Read the complete content of the selected op_host and op_kernel templates
  2. Write the content to ascend-kernel/csrc/ops/<op_name>/op_host/<op_name>.cpp and op_kernel/<op_name>.cpp(overwrite the skeleton placeholder files)
  3. Replace all placeholders (<op_name>, <OpName> , etc.)
  4. Modify according to the design document:
    • op_host: function signature, input validation,bufferCoefficient, EXEC_KERNEL_CMD parameters
    • op_kernel: Init's GM parameters, InitBuffer size, Compute computation logic

FP16/FP32 template branches

The Kernel template uses template <typename T> generics + if constexpr branches to handle different data types:

if constexpr (sizeof(T) == sizeof(float)) {
    // float32 直接计算
} else {
    // fp16/bf16: Cast 升精度 → fp32 计算 → Cast 降精度
}

The AscendC compiler supports C++17 if constexpr, which can be safely used for template branches. If compatibility with older compilers is needed, ordinary if with sizeof(T) == 4 can also be used (the compiler will optimize away dead branches).

Key Points for Template Modification

Hardware parameter acquisition (built into the template)

The template uses platform APIs to dynamically obtain hardware parameters, no longer hardcoded:

#include "tiling/platform/platform_ascendc.h"

auto ascendc_platform = platform_ascendc::PlatformAscendCManager::GetInstance();
int64_t coreNum = static_cast<int64_t>(ascendc_platform->GetCoreNumAiv());
uint64_t ubSize;
ascendc_platform->GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
// 需要 workspace 时:
// uint64_t sysWorkspaceSize = ascendc_platform->GetLibApiWorkSpaceSize();

bufferCoefficient derivation

Derive from the UB allocation table in the design document—the total size of all Buffers = bufferCoefficient * tileLength：

  • Sum the "Total Size" column of all Buffers in the UB allocation table
  • The result has the form tileLength * N, then bufferCoefficient = N
  • When the coefficients differ for different dtypes, use if (dtypeSize == 2) {...} else {...} branch

FP16/BF16 precision promotion

Both FP16 and BF16 must be promoted to FP32 for computation and then converted back. The template has already marked the insertion locations for the precision promotion code; just uncomment and fill in the computation logic.

GM ↔ UB transfer (DataCopyPad)

Production code must use DataCopyPad(do not use DataCopy for GM↔UB transfers):

// CopyIn: GM → UB
AscendC::DataCopyExtParams copyInParams{1,
    static_cast<uint32_t>(curTileLength * sizeof(T)), 0, 0, 0};
AscendC::DataCopyPadExtParams<T> padParams{false, 0, 0, static_cast<T>(0)};
AscendC::DataCopyPad(xLocal, xGm[progress * tileLength], copyInParams, padParams);

// CopyOut: UB → GM
AscendC::DataCopyExtParams copyOutParams{1,
    static_cast<uint32_t>(curTileLength * sizeof(T)), 0, 0, 0};
AscendC::DataCopyPad(yGm[progress * tileLength], yLocal, copyOutParams);

Note: The template files may use DataCopy to simplify examples; production code should replace it with DataCopyPad.

ReduceSum/ReduceMax note

Reduction operations may modify the source tensor, so it must be backed up first:

AscendC::Adds(backup, src, 0.0f, len);  // 备份
AscendC::ReduceSum<float, true>(result, backup, sharedTmp, dimLen);

Phase 4: Framework adaptation

4.1 Update csrc/ops.h

In namespace ascend_kernel , add the function declaration:

<返回类型> <op_name>(<参数列表>);

4.2 Update csrc/register.cpp

In TORCH_LIBRARY_FRAGMENT(npu, m) , add the schema:

m.def("<op_name>(<schema_参数>) -> <返回类型>");

In TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) , add the implementation:

m.impl("<op_name>", TORCH_FN(ascend_kernel::<op_name>));

Schema syntax quick reference:

C++ type Schema type Example
const at::Tensor & Tensor Tensor self
at::IntArrayRef int[] int[] kernel_size
int64_t int int dim=-1
double float float eps=1e-5
bool bool bool flag=False
c10::optional<at::Tensor> Tensor? Tensor? weight=None
c10::optional<int64_t> int? int? divisor=None

4.3 Update csrc/CMakeLists.txt

In FILE(GLOB OP_SRCS ...) , add the host source file:

${PROJECT_OP_SRC_BASE}/ops/<op_name>/op_host/<op_name>.cpp

In ascendc_library(no_workspace_kernel STATIC ...) , add the kernel source file:

${PROJECT_OP_SRC_BASE}/ops/<op_name>/op_kernel/<op_name>.cpp

Phase 5: Compilation, installation, and testing

MANDATORY — follow the ascendc-operator-compile-debug skill process：

  1. Load the environment: source the CANN environment variables, activate the conda environment
  2. Ensure build.sh has execute permission:chmod +x build.sh
  3. Compile the project:bash build.sh
  4. Verify that compilation succeeds:ls output/ascend_kernel*.whl
  5. Install the whl package:pip install output/ascend_kernel*.whl --force-reinstall --no-deps
  6. Generate the test file tests/test_<op_name>.py
  7. Run the functional test:python tests/test_<op_name>.py
  8. Run the precision test:pytest tests/test_<op_name>.py -v
  9. If compilation/testing fails, enter the troubleshooting loop (up to 3 times)
  10. When the precision test fails: if allclose does not pass, the output deviates, or it is all zeros or NaN,MUST read and follow the ascendc-operator-precision-debug skill process to troubleshoot the root cause (error analysis → code review → experiment isolation → instrumentation localization → fix verification)

Practical experience: source the environment variables before every shell command, otherwise the compilation tools or Python packages will not be found.

Post-generation checklist

op_host

  • namespace ascend_kernel，include torch_kernel_helper.h + tiling/platform/platform_ascendc.h + aclrtlaunch_<op_name>.h
  • Use platform APIs to obtain coreNum and ubSize (do not hardcode)
  • bufferCoefficient is consistent with the UB allocation table in the design document
  • All parameters of EXEC_KERNEL_CMD are lvalues
  • [Row processing] padding/de-padding is correct

op_kernel

  • include kernel_operator.h，BUFFER_NUM = 2
  • Init whole-core/tail-core offsets are correct, and the InitBuffer size is consistent with the UB allocation table
  • Process tail tile alignment (alignedTailLen)
  • AllocTensor/FreeTensor are paired, and EnQue/DeQue are paired
  • FP16/BF16 must be promoted to FP32 for computation
  • Back up the source data before ReduceSum

Framework adaptation

  • The ops.h declaration matches the op_host function signature
  • The register.cpp schema parameter types/default values are correct
  • csrc/CMakeLists.txt has added the host and kernel source files

Anti-pattern list

  • NEVER Letting FP16/BF16 participate directly in complex mathematical computations (Mul/Div/Exp/Tanh, etc.); they must first be Cast to FP32
  • NEVER Passing rvalues (temporary objects, literals, expression results) in EXEC_KERNEL_CMD
  • NEVER Use the bool parameter type in the kernel; replace it with int64_t
  • NEVER Use DataCopyPad for GM↔UB transfers instead of DataCopy
  • NEVER Reuse the source tensor directly after ReduceSum/ReduceMax (reduction may modify the source data)
  • NEVER Use std::min/max/abs/sqrt/exp and other standard library functions in the kernel
  • NEVER Pass repeatTime > 255 to the high-dimensional split API (uint8_t will silently truncate to 0)
  • NEVER Modify cmake/ or csrc/utils/ files under
  • NEVER Omit default values in the schema in register.cpp (e.g. int dim should be written as int dim=-1）
  • NEVER Hard-code the number of cores or UB size; it must be obtained through the platform API
  • NEVER Make the dst and tmpBuffer of ReduceMax/Sum point to the same memory block