SuperAPI AI 工具集

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

智能体与模型

查看项目

更新于 10-15 16:58

技能介绍

本技能用于根据设计文档生成 AscendC 算子的完整代码,并完成框架适配。触发条件为:设计文档已完成,需要生成 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(编译 + 安装 + 测试)。

阶段 1:加载参考文档 读取 references/GUIDE.md,并根据设计文档中的算子类型加载对应的 reference 文件,该步骤标记为必须执行。

阶段 2:读取设计文档 从 ascend-kernel/csrc/ops/<op_name>/design.md 提取信息,包括:原型设计中的函数签名与支持的数据类型,用于 op_host 函数原型与 kernel 模板参数;Tiling 切分中的算子类型,用于选择 elementwise 或 row 模板;UB 空间分配中的 UB 分配表,用于推导 bufferCoefficient 与 InitBuffer 大小;Kernel 实现中的计算伪代码,用于确定 Compute 函数逻辑。

阶段 3:选择模板并生成代码 根据算子类型选择模板。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。需读取对应模板文件的完整内容,复制到目标路径后修改。

生成步骤:读取选中的 op_host 和 op_kernel 模板完整内容;将内容写入 ascend-kernel/csrc/ops/<op_name>/op_host/<op_name>.cpp 和 op_kernel/<op_name>.cpp(覆盖骨架占位文件);替换所有占位符(<op_name>、<OpName> 等);根据设计文档修改 op_host(函数签名、输入校验、bufferCoefficient、EXEC_KERNEL_CMD 参数)和 op_kernel(Init 的 GM 参数、InitBuffer 大小、Compute 计算逻辑)。

FP16/FP32 模板分支:Kernel 模板使用 template <typename T> 泛型加 if constexpr 分支处理不同数据类型。若 sizeof(T) == sizeof(float),float32 直接计算;否则 fp16/bf16 先 Cast 升精度到 fp32 计算,再 Cast 降精度。AscendC 编译器支持 C++17 if constexpr,可安全用于模板分支;如需兼容旧编译器,也可用 sizeof(T) == 4 的普通 if。

硬件参数获取:模板使用平台 API 动态获取硬件参数,不硬编码,通过 platform_ascendc::PlatformAscendCManager::GetInstance() 获取 coreNum 与 ubSize,需要 workspace 时可获取 sysWorkspaceSize。

bufferCoefficient 推导:从设计文档 UB 分配表推导,将所有 Buffer 的“总大小”列加和,结果形如 tileLength * N,则 bufferCoefficient = N。不同 dtype 系数不同时用 if (dtypeSize == 2) {...} else {...} 分支。

FP16/BF16 升精度:FP16 和 BF16 都必须升精度到 FP32 计算再转回,模板中已标注升精度代码的插入位置。

GM ↔ UB 搬运:生产代码必须使用 DataCopyPad,不要使用 DataCopy 进行 GM↔UB 搬运。模板文件中可能使用 DataCopy 以简化示例,生产代码应替换为 DataCopyPad。

ReduceSum/ReduceMax 注意:归约操作可能修改源 tensor,必须先备份。

阶段 4:框架适配 4.1 更新 csrc/ops.h:在 namespace ascend_kernel 中添加函数声明。 4.2 更新 csrc/register.cpp:在 TORCH_LIBRARY_FRAGMENT(npu, m) 中添加 schema,在 TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) 中添加实现。Schema 语法映射包括:const at::Tensor & 对应 Tensor;at::IntArrayRef 对应 int[];int64_t 对应 int;double 对应 float;bool 对应 bool;c10::optionalat::Tensor 对应 Tensor?;c10::optional<int64_t> 对应 int?。 4.3 更新 csrc/CMakeLists.txt:在 FILE(GLOB OP_SRCS ...) 中添加 host 源文件,在 ascendc_library(no_workspace_kernel STATIC ...) 中添加 kernel 源文件。

阶段 5:编译安装与测试 按 ascendc-operator-compile-debug skill 流程执行:加载环境(source CANN 环境变量,激活 conda 环境);确保 build.sh 有执行权限;编译工程 bash build.sh;验证编译成功 ls output/ascend_kernel*.whl;安装 whl 包 pip install output/ascend_kernel*.whl --force-reinstall --no-deps;生成测试文件 tests/test_<op_name>.py;运行功能测试 python tests/test_<op_name>.py;运行精度测试 pytest tests/test_<op_name>.py -v。编译/测试失败进入排错循环(最多 3 次)。精度测试失败时,若 allclose 不通过、输出偏差、全零或 NaN,须按 ascendc-operator-precision-debug skill 流程排查根因(误差分析 → 代码审查 → 实验隔离 → 插桩定位 → 修复验证)。实战经验:每次 shell 命令前都要 source 环境变量,否则会找不到编译工具或 Python 包。

生成后检查清单 op_host:namespace ascend_kernel,include torch_kernel_helper.h + tiling/platform/platform_ascendc.h + aclrtlaunch_<op_name>.h;使用平台 API 获取 coreNum 和 ubSize(不硬编码);bufferCoefficient 与设计文档 UB 分配表一致;EXEC_KERNEL_CMD 所有参数均为左值;行处理时 padding/去 padding 正确。 op_kernel:include kernel_operator.h,BUFFER_NUM = 2;Init 整核/尾核偏移正确,InitBuffer 大小与 UB 分配表一致;Process 尾 tile 对齐(alignedTailLen);AllocTensor/FreeTensor 配对,EnQue/DeQue 配对;FP16/BF16 必须升精度到 FP32 计算;ReduceSum 前备份源数据。 框架适配:ops.h 声明与 op_host 函数签名一致;register.cpp schema 参数类型/默认值正确;csrc/CMakeLists.txt 添加了 host 和 kernel 源文件。

反模式清单 不要让 FP16/BF16 直接参与复杂数学计算(Mul/Div/Exp/Tanh 等),必须先 Cast 到 FP32;不要在 EXEC_KERNEL_CMD 中传右值(临时对象、字面量、表达式结果);不要在 kernel 中使用 bool 参数类型,用 int64_t 替代;不要对 GM↔UB 搬运使用 DataCopy,必须用 DataCopyPad;不要在 ReduceSum/ReduceMax 后直接复用源 tensor(归约可能修改源数据);不要在 kernel 中使用 std::min/max/abs/sqrt/exp 等标准库函数;不要向高维切分 API 传入 repeatTime > 255(uint8_t 会静默截断为 0);不要修改 cmake/ 或 csrc/utils/ 下的文件;不要在 register.cpp 的 schema 中遗漏默认值(如 int dim 应写为 int dim=-1);不要硬编码核数或 UB 大小,必须通过平台 API 获取;不要让 ReduceMax/Sum 的 dst 与 tmpBuffer 指向同一块内存。