asc-devkit Ascend C 算子接入 PyTorch 图模式：TorchAir 转换 GE 节点样例实战-平芜编程栈

人工智能
深度学习
算子库
CANN
Ascend

【免费下载链接】asc-devkit

本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言，原生支持C和C++标准规范，主要由类库和语言扩展层构成，提供多层级API，满足多维场景算子开发诉求。

项目地址：https://gitcode.com/cann/asc-devkit

本文基于 asc-devkit 仓库中的ge_torchair样例，讲解 Ascend C 自定义算子如何经 TorchAir 接入 PyTorch 图模式（torch.compile）执行链路。读完后，你将掌握从 Ascend C kernel 编写、GE 算子原型注册、PyTorch schema/Meta 注册，到 TorchAir converter 转换与图模式验证测试的完整接入方法，并能复用该样例接入自己的自定义算子。

样例定位与核心链路

该样例以AddCustom算子为例，演示模型中调用torch.ops.ascendc_ops.ascendc_add后，torch.compile捕获该调用、TorchAir 将其转换为 GEAddCustom节点、最终由 GE 执行入口启动 Ascend C kernel 的全过程。其中 TorchAir 承担 PyTorch 图到 GE 图的转换工作，是 PyTorch 自定义算子进入 GE 图执行链路的关键桥梁。

核心链路按交付件职责拆解为 5 步：

在 add_custom_kernel.asc 中实现 Ascend C Add kernel；
在 add_custom_ge.asc 中实现 GEAddCustom交付件，包括算子原型、shape/dtype 推导和Execute执行入口；
add_custom_test.py 通过torch.ops.load_library加载编译生成的libascendc_ops.so，使 PyTorch 侧接口和 GEAddCustom交付件注册生效；
在 add_custom_pytorch.asc 和add_custom_test.py中补充 PyTorch 侧 schema/Meta 注册和 TorchAir converter，使模型中的ascendc_add节点可以映射到 GEAddCustom节点；
通过add_custom_test.py使用 TorchAir 后端执行torch.compile图模式测试，并校验 GE 图执行结果。

该样例隶属于 PyTorch 框架接入目录，该目录下还有pybind、torch_library（Eager 模式接入）等同类样例，ge_torchair是其中的图模式接入方案。

支持的产品与 CANN 版本

样例支持的产品及 CANN 软件版本如下：

产品	CANN 软件版本
Ascend 950PR & 950DT 系列产品	>= CANN 9.1.0
Atlas A3 系列产品	>= CANN 9.1.0
Atlas A2 系列产品	>= CANN 9.1.0

目录结构

├── ge_torchair │ ├── CMakeLists.txt // 编译工程文件 │ ├── add_custom_kernel.asc // Ascend C kernel实现 │ ├── add_custom_pytorch.asc // PyTorch侧schema和Meta注册 │ ├── add_custom_ge.asc // GE AddCustom原型注册和执行实现 │ ├── add_custom_test.py // TorchAir图模式验证脚本 │ └── README.md // 样例说明文档

第一步：Ascend C 计算实现

add_custom_kernel.asc 实现 Add 计算对应的 Ascend C kernel，计算公式为z = x + y：

x：输入，形状为[8, 2048]，数据类型为 float，格式为 ND；
y：输入，形状为[8, 2048]，数据类型为 float，格式为 ND；
z：输出，形状为[8, 2048]，数据类型为 float，格式为 ND。

kernel 的完整实现如下，可以看到它采用典型的“GM → UB 搬运、UB 计算、结果写回 GM”结构：

