ATVC 实践通过 PyTorch 调用 ATVC 模板开发 Add 自定义 Vector 算子【免费下载链接】atvcATVCAscend C Templates for Vector Compute是为基于Ascend C开发的典型Vector算子封装的一系列模板头文件的集合可帮助用户快速开发典型Vector算子。项目地址: https://gitcode.com/cann/atvc本文以 ATVCAscend C Templates for Vector Compute仓库中的examples/ops_pytorch/add样例为主线完整讲解如何基于 ATVC 的 EleWise逐元素算子模板把一个自定义 Add 算子从 kernel 侧实现、PyTorch C 入口、Python 测试用例到编译脚本全链路打通并深入剖析CalcEleWiseTiling与EleWiseOpTemplate的底层运行机制。读完本文读者可以掌握“PyTorch 框架 ATVC 模板 核函数调用”这一 Vector 算子开发范式的完整落地方法。样例概述与目录结构本样例基于 AddCustom 算子工程其 ACLNN 形态可参考 examples/ops_aclnn/add介绍了基于 ATVC 的 PyTorch 工程搭建与调用方式。样例目录结构如下见 examples/ops_pytorch/add/README.mdadd/ ├── add_custom_impl.h // 通过PyTorch调用的方式调用Add算子 ├── pytorch_ascendc_extension.cpp // PyTorch调用入口 ├── run_op.py // PyTorch的测试用例 └── run.sh // 脚本编译需要的二进制文件并测试算子描述与规格Add 算子实现了两个数据相加、返回相加结果的功能对应的数学表达式为z x y算子规格如下项目名称shapedata typeformat算子类型OpTypeAdd---算子输入x8 × 2048int32_t、floatND算子输入y8 × 2048int32_t、floatND算子输出z8 × 2048int32_t、floatND核函数名AddCustom---也就是说本样例需要同时支持float与int两种数据类型的逐元素加法这也对应了源码中两套编译态类型参数Traits的设计。准备获取源码包与环境配置1. 获取源码包与基础环境编译运行此样例前请先参考 PyTorch 调用样例总览 中的“准备获取样例代码”一节完成 CANN 软件包安装、环境变量配置以及 ATVC 源码的获取对应 docs/01_quick_start.md 中的环境准备与源码下载章节。2. 安装 PyTorch 环境运行该样例要求torch、torch_npu版本支持2.7.1 及以上。按照原文档需要额外准备两类 CANN 软件包包名以实际 CANN 版本${cann_version}和机器架构${arch}为准cann-hccl_${cann_version}_linux-${arch}.runHCCL 软件包提供 x86_64 与 aarch64 两个版本Ascend-cann-A3-ops_${cann_version}_linux-${arch}.runA3 芯片算子包同样提供 x86_64 与 aarch64 两个版本。Kernel 侧实现基于 EleWise 模板的 AddCustom 核函数kernel 侧代码位于 add_custom_impl.h它是“ATVC 模板 自定义 Compute”范式的典型体现共三步定义编译态参数、定义计算逻辑、定义核函数入口。1. 定义编译态参数OpTraitsusing AddOpTraitsFloat ATVC::OpTraitsATVC::OpInputsfloat, float, ATVC::OpOutputsfloat; using AddOpTraitsInt ATVC::OpTraitsATVC::OpInputsint, int, ATVC::OpOutputsint;OpTraits通过编译期类型列表描述了算子原型两个输入、一个输出及其数据类型。Host 侧的 Tiling 计算与 Kernel 侧的模板实例化都会以它作为模板参数这正是 EleWise 数据流图中“提供 OpTraits 编译态参数描述算子原型”的落点。2. 定义计算逻辑自定义 Compute 仿函数// 传入编译态参数ATVC::OpTraits templatetypename Traits struct AddComputeFunc { // 函数说明 z x y templatetypename T // 重载operator提供给算子模板类调用 __aicore__ inline void operator()(AscendC::LocalTensorT x, AscendC::LocalTensorT y, AscendC::LocalTensorT z) { AscendC::Add(z, x, y, z.GetSize()); // 通过z.GetSize()获取单次计算的元素数量 } };用户只需实现一个重载了operator()的仿函数内部调用 Ascend C API此处为AscendC::Add完成实际数学运算。模板层会在每一“块”数据搬运完成后以LocalTensor形式传入 x、y、z用户通过z.GetSize()感知本次计算的元素数量从而天然兼容模板切分后的“尾块”。3. 定义核函数入口// 该函数为Add算子核函数入口 // x/y/z: Device上的gm地址分别指向Add算子输入1、输入2、输出1 // param: ATVC::EleWiseParam数据 template class Traits __global__ __aicore__ void AddCustom(GM_ADDR x, GM_ADDR y, GM_ADDR z, ATVC::EleWiseParam param) { KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); // 将AddComputeFunc仿函数作为模板参数传入实例化EleWiseOpTemplate模板类 auto op ATVC::Kernel::EleWiseOpTemplateAddComputeFuncTraits(); op.Run(x, y, z, param); }核函数本体只有三行有效代码声明 AIV 任务类型、实例化EleWiseOpTemplate并传入用户的AddComputeFunc、调用Run。数据搬运、多核切分、UB 缓冲管理全部由模板托管。PyTorch 调用入口通过 调度核函数PyTorch 入口位于 pytorch_ascendc_extension.cpp头文件引入部分是整个工程约定的关键——需要引入 PyTorch 扩展头、NPU 流接口以及保护核函数声明所在的{kernel_name}_impl.h#include torch/extension.h #include torch_npu/csrc/core/npu/NPUStream.h #include add_custom_impl.h核心实现函数按“取流 → 分配输出 → 计算 Tiling → 启动核函数”四步组织namespace ascendc_elewise_ops { at::Tensor op_add_custom(const at::Tensor x, const at::Tensor y) { // 运行资源申请通过c10_npu::getCurrentNPUStream()获取当前NPU上的流 auto stream c10_npu::getCurrentNPUStream().stream(false); // 分配Device侧输出内存 at::Tensor z at::empty_like(x); int32_t totalLength 1; for (int32_t size : x.sizes()) { totalLength * size; // 输入x展平后的总元素个数 } // 声明运行态参数param ATVC::EleWiseParam param; if (x.scalar_type() at::kFloat) { // Host侧调用Tiling API完成相关运行态参数的运算 (void)ATVC::Host::CalcEleWiseTilingAddOpTraitsFloat(totalLength, param); // 使用方式调用核函数完成指定的运算 AddCustomAddOpTraitsFloatparam.tilingData.blockNum, nullptr, stream( (uint8_t *)(x.storage().data()), (uint8_t *)(y.storage().data()), (uint8_t *)(z.storage().data()), param); } else if (x.scalar_type() at::kInt) { (void)ATVC::Host::CalcEleWiseTilingAddOpTraitsInt(totalLength, param); AddCustomAddOpTraitsIntparam.tilingData.blockNum, nullptr, stream( (uint8_t *)(x.storage().data()), (uint8_t *)(y.storage().data()), (uint8_t *)(z.storage().data()), param); } return z; } TORCH_LIBRARY(ascendc_ops, m) { m.def(add, ascendc_elewise_ops::op_add_custom); // 将算子与PyTorch绑定 } } // namespace ascendc_elewise_ops几点值得注意gridSize, nullptr, stream是 CCE 提供的类 CUDA 核函数启动语法第一个参数param.tilingData.blockNum由 Host 侧 Tiling API 计算得出直接决定在 NPU 上启动的核数输入张量 x、y 的 Device 内存由 Python 测试脚本通过x.npu()分配并拷入C 侧仅用at::empty_like(x)分配输出内存TORCH_LIBRARY(ascendc_ops, m)将 C 函数注册到名为ascendc_ops的算子命名空间Python 侧即可通过torch.ops.ascendc_ops.add(...)调用。Python 调用与测试用例测试用例位于 run_op.py基于torch_npu的TestCase框架实际包含 float 与 int 两组用例对应算子规格表中的两种数据类型import torch import torch_npu from torch_npu.testing.testcase import TestCase, run_tests torch.npu.config.allow_internal_format False # 关闭内部格式保证ND格式直传 torch.ops.load_library(./libascendc_pytorch.so) # 加载编译产物 class TestAscendCOps(TestCase): def test_add_custom_ops_float(self): # 分配Host侧输入内存并进行数据的初始化 length [8, 2048] x torch.rand(length, devicecpu, dtypetorch.float32) y torch.rand(length, devicecpu, dtypetorch.float32) # 将数据从Host拷贝到Device上并调用自定义算子 npuout torch.ops.ascendc_ops.add(x.npu(), y.npu()) cpuout torch.add(x, y) self.assertRtolEqual(npuout, cpuout) def test_add_custom_ops_int(self): length [8, 2048] x torch.randint(-10, 10, length, devicecpu, dtypetorch.int32) y torch.randint(-10, 10, length, devicecpu, dtypetorch.int32) npuout torch.ops.ascendc_ops.add(x.npu(), y.npu()) cpuout torch.add(x, y) self.assertRtolEqual(npuout, cpuout) if __name__ __main__: run_tests()用例流程为Host 侧构造随机输入 →x.npu()将数据搬运到 Device → 调用自定义算子 → 与 PyTorch 自带的torch.add结果做相对误差比对assertRtolEqual验证正确性。编译运行样例1. 编译脚本 run.sh 详解run.sh 自动完成“环境探测 → 编译 → 测试 → 清理”的完整流程关键逻辑如下# 动态探测torch、torch_npu、python的lib和include路径 torch_location$(python3 -c import torch; print(torch.__path__[0])) torch_npu_location$(python3 -c import torch_npu; print(torch_npu.__path__[0])) python_include$(python3 -c import sysconfig; print(sysconfig.get_path(include))) python_lib$(python3 -c import sysconfig; print(sysconfig.get_path(stdlib))) lib_path$(dirname $python_lib) export LD_LIBRARY_PATH${torch_npu_location}/lib/:$LD_LIBRARY_PATH export LD_LIBRARY_PATH${torch_location}/lib/:$LD_LIBRARY_PATH # ATVC头文件路径优先使用环境变量ATVC_PATH否则使用相对路径 ../../../include if [ -z $ATVC_PATH ]; then atvc_path$(realpath ../../../include) else atvc_path$ATVC_PATH fi脚本还会按ASCEND_INSTALL_PATH→ASCEND_HOME_PATH→~/Ascend/ascend-toolkit/latest→/usr/local/Ascend/ascend-toolkit/latest的顺序探测 CANN 安装路径然后使用bisheng编译器-x cce模式编译生成libascendc_pytorch.sobisheng -x cce pytorch_ascendc_extension.cpp \ -D_GLIBCXX_USE_CXX11_ABI1 \ -I${torch_location}/include \ -I${torch_location}/include/torch/csrc/api/include \ -I${python_include} \ -I${atvc_path} \ -I${torch_npu_location}/include \ -L${torch_location}/lib \ -L${torch_npu_location}/lib \ -L${python_lib} \ -L${lib_path} \ -L${_ASCEND_INSTALL_PATH}/lib64 \ -ltorch -ltorch_cpu -lc10 -ltorch_npu -lpython3 -ltorch_python \ -shared -cce-enable-plugin --cce-aicore-archdav-c220 -fPIC \ -ltiling_api -lplatform -lm -ldl \ -o libascendc_pytorch.so python3 run_op.py编译参数要点-cce-enable-plugin启用 CCE 插件使__global__ __aicore__核函数能够被同一编译单元编译进共享库--cce-aicore-archdav-c220指定 AICore 目标架构样例面向 A3 系列c220 架构-ltiling_api -lplatform链接 Tiling 与平台库支撑 Host 侧 Tiling API 与核函数启动-D_GLIBCXX_USE_CXX11_ABI1与 PyTorch 预编译库的 C ABI 保持一致避免符号不匹配。编译结束后脚本执行python3 run_op.py运行测试用例并将结果非零视为失败最后清理*.json与libascendc_pytorch.so中间产物。2. 执行验证按原文档给出的操作流程# 如果不导入默认使用./atvc/include路径 export ATVC_PATH${atvc}/include # 调用脚本编译生成PyTorch算子并运行测试用例 cd ./examples/ops_pytorch/add bash run.sh ... OK测试全部通过时run_tests()会输出OK即表示自定义 Add 算子在 NPU 上的结果与 CPU 参考实现一致。原理剖析ATVC EleWise 模板如何驱动数据搬运Host 侧CalcEleWiseTiling 的运行态参数ATVC::Host::CalcEleWiseTilingOpTraits定义在 include/elewise/host/elewise_host.h其签名带有默认超参调用方不传参时即使用框架内置的调优经验值struct EleWiseTilingHyperParam { uint32_t singleCoreBaseLine 512; // 单核数据量基线取值范围[256, 128*1024] float ubSizeLimitThreshold 0.95f; // UB内存占用上限决定basicBlock最大值 uint32_t nBufferNum 2; // 双缓冲数量取值范围[1, 2] uint32_t splitDataShape[3] {1024, 32*1024, 64*1024}; // 形状分段节点 uint32_t dataSplitFactor[4] {4, 4, 8, 6}; // 各分段节点的切分系数 uint32_t rsvLiveCnt 0; // 额外预留的UB空间节点数 };从源码看该函数的核心计算逻辑为blockNum totalCnt / singleCoreBaseLine小于基线时为 1并受vectorCoreNum上限约束——决定启动多少个 AIV 核按 UB 容量ubSize * ubSizeLimitThreshold与“输入输出字节数 × 缓冲数”折算出单块可容纳的元素上限ubufLimitCntGetEleWiseBasicCnt结合分段节点与切分系数得到每个 block 的基本数据块大小tiledCnt并强制按32 字节对齐basicCnt basicCnt / 32 * 32最小 32 个元素因 UB 需要 32B 对齐将blockNum、numPerBlock每核循环次数、tailBlockCnt、tailElemCnt尾块元素数、tiledCnt写入EleWiseParam传给核函数。这些运行态参数定义在 include/elewise/common/elewise_common.hstruct EleWiseTilingData { uint32_t tailBlockCnt; // 需要多执行一次循环的核数 uint32_t tailElemCnt; // 尾块元素个数 uint32_t numPerBlock; // 每个核计算的基本块数 uint32_t tiledCnt; // 基本块元素个数 uint32_t blockNum; // 执行的核数 }; struct EleWiseParam { EleWiseTilingData tilingData; // 影响数据处理的相关参数 uint32_t totalCnt 0; // 单个Tensor的元素个数 uint32_t nBufferNum 2; // 每个队列中Tensor的数量 };这正是 PyTorch 入口中param.tilingData.blockNum, nullptr, stream的 block 数来源。Kernel 侧EleWiseOpTemplate 的主循环include/elewise/kernel/elewise_op_template.h 中的EleWiseOpTemplate是逐元素算子的通用运行时Run(x, y, z, param)之后模板内部完成各核分工计算根据GetBlockIdx()与tailBlockCnt/numPerBlock每个核算出自己负责的元素区间curCoreStartCnt_和长度curCoreCnt_最后一个 block 额外承担tailElemCnt个尾块元素UB 缓冲初始化Init()按nBufferNum双缓冲初始化inQueueVECIN、outQueueVECOUT与 temp 缓冲区VECCALC实现搬运与计算重叠主循环Process()对每个核按repeat curCoreCnt_ / tiledCnt整块循环执行CopyIn → Compute → CopyOut若有余数则额外处理一次尾块tailCnt curCoreCnt_ % tiledCnt对齐搬运CopyIn/CopyOut内部按 32B 对齐切分主段用AscendC::DataCopy不足对齐的尾段用DataCopyPad补齐保证任意tiledCnt下搬运高效且不越界用户 Compute 调用Compute将各输入/输出LocalTensor按caclCnt_本块实际计算元素数解引用后交给用户仿函数即回调AddComputeFunc::operator()里的AscendC::Add。对于本样例的 float 版本每核 UB 占用为(2×4B 输入 4B 输出) × 2 缓冲 × tiledCnt 元素模板据此保证多核并行时 UB 不越界。小结与扩展路径本样例展示了 ATVC 面向 PyTorch 场景的标准开发路径kernel 侧定义OpTraits 自定义 Compute 仿函数 三行核函数入口add_custom_impl.hHost 侧CalcEleWiseTiling计算运行态参数以blockNum启动核函数pytorch_ascendc_extension.cpp验证侧torch_npuTestCase 对比 CPU 参考结果run_op.py工程侧bisheng -x cce一把编译出libascendc_pytorch.so并自动跑测run.sh。需要自行扩展时注意三点前提torch/torch_npu 2.7.1 及以上、编译目标架构与 NPU 硬件匹配样例为dav-c220、ATVC_PATH指向 ATVC 的 include 目录。更多逐元素/广播/归约算子的模板用法可参考 开发指南 与 PyTorch 调用样例总览其中 reduce_sum 展示了另一类算子模板的 PyTorch 集成方式。【免费下载链接】atvcATVCAscend C Templates for Vector Compute是为基于Ascend C开发的典型Vector算子封装的一系列模板头文件的集合可帮助用户快速开发典型Vector算子。项目地址: https://gitcode.com/cann/atvc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
ATVC 实践通过 PyTorch 调用 ATVC 模板开发 Add 自定义 Vector 算子【免费下载链接】atvcATVCAscend C Templates for Vector Compute是为基于Ascend C开发的典型Vector算子封装的一系列模板头文件的集合可帮助用户快速开发典型Vector算子。项目地址: https://gitcode.com/cann/atvc本文以 ATVCAscend C Templates for Vector Compute仓库中的examples/ops_pytorch/add样例为主线完整讲解如何基于 ATVC 的 EleWise逐元素算子模板把一个自定义 Add 算子从 kernel 侧实现、PyTorch C 入口、Python 测试用例到编译脚本全链路打通并深入剖析CalcEleWiseTiling与EleWiseOpTemplate的底层运行机制。读完本文读者可以掌握“PyTorch 框架 ATVC 模板 核函数调用”这一 Vector 算子开发范式的完整落地方法。样例概述与目录结构本样例基于 AddCustom 算子工程其 ACLNN 形态可参考 examples/ops_aclnn/add介绍了基于 ATVC 的 PyTorch 工程搭建与调用方式。样例目录结构如下见 examples/ops_pytorch/add/README.mdadd/ ├── add_custom_impl.h // 通过PyTorch调用的方式调用Add算子 ├── pytorch_ascendc_extension.cpp // PyTorch调用入口 ├── run_op.py // PyTorch的测试用例 └── run.sh // 脚本编译需要的二进制文件并测试算子描述与规格Add 算子实现了两个数据相加、返回相加结果的功能对应的数学表达式为z x y算子规格如下项目名称shapedata typeformat算子类型OpTypeAdd---算子输入x8 × 2048int32_t、floatND算子输入y8 × 2048int32_t、floatND算子输出z8 × 2048int32_t、floatND核函数名AddCustom---也就是说本样例需要同时支持float与int两种数据类型的逐元素加法这也对应了源码中两套编译态类型参数Traits的设计。准备获取源码包与环境配置1. 获取源码包与基础环境编译运行此样例前请先参考 PyTorch 调用样例总览 中的“准备获取样例代码”一节完成 CANN 软件包安装、环境变量配置以及 ATVC 源码的获取对应 docs/01_quick_start.md 中的环境准备与源码下载章节。2. 安装 PyTorch 环境运行该样例要求torch、torch_npu版本支持2.7.1 及以上。按照原文档需要额外准备两类 CANN 软件包包名以实际 CANN 版本${cann_version}和机器架构${arch}为准cann-hccl_${cann_version}_linux-${arch}.runHCCL 软件包提供 x86_64 与 aarch64 两个版本Ascend-cann-A3-ops_${cann_version}_linux-${arch}.runA3 芯片算子包同样提供 x86_64 与 aarch64 两个版本。Kernel 侧实现基于 EleWise 模板的 AddCustom 核函数kernel 侧代码位于 add_custom_impl.h它是“ATVC 模板 自定义 Compute”范式的典型体现共三步定义编译态参数、定义计算逻辑、定义核函数入口。1. 定义编译态参数OpTraitsusing AddOpTraitsFloat ATVC::OpTraitsATVC::OpInputsfloat, float, ATVC::OpOutputsfloat; using AddOpTraitsInt ATVC::OpTraitsATVC::OpInputsint, int, ATVC::OpOutputsint;OpTraits通过编译期类型列表描述了算子原型两个输入、一个输出及其数据类型。Host 侧的 Tiling 计算与 Kernel 侧的模板实例化都会以它作为模板参数这正是 EleWise 数据流图中“提供 OpTraits 编译态参数描述算子原型”的落点。2. 定义计算逻辑自定义 Compute 仿函数// 传入编译态参数ATVC::OpTraits templatetypename Traits struct AddComputeFunc { // 函数说明 z x y templatetypename T // 重载operator提供给算子模板类调用 __aicore__ inline void operator()(AscendC::LocalTensorT x, AscendC::LocalTensorT y, AscendC::LocalTensorT z) { AscendC::Add(z, x, y, z.GetSize()); // 通过z.GetSize()获取单次计算的元素数量 } };用户只需实现一个重载了operator()的仿函数内部调用 Ascend C API此处为AscendC::Add完成实际数学运算。模板层会在每一“块”数据搬运完成后以LocalTensor形式传入 x、y、z用户通过z.GetSize()感知本次计算的元素数量从而天然兼容模板切分后的“尾块”。3. 定义核函数入口// 该函数为Add算子核函数入口 // x/y/z: Device上的gm地址分别指向Add算子输入1、输入2、输出1 // param: ATVC::EleWiseParam数据 template class Traits __global__ __aicore__ void AddCustom(GM_ADDR x, GM_ADDR y, GM_ADDR z, ATVC::EleWiseParam param) { KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); // 将AddComputeFunc仿函数作为模板参数传入实例化EleWiseOpTemplate模板类 auto op ATVC::Kernel::EleWiseOpTemplateAddComputeFuncTraits(); op.Run(x, y, z, param); }核函数本体只有三行有效代码声明 AIV 任务类型、实例化EleWiseOpTemplate并传入用户的AddComputeFunc、调用Run。数据搬运、多核切分、UB 缓冲管理全部由模板托管。PyTorch 调用入口通过 调度核函数PyTorch 入口位于 pytorch_ascendc_extension.cpp头文件引入部分是整个工程约定的关键——需要引入 PyTorch 扩展头、NPU 流接口以及保护核函数声明所在的{kernel_name}_impl.h#include torch/extension.h #include torch_npu/csrc/core/npu/NPUStream.h #include add_custom_impl.h核心实现函数按“取流 → 分配输出 → 计算 Tiling → 启动核函数”四步组织namespace ascendc_elewise_ops { at::Tensor op_add_custom(const at::Tensor x, const at::Tensor y) { // 运行资源申请通过c10_npu::getCurrentNPUStream()获取当前NPU上的流 auto stream c10_npu::getCurrentNPUStream().stream(false); // 分配Device侧输出内存 at::Tensor z at::empty_like(x); int32_t totalLength 1; for (int32_t size : x.sizes()) { totalLength * size; // 输入x展平后的总元素个数 } // 声明运行态参数param ATVC::EleWiseParam param; if (x.scalar_type() at::kFloat) { // Host侧调用Tiling API完成相关运行态参数的运算 (void)ATVC::Host::CalcEleWiseTilingAddOpTraitsFloat(totalLength, param); // 使用方式调用核函数完成指定的运算 AddCustomAddOpTraitsFloatparam.tilingData.blockNum, nullptr, stream( (uint8_t *)(x.storage().data()), (uint8_t *)(y.storage().data()), (uint8_t *)(z.storage().data()), param); } else if (x.scalar_type() at::kInt) { (void)ATVC::Host::CalcEleWiseTilingAddOpTraitsInt(totalLength, param); AddCustomAddOpTraitsIntparam.tilingData.blockNum, nullptr, stream( (uint8_t *)(x.storage().data()), (uint8_t *)(y.storage().data()), (uint8_t *)(z.storage().data()), param); } return z; } TORCH_LIBRARY(ascendc_ops, m) { m.def(add, ascendc_elewise_ops::op_add_custom); // 将算子与PyTorch绑定 } } // namespace ascendc_elewise_ops几点值得注意gridSize, nullptr, stream是 CCE 提供的类 CUDA 核函数启动语法第一个参数param.tilingData.blockNum由 Host 侧 Tiling API 计算得出直接决定在 NPU 上启动的核数输入张量 x、y 的 Device 内存由 Python 测试脚本通过x.npu()分配并拷入C 侧仅用at::empty_like(x)分配输出内存TORCH_LIBRARY(ascendc_ops, m)将 C 函数注册到名为ascendc_ops的算子命名空间Python 侧即可通过torch.ops.ascendc_ops.add(...)调用。Python 调用与测试用例测试用例位于 run_op.py基于torch_npu的TestCase框架实际包含 float 与 int 两组用例对应算子规格表中的两种数据类型import torch import torch_npu from torch_npu.testing.testcase import TestCase, run_tests torch.npu.config.allow_internal_format False # 关闭内部格式保证ND格式直传 torch.ops.load_library(./libascendc_pytorch.so) # 加载编译产物 class TestAscendCOps(TestCase): def test_add_custom_ops_float(self): # 分配Host侧输入内存并进行数据的初始化 length [8, 2048] x torch.rand(length, devicecpu, dtypetorch.float32) y torch.rand(length, devicecpu, dtypetorch.float32) # 将数据从Host拷贝到Device上并调用自定义算子 npuout torch.ops.ascendc_ops.add(x.npu(), y.npu()) cpuout torch.add(x, y) self.assertRtolEqual(npuout, cpuout) def test_add_custom_ops_int(self): length [8, 2048] x torch.randint(-10, 10, length, devicecpu, dtypetorch.int32) y torch.randint(-10, 10, length, devicecpu, dtypetorch.int32) npuout torch.ops.ascendc_ops.add(x.npu(), y.npu()) cpuout torch.add(x, y) self.assertRtolEqual(npuout, cpuout) if __name__ __main__: run_tests()用例流程为Host 侧构造随机输入 →x.npu()将数据搬运到 Device → 调用自定义算子 → 与 PyTorch 自带的torch.add结果做相对误差比对assertRtolEqual验证正确性。编译运行样例1. 编译脚本 run.sh 详解run.sh 自动完成“环境探测 → 编译 → 测试 → 清理”的完整流程关键逻辑如下# 动态探测torch、torch_npu、python的lib和include路径 torch_location$(python3 -c import torch; print(torch.__path__[0])) torch_npu_location$(python3 -c import torch_npu; print(torch_npu.__path__[0])) python_include$(python3 -c import sysconfig; print(sysconfig.get_path(include))) python_lib$(python3 -c import sysconfig; print(sysconfig.get_path(stdlib))) lib_path$(dirname $python_lib) export LD_LIBRARY_PATH${torch_npu_location}/lib/:$LD_LIBRARY_PATH export LD_LIBRARY_PATH${torch_location}/lib/:$LD_LIBRARY_PATH # ATVC头文件路径优先使用环境变量ATVC_PATH否则使用相对路径 ../../../include if [ -z $ATVC_PATH ]; then atvc_path$(realpath ../../../include) else atvc_path$ATVC_PATH fi脚本还会按ASCEND_INSTALL_PATH→ASCEND_HOME_PATH→~/Ascend/ascend-toolkit/latest→/usr/local/Ascend/ascend-toolkit/latest的顺序探测 CANN 安装路径然后使用bisheng编译器-x cce模式编译生成libascendc_pytorch.sobisheng -x cce pytorch_ascendc_extension.cpp \ -D_GLIBCXX_USE_CXX11_ABI1 \ -I${torch_location}/include \ -I${torch_location}/include/torch/csrc/api/include \ -I${python_include} \ -I${atvc_path} \ -I${torch_npu_location}/include \ -L${torch_location}/lib \ -L${torch_npu_location}/lib \ -L${python_lib} \ -L${lib_path} \ -L${_ASCEND_INSTALL_PATH}/lib64 \ -ltorch -ltorch_cpu -lc10 -ltorch_npu -lpython3 -ltorch_python \ -shared -cce-enable-plugin --cce-aicore-archdav-c220 -fPIC \ -ltiling_api -lplatform -lm -ldl \ -o libascendc_pytorch.so python3 run_op.py编译参数要点-cce-enable-plugin启用 CCE 插件使__global__ __aicore__核函数能够被同一编译单元编译进共享库--cce-aicore-archdav-c220指定 AICore 目标架构样例面向 A3 系列c220 架构-ltiling_api -lplatform链接 Tiling 与平台库支撑 Host 侧 Tiling API 与核函数启动-D_GLIBCXX_USE_CXX11_ABI1与 PyTorch 预编译库的 C ABI 保持一致避免符号不匹配。编译结束后脚本执行python3 run_op.py运行测试用例并将结果非零视为失败最后清理*.json与libascendc_pytorch.so中间产物。2. 执行验证按原文档给出的操作流程# 如果不导入默认使用./atvc/include路径 export ATVC_PATH${atvc}/include # 调用脚本编译生成PyTorch算子并运行测试用例 cd ./examples/ops_pytorch/add bash run.sh ... OK测试全部通过时run_tests()会输出OK即表示自定义 Add 算子在 NPU 上的结果与 CPU 参考实现一致。原理剖析ATVC EleWise 模板如何驱动数据搬运Host 侧CalcEleWiseTiling 的运行态参数ATVC::Host::CalcEleWiseTilingOpTraits定义在 include/elewise/host/elewise_host.h其签名带有默认超参调用方不传参时即使用框架内置的调优经验值struct EleWiseTilingHyperParam { uint32_t singleCoreBaseLine 512; // 单核数据量基线取值范围[256, 128*1024] float ubSizeLimitThreshold 0.95f; // UB内存占用上限决定basicBlock最大值 uint32_t nBufferNum 2; // 双缓冲数量取值范围[1, 2] uint32_t splitDataShape[3] {1024, 32*1024, 64*1024}; // 形状分段节点 uint32_t dataSplitFactor[4] {4, 4, 8, 6}; // 各分段节点的切分系数 uint32_t rsvLiveCnt 0; // 额外预留的UB空间节点数 };从源码看该函数的核心计算逻辑为blockNum totalCnt / singleCoreBaseLine小于基线时为 1并受vectorCoreNum上限约束——决定启动多少个 AIV 核按 UB 容量ubSize * ubSizeLimitThreshold与“输入输出字节数 × 缓冲数”折算出单块可容纳的元素上限ubufLimitCntGetEleWiseBasicCnt结合分段节点与切分系数得到每个 block 的基本数据块大小tiledCnt并强制按32 字节对齐basicCnt basicCnt / 32 * 32最小 32 个元素因 UB 需要 32B 对齐将blockNum、numPerBlock每核循环次数、tailBlockCnt、tailElemCnt尾块元素数、tiledCnt写入EleWiseParam传给核函数。这些运行态参数定义在 include/elewise/common/elewise_common.hstruct EleWiseTilingData { uint32_t tailBlockCnt; // 需要多执行一次循环的核数 uint32_t tailElemCnt; // 尾块元素个数 uint32_t numPerBlock; // 每个核计算的基本块数 uint32_t tiledCnt; // 基本块元素个数 uint32_t blockNum; // 执行的核数 }; struct EleWiseParam { EleWiseTilingData tilingData; // 影响数据处理的相关参数 uint32_t totalCnt 0; // 单个Tensor的元素个数 uint32_t nBufferNum 2; // 每个队列中Tensor的数量 };这正是 PyTorch 入口中param.tilingData.blockNum, nullptr, stream的 block 数来源。Kernel 侧EleWiseOpTemplate 的主循环include/elewise/kernel/elewise_op_template.h 中的EleWiseOpTemplate是逐元素算子的通用运行时Run(x, y, z, param)之后模板内部完成各核分工计算根据GetBlockIdx()与tailBlockCnt/numPerBlock每个核算出自己负责的元素区间curCoreStartCnt_和长度curCoreCnt_最后一个 block 额外承担tailElemCnt个尾块元素UB 缓冲初始化Init()按nBufferNum双缓冲初始化inQueueVECIN、outQueueVECOUT与 temp 缓冲区VECCALC实现搬运与计算重叠主循环Process()对每个核按repeat curCoreCnt_ / tiledCnt整块循环执行CopyIn → Compute → CopyOut若有余数则额外处理一次尾块tailCnt curCoreCnt_ % tiledCnt对齐搬运CopyIn/CopyOut内部按 32B 对齐切分主段用AscendC::DataCopy不足对齐的尾段用DataCopyPad补齐保证任意tiledCnt下搬运高效且不越界用户 Compute 调用Compute将各输入/输出LocalTensor按caclCnt_本块实际计算元素数解引用后交给用户仿函数即回调AddComputeFunc::operator()里的AscendC::Add。对于本样例的 float 版本每核 UB 占用为(2×4B 输入 4B 输出) × 2 缓冲 × tiledCnt 元素模板据此保证多核并行时 UB 不越界。小结与扩展路径本样例展示了 ATVC 面向 PyTorch 场景的标准开发路径kernel 侧定义OpTraits 自定义 Compute 仿函数 三行核函数入口add_custom_impl.hHost 侧CalcEleWiseTiling计算运行态参数以blockNum启动核函数pytorch_ascendc_extension.cpp验证侧torch_npuTestCase 对比 CPU 参考结果run_op.py工程侧bisheng -x cce一把编译出libascendc_pytorch.so并自动跑测run.sh。需要自行扩展时注意三点前提torch/torch_npu 2.7.1 及以上、编译目标架构与 NPU 硬件匹配样例为dav-c220、ATVC_PATH指向 ATVC 的 include 目录。更多逐元素/广播/归约算子的模板用法可参考 开发指南 与 PyTorch 调用样例总览其中 reduce_sum 展示了另一类算子模板的 PyTorch 集成方式。【免费下载链接】atvcATVCAscend C Templates for Vector Compute是为基于Ascend C开发的典型Vector算子封装的一系列模板头文件的集合可帮助用户快速开发典型Vector算子。项目地址: https://gitcode.com/cann/atvc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
