ATVC 实践：通过 PyTorch 调用 ATVC 模板开发 Add 自定义 Vector 算子

ATVC 实践：通过 PyTorch 调用 ATVC 模板开发 Add 自定义 Vector 算子

【免费下载链接】atvcATVC（Ascend C Templates for Vector Compute），是为基于Ascend C开发的典型Vector算子封装的一系列模板头文件的集合，可帮助用户快速开发典型Vector算子。项目地址: https://gitcode.com/cann/atvc

本文以 ATVC（Ascend 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.md）：

add/ ├── add_custom_impl.h // 通过PyTorch调用的方式调用Add算子 ├── pytorch_ascendc_extension.cpp // PyTorch调用入口 ├── run_op.py // PyTorch的测试用例 └── run.sh // 脚本，编译需要的二进制文件，并测试

算子描述与规格

Add 算子实现了两个数据相加、返回相加结果的功能，对应的数学表达式为：

z = x + y

算子规格如下：

也就是说，本样例需要同时支持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}.run：HCCL 软件包（提供 x86_64 与 aarch64 两个版本）；
• Ascend-cann-A3-ops_${cann_version}_linux-${arch}.run：A3 芯片算子包（同样提供 x86_64 与 aarch64 两个版本）。

Kernel 侧实现：基于 EleWise 模板的 AddCustom 核函数

kernel 侧代码位于 add_custom_impl.h，它是“ATVC 模板 + 自定义 Compute”范式的典型体现，共三步：定义编译态参数、定义计算逻辑、定义核函数入口。

1. 定义编译态参数（OpTraits）

using AddOpTraitsFloat = ATVC::OpTraits<ATVC::OpInputs<float, float>, ATVC::OpOutputs<float>>; using AddOpTraitsInt = ATVC::OpTraits<ATVC::OpInputs<int, int>, ATVC::OpOutputs<int>>;

OpTraits通过编译期类型列表描述了算子原型：两个输入、一个输出及其数据类型。Host 侧的 Tiling 计算与 Kernel 侧的模板实例化都会以它作为模板参数，这正是 EleWise 数据流图中“提供 OpTraits 编译态参数，描述算子原型”的落点。

2. 定义计算逻辑（自定义 Compute 仿函数）

// 传入编译态参数ATVC::OpTraits template<typename Traits> struct AddComputeFunc { // 函数说明： z = x + y template<typename T> // 重载operator，提供给算子模板类调用 __aicore__ inline void operator()(AscendC::LocalTensor<T> x, AscendC::LocalTensor<T> y, AscendC::LocalTensor<T> 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::EleWiseOpTemplate<AddComputeFunc<Traits>>(); 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::CalcEleWiseTiling<AddOpTraitsFloat>(totalLength, param); // 使用<<<>>方式调用核函数完成指定的运算 AddCustom<AddOpTraitsFloat><<<param.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::CalcEleWiseTiling<AddOpTraitsInt>(totalLength, param); AddCustom<AddOpTraitsInt><<<param.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, device='cpu', dtype=torch.float32) y = torch.rand(length, device='cpu', dtype=torch.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, device='cpu', dtype=torch.int32) y = torch.randint(-10, 10, length, device='cpu', dtype=torch.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.so：

bisheng -x cce pytorch_ascendc_extension.cpp \ -D_GLIBCXX_USE_CXX11_ABI=1 \ -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-arch=dav-c220 -fPIC \ -ltiling_api -lplatform -lm -ldl \ -o libascendc_pytorch.so python3 run_op.py

编译参数要点：

• -cce-enable-plugin：启用 CCE 插件，使__global__ __aicore__核函数能够被同一编译单元编译进共享库；
• --cce-aicore-arch=dav-c220：指定 AICore 目标架构（样例面向 A3 系列，c220 架构）；
• -ltiling_api -lplatform：链接 Tiling 与平台库，支撑 Host 侧 Tiling API 与核函数启动；
• -D_GLIBCXX_USE_CXX11_ABI=1：与 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::CalcEleWiseTiling<OpTraits>定义在 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空间节点数 };

从源码看，该函数的核心计算逻辑为：

1. blockNum = totalCnt / singleCoreBaseLine（小于基线时为 1），并受vectorCoreNum上限约束——决定启动多少个 AIV 核；
2. 按 UB 容量（ubSize * ubSizeLimitThreshold）与“输入+输出字节数 × 缓冲数”折算出单块可容纳的元素上限ubufLimitCnt；
3. GetEleWiseBasicCnt结合分段节点与切分系数得到每个 block 的基本数据块大小tiledCnt，并强制按32 字节对齐（basicCnt = basicCnt / 32 * 32，最小 32 个元素，因 UB 需要 32B 对齐）；
4. 将blockNum、numPerBlock（每核循环次数）、tailBlockCnt、tailElemCnt（尾块元素数）、tiledCnt写入EleWiseParam传给核函数。

这些运行态参数定义在 include/elewise/common/elewise_common.h：

struct 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双缓冲初始化inQueue（VECIN）、outQueue（VECOUT）与 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 场景的标准开发路径：

1. kernel 侧：定义OpTraits+ 自定义 Compute 仿函数 + 三行核函数入口（add_custom_impl.h）；
2. Host 侧：CalcEleWiseTiling计算运行态参数，<<<>>>以blockNum启动核函数（pytorch_ascendc_extension.cpp）；
3. 验证侧：torch_npuTestCase 对比 CPU 参考结果（run_op.py）；
4. 工程侧：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 集成方式）。

【免费下载链接】atvcATVC（Ascend C Templates for Vector Compute），是为基于Ascend C开发的典型Vector算子封装的一系列模板头文件的集合，可帮助用户快速开发典型Vector算子。项目地址: https://gitcode.com/cann/atvc