__global__ __vector__ void add_custom(__gm__ float* x, __gm__ float* y, __gm__ float* z) { AscendC::InitSocState(); AscendC::GlobalTensor<float> xGm; AscendC::GlobalTensor<float> yGm; AscendC::GlobalTensor<float> zGm; xGm.SetGlobalBuffer(x + block_idx * BLOCK_LENGTH, BLOCK_LENGTH); yGm.SetGlobalBuffer(y + block_idx * BLOCK_LENGTH, BLOCK_LENGTH); zGm.SetGlobalBuffer(z + block_idx * BLOCK_LENGTH, BLOCK_LENGTH); AscendC::LocalMemAllocator<AscendC::Hardware::UB> ubAllocator; AscendC::LocalTensor<float> xLocal = ubAllocator.Alloc<float>(BLOCK_LENGTH); AscendC::LocalTensor<float> yLocal = ubAllocator.Alloc<float>(BLOCK_LENGTH); AscendC::LocalTensor<float> zLocal = ubAllocator.Alloc<float>(BLOCK_LENGTH); AscendC::DataCopy(xLocal, xGm, BLOCK_LENGTH); AscendC::DataCopy(yLocal, yGm, BLOCK_LENGTH); AscendC::PipeBarrier<PIPE_ALL>(); AscendC::Add(zLocal, xLocal, yLocal, BLOCK_LENGTH); AscendC::PipeBarrier<PIPE_ALL>(); AscendC::DataCopy(zGm, zLocal, BLOCK_LENGTH); AscendC::PipeBarrier<PIPE_ALL>(); }

从源码结构看，BLOCK_LENGTH在 kernel 中定义为2048U，即每个 block 处理 2048 个 float 元素；kernel 通过block_idx * BLOCK_LENGTH对三个 GlobalTensor 做偏移，由 8 个 block 恰好覆盖[8, 2048]的全部数据，各 block 间无交叉。计算前后均使用AscendC::PipeBarrier<PIPE_ALL>()保障搬运与计算的时序正确性。

第二步：GE 交付件实现

add_custom_ge.asc 提供 GE 侧需要的AddCustom原型、shape/dtype 推导和执行入口。GE 运行到AddCustom节点时进入执行入口，并启动add_custom_kernel.asc中的 Ascend C kernel。

该类继承EagerExecuteOp与ShapeInferOp，分别对应运行时执行和图编译阶段的推导两个职责：

