with one click
ascendc-operator-code-gen
// 根据设计文档生成 AscendC 算子完整代码实现并完成框架适配。TRIGGER when: 设计文档已完成,需要生成 op_host/op_kernel 代码、注册到 PyTorch 框架、编译测试。关键词:代码生成、op_host、op_kernel、tiling、kernel、框架适配、算子注册。
// 根据设计文档生成 AscendC 算子完整代码实现并完成框架适配。TRIGGER when: 设计文档已完成,需要生成 op_host/op_kernel 代码、注册到 PyTorch 框架、编译测试。关键词:代码生成、op_host、op_kernel、tiling、kernel、框架适配、算子注册。
[HINT] Download the complete skill directory including SKILL.md and all related files
| name | ascendc-operator-code-gen |
| description | 根据设计文档生成 AscendC 算子完整代码实现并完成框架适配。TRIGGER when: 设计文档已完成,需要生成 op_host/op_kernel 代码、注册到 PyTorch 框架、编译测试。关键词:代码生成、op_host、op_kernel、tiling、kernel、框架适配、算子注册。 |
根据设计文档生成 op_host + op_kernel 代码,注册到 PyTorch 框架,调用 ascendc-operator-compile-debug skill 完成编译安装和精度测试。
前置条件: 设计文档 ascend-kernel/csrc/ops/<op_name>/design.md 已就绪,目录结构已创建。
读取设计文档 → 加载 reference → 选择模板 → 生成 op_host + op_kernel
→ 框架适配 (ops.h + register.cpp + csrc/CMakeLists.txt)
→ 调用 ascendc-operator-compile-debug skill (编译 + 安装 + 测试)
MANDATORY — READ BEFORE CODING: 读取 references/GUIDE.md,根据设计文档中的算子类型,加载对应的 reference 文件。绝对不要跳过此步骤。
从 ascend-kernel/csrc/ops/<op_name>/design.md 提取:
| 提取项 | 设计文档章节 | 用途 |
|---|---|---|
| 函数签名 + 支持的数据类型 | 原型设计 | op_host 函数原型、kernel 模板参数 |
| 算子类型 | Tiling 切分 → 步骤1 | 选择模板 (elementwise / row) |
| UB 分配表 | UB 空间分配 | 推导 bufferCoefficient、InitBuffer 大小 |
| 计算伪代码 | Kernel 实现 | Compute 函数逻辑 |
根据算子类型选择对应模板,复制到工程目录后修改。
| 算子类型 | op_host 模板 | op_kernel 模板 |
|---|---|---|
| Elementwise (ReLU, GELU, Add...) | templates/elementwise_op_host.cpp | templates/elementwise_op_kernel.cpp |
| 行处理 (LayerNorm, Softmax...) | templates/row_op_host.cpp | templates/row_op_kernel.cpp |
| Index-PerRow (IndexSelect, 共享1D索引) | templates/index_op_host.cpp | templates/index_op_kernel.cpp |
| Index-PerElem (Gather, Scatter, 逐行索引) | templates/index_op_per_elem_host.cpp | templates/index_op_per_elem_kernel.cpp |
| Sort (Sort, TopK...) | templates/sort_op_host.cpp | templates/sort_op_kernel.cpp |
| Pool (MaxPool, AvgPool...) | templates/pool_ndhwc_op_host.cpp | templates/pool_ndhwc_op_kernel.cpp |
MANDATORY: 读取对应模板文件的完整内容,复制到目标路径后修改。
ascend-kernel/csrc/ops/<op_name>/op_host/<op_name>.cpp 和 op_kernel/<op_name>.cpp(覆盖骨架占位文件)<op_name>, <OpName> 等)bufferCoefficient、EXEC_KERNEL_CMD 参数Kernel 模板使用 template <typename T> 泛型 + if constexpr 分支处理不同数据类型:
if constexpr (sizeof(T) == sizeof(float)) {
// float32 直接计算
} else {
// fp16/bf16: Cast 升精度 → fp32 计算 → Cast 降精度
}
AscendC 编译器支持 C++17 if constexpr,可安全用于模板分支。如需兼容旧编译器,也可用 sizeof(T) == 4 的普通 if(编译器会优化掉死分支)。
模板使用平台 API 动态获取硬件参数,不再硬编码:
#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();
从设计文档 UB 分配表推导——所有 Buffer 总大小 = bufferCoefficient * tileLength:
tileLength * N,则 bufferCoefficient = Nif (dtypeSize == 2) {...} else {...} 分支FP16 和 BF16 都必须升精度到 FP32 计算再转回。模板中已标注升精度代码的插入位置,取消注释并填充计算逻辑即可。
生产代码必须使用 DataCopyPad(不要使用 DataCopy 进行 GM↔UB 搬运):
// 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);
注意: 模板文件中可能使用 DataCopy 以简化示例,生产代码应替换为 DataCopyPad。
归约操作可能修改源 tensor,必须先备份:
AscendC::Adds(backup, src, 0.0f, len); // 备份
AscendC::ReduceSum<float, true>(result, backup, sharedTmp, dimLen);
csrc/ops.h在 namespace ascend_kernel 中添加函数声明:
<返回类型> <op_name>(<参数列表>);
csrc/register.cpp在 TORCH_LIBRARY_FRAGMENT(npu, m) 中添加 schema:
m.def("<op_name>(<schema_参数>) -> <返回类型>");
在 TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) 中添加实现:
m.impl("<op_name>", TORCH_FN(ascend_kernel::<op_name>));
Schema 语法速查:
| C++ 类型 | Schema 类型 | 示例 |
|---|---|---|
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 |
csrc/CMakeLists.txt在 FILE(GLOB OP_SRCS ...) 中添加 host 源文件:
${PROJECT_OP_SRC_BASE}/ops/<op_name>/op_host/<op_name>.cpp
在 ascendc_library(no_workspace_kernel STATIC ...) 中添加 kernel 源文件:
${PROJECT_OP_SRC_BASE}/ops/<op_name>/op_kernel/<op_name>.cpp
MANDATORY — 按 ascendc-operator-compile-debug skill 流程执行:
build.sh 有执行权限:chmod +x build.shbash build.shls output/ascend_kernel*.whlpip install output/ascend_kernel*.whl --force-reinstall --no-depstests/test_<op_name>.pypython tests/test_<op_name>.pypytest tests/test_<op_name>.py -vascendc-operator-precision-debug skill 流程排查根因(误差分析 → 代码审查 → 实验隔离 → 插桩定位 → 修复验证)实战经验:每次 shell 命令前都要 source 环境变量,否则会找不到编译工具或 Python 包。
ascend_kernel,include torch_kernel_helper.h + tiling/platform/platform_ascendc.h + aclrtlaunch_<op_name>.hkernel_operator.h,BUFFER_NUM = 2std::min/max/abs/sqrt/exp 等标准库函数cmake/ 或 csrc/utils/ 下的文件int dim 应写为 int dim=-1)