class AddCustom : public EagerExecuteOp, public ShapeInferOp { public: // GE运行到AddCustom节点时调用Execute，完成输入校验、输出申请和Ascend C核函数启动 graphStatus Execute(gert::EagerOpExecutionContext* ctx) override { // ... // 申请输出Tensor，shape、format和dtype与输入x保持一致 gert::Tensor* z = ctx->MallocOutputTensor(OUTPUT_Z_INDEX, x->GetShape(), x->GetFormat(), x->GetDataType()); // ... // 使用GE提供的stream启动Ascend C核函数执行加法计算 ascendc_ops::add_custom<<<ascendc_ops::BLOCK_NUM, 0, ctx->GetStream()>>>( static_cast<float*>(const_cast<void*>(x->GetAddr())), static_cast<float*>(const_cast<void*>(y->GetAddr())), static_cast<float*>(z->GetAddr())); return GRAPH_SUCCESS; } // 推导输出shape，本样例为逐元素加法，输出shape与输入x一致 graphStatus InferShape(gert::InferShapeContext* ctx) override { // ... *zShape = *xShape; return GRAPH_SUCCESS; } // 推导输出数据类型，本样例输出dtype与输入x一致 graphStatus InferDataType(gert::InferDataTypeContext* ctx) override { // ... return ctx->SetOutputDataType(OUTPUT_Z_INDEX, ctx->GetInputDataType(INPUT_X_INDEX)); } };

结合源码可以看到几个实现细节：

输入校验：Execute入口先通过ctx->GetInputTensor获取输入，随后校验两个输入均非空、dtype 均为DT_FLOAT、shape 均为[8, 2048]（校验逻辑封装在IsSupportedShape中），任一不满足即打印错误并返回GRAPH_FAILED。这说明该样例的 kernel 规格是写死的，接入实际算子时需要扩展为动态规格；
输出申请：通过ctx->MallocOutputTensor按输入 x 的 shape、format、dtype 申请输出，保证输出与输入规格一致；
kernel 启动：使用<<<BLOCK_NUM, 0, ctx->GetStream()>>>三尖括号语法在 GE 提供的 stream 上启动 kernel，BLOCK_NUM = 8与 kernel 的block_idx切分逻辑对应；
原型注册：文件末尾通过REG_OP(AddCustom)声明x/y/z三个端口（均为DT_FLOAT），使 TorchAir 转换出的 GE 节点能被 GE 识别，再通过REG_AUTO_MAPPING_OP(AddCustom)将AddCustom类注册为该原型的执行实现——GE 运行时遇到AddCustom节点会自动调用AddCustom::Execute()。

REG_OP(AddCustom) .INPUT(x, TensorType({DT_FLOAT})) .INPUT(y, TensorType({DT_FLOAT})) .OUTPUT(z, TensorType({DT_FLOAT})) .OP_END_FACTORY_REG(AddCustom); REG_AUTO_MAPPING_OP(AddCustom);

第三步：PyTorch 侧 schema 与 Meta 注册

add_custom_pytorch.asc 提供 PyTorch 侧可见的ascendc_ops::ascendc_add接口。图捕获阶段主要需要两类信息：接口签名和 Meta 实现。

接口签名用于让 Python 侧能够通过torch.ops.ascendc_ops.ascendc_add(x, y)找到该接口：

TORCH_LIBRARY_FRAGMENT(ascendc_ops, m) { m.def("ascendc_add(Tensor x, Tensor y) -> Tensor"); }

Meta 实现：torch.compile捕获图时不会执行真实 NPU 计算，但需要推导输出 Tensor 信息，因此还需要注册 Meta 实现：

TORCH_LIBRARY_IMPL(ascendc_ops, Meta, m) { m.impl("ascendc_add", TORCH_FN(AscendcAddMeta)); }

从源码结构看，AscendcAddMeta在推导前会做与 GE 侧一致的规格校验（x/y形状相同、维度为[8, 2048]、dtype 为Float），校验通过后返回at::empty_like(x)作为输出描述：

at::Tensor AscendcAddMeta(const at::Tensor& x, const at::Tensor& y) { TORCH_CHECK(x.sizes() == y.sizes(), "x and y must have the same shape."); TORCH_CHECK( x.dim() == 2 && x.size(0) == ROW_NUM && x.size(1) == COL_NUM, "ascendc_add only supports shape [8, 2048] in this sample."); TORCH_CHECK( x.scalar_type() == at::ScalarType::Float && y.scalar_type() == at::ScalarType::Float, "ascendc_add only supports float tensors in this sample."); return at::empty_like(x); }

这意味着 PyTorch 侧（图捕获阶段）和 GE 侧（执行阶段）存在双重校验，两侧约束需保持一致，否则会出现图能捕获但执行失败的问题。

第四步：TorchAir converter 与图模式执行

add_custom_test.py 同时承担 TorchAir converter 注册与图模式验证两个角色。converter 将torch.ops.ascendc_ops.ascendc_add转换为 GEAddCustom节点：

@torchair.register_fx_node_ge_converter(torch.ops.ascendc_ops.ascendc_add.default) def convert_ascendc_add(x, y, z=None, meta_outputs=None): return torchair.ge.custom_op( "AddCustom", inputs={"x": x, "y": y}, outputs=["z"], )

converter 的签名中z与meta_outputs为 TorchAir 框架约定的占位参数；返回值torchair.ge.custom_op以算子名"AddCustom"和输入/输出端口字典的形式构造 GE 节点，端口名必须与REG_OP(AddCustom)中声明的x、y、z一致。

converter 注册后，加载与执行链路如下：

# 加载编译生成的动态库，使PyTorch侧接口和GE AddCustom交付件注册生效 torch.ops.load_library(LIB_PATH) class AddCustomModel(torch.nn.Module): def forward(self, x, y): return torch.ops.ascendc_ops.ascendc_add(x, y) opt_model = torch.compile(model, fullgraph=True, backend=npu_backend, dynamic=False)

其中LIB_PATH支持通过环境变量ASCENDC_OPS_LIB指定，默认为当前目录下的libascendc_ops.so。torch.compile的fullgraph=True确保 forward 中的ascendc_add被整图捕获而非拆分为动态执行；npu_backend由torchair.get_npu_backend(compiler_config=torchair.CompilerConfig())构造。编译后的模型在 NPU 侧执行，GE 运行到AddCustom节点时进入add_custom_ge.asc中的执行入口，并启动add_custom_kernel.asc中的 Ascend C kernel。

编译算子库

本样例的构建工程为 CMakeLists.txt，在样例根目录下执行：

# 创建并进入build目录 mkdir -p build && cd build; # 编译单算子，生成libascendc_ops.so cmake -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..;make -j; # 加载libascendc_ops.so，并执行TorchAir图模式测试 python3 ../add_custom_test.py

编译选项说明：

选项	可选值	说明
`CMAKE_ASC_ARCHITECTURES`	`dav-2201`（默认）、`dav-3510`	NPU 架构：`dav-2201`对应 Atlas A2 系列产品与 Atlas A3 系列产品，`dav-3510`对应 Ascend 950PR & 950DT 系列产品

从 CMakeLists.txt 源码可以进一步确认该构建的依赖关系：

find_package(ASC REQUIRED)引入 Ascend C 工具链，project(ge_torchair LANGUAGES ASC CXX)同时启用 ASC 与 CXX 两种语言，C++ 标准为 C++17；
三个.asc源文件（add_custom_kernel.asc、add_custom_pytorch.asc、add_custom_ge.asc）共同编译进ascendc_ops共享库，即最终的libascendc_ops.so；
头文件搜索路径包含${ASCEND_HOME_PATH}/include下的graph、register、external子目录，分别对应 GE 图接口（graph/custom_op.h、graph/operator_reg.h）、算子注册接口和外部扩展接口；
链接阶段链接Python3::Python、gert与 Torch 库，并使用-Wl,--no-as-needed强制保留这些依赖——这是确保动态库加载时 GE 交付件与 PyTorch 算子注册全局构造代码生效的关键；
工程通过--gcc-toolchain编译选项将 CMake 检测到的 GCC 工具链根目录传递给编译器，Ascend C 编译与 C++ 编译共用同一工具链，保证 host 侧交付件代码的 ABI 一致性。

运行图模式验证测试

add_custom_test.py 的测试逻辑为：

加载libascendc_ops.so并注册 TorchAir converter；
在 CPU 侧用torch.rand生成[8, 2048]的 float32 输入x_cpu、y_cpu（torch.manual_seed(0)固定随机种子）；
将模型移到 NPU（AddCustomModel().npu()），通过 TorchAir 后端编译后执行opt_model(x_cpu.npu(), y_cpu.npu())；
将 NPU 输出拷回 CPU，用torch.testing.assert_close(output, golden, rtol=1e-4, atol=1e-4)对比图模式输出与 CPU 标准加法torch.add(x_cpu, y_cpu)的结果。

执行完成后输出如下，说明 TorchAir 图模式执行结果与 CPU 标准加法结果对齐：

Ran 1 test in **s OK

前置环境要求

运行前需完成以下准备：

安装 PyTorch、Ascend Extension for PyTorch 插件以及 TorchAir：选取支持的 Python 版本配套发行版，完成torch、torch-npu和torchair的安装；
配置环境变量：根据当前环境上 CANN 开发套件包的安装方式配置环境变量：
```
source ${install_path}/cann/set_env.sh
```
${install_path}为 CANN 包安装目录，未指定安装目录时默认安装至/usr/local/Ascend下。

小结与扩展要点

该样例展示了 Ascend C 算子接入 PyTorch 图模式的四个交付件分工：kernel 负责计算、GE 交付件负责算子原型与执行入口、PyTorch 侧负责 schema 与 Meta、converter 负责 PyTorch 算子到 GE 节点的映射。基于此样例扩展实际算子时，需要关注：

当前样例的输入规格固定为[8, 2048]float 且 ND 格式，PyTorch Meta 校验（AscendcAddMeta）、GE 执行校验（IsSupportedShape）与 kernel 的BLOCK_LENGTH/BLOCK_NUM三处需同步扩展为动态规格；
编译时通过CMAKE_ASC_ARCHITECTURES指定目标架构，dav-2201面向 Atlas A2/A3 系列，dav-3510面向 Ascend 950PR & 950DT 系列；
fullgraph=True保证整图捕获，若模型中存在无法被图捕获的动态行为，需要将其改为dynamic或拆分执行，这是 PyTorch 自定义算子图模式接入中常见的调试起点。