企业官网建设流程全解析

PTO 算子集成到推理框架：从 PyTorch 到 TensorFlow/ONNX Runtime 的完整接入指南

【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa

导读：本文以 CANN pto-isa 仓库中 docs/coding/framework-integration_zh.md 为核心脉络，系统讲解如何将基于 PTO（Parallel Tile Operation）虚拟指令集编写的 Tile 级内核接入主流深度学习框架运行时。文中给出从算子 Schema 定义、Kernel 实现、框架注册、扩展编译到测试验证的完整闭环，并结合仓库内 demos/baseline/add 的真实示例代码与 include/pto/pto-inst.hpp 头文件进行源码级佐证。读完本文，你将掌握在 PyTorch（含 torch_npu 与 torch.library 两种路径）、TensorFlow、ONNX Runtime、MindSpore Lite 中接入 PTO 内核的通用实现模式，并了解算子融合、内存优化、异步执行等性能优化手段。

1. 集成概述

1.1 集成架构

基于 PTO 的内核接入框架时，整体调用链自上而下分为四层：

应用层 ↓ 框架运行时与算子注册层 ↓ 基于 PTO 的内核实现 ↓ 目标 Ascend AI Core 执行

• 应用层：用户调用框架算子 API（如torch.ops.npu.my_add）；
• 框架运行时与算子注册层：负责算子 Schema 解析、Dispatch 分发、形状/类型推导、Autograd 调度；
• 基于 PTO 的内核实现：以__global__ __aicore__内核函数 + Tile 指令（如TLOAD/TADD/TSTORE）组成的计算逻辑，统一入口见 include/pto/pto-inst.hpp，该头文件根据编译宏（__CPU_SIM、__CCE_AICORE__、__COSTMODEL）分别引入 CPU 模拟、AI Core 编译或成本模型对应的指令实现；
• 目标 Ascend AI Core 执行：内核被编译为二进制后在 Ascend 芯片的 AI Core（含 Vector Core / Cube Core）上执行。

1.2 集成方式对比

1.3 集成流程

一个完整的算子接入通常包含五个阶段：

1. 定义算子接口 ├─ 输入/输出 Tensor 规格 ├─ 参数类型和默认值 └─ 算子属性（inplace、deterministic） 2. 实现算子逻辑 ├─ 前向计算 ├─ 反向传播（训练） └─ 形状推导 3. 注册算子 ├─ 框架算子注册 ├─ 后端绑定 └─ 类型推导 4. 测试验证 ├─ 单元测试 ├─ 数值正确性 └─ 性能基准测试 5. 文档和示例 ├─ API 文档 ├─ 使用示例 └─ 性能报告

注意：不同框架版本、torch_npu集成方式以及产品发布分支的注册接口可能不同，下文所有示例均应视为"实现思路"，而不是可直接照搬的模板代码。

2. PyTorch 集成

2.1 通过 torch_npu 集成

步骤 1：定义算子 Schema

PyTorch 使用TORCH_LIBRARY/TORCH_LIBRARY_FRAGMENT声明算子 schema，使其可从 Python 通过torch.ops.<namespace>.<op_name>调用。以下示例在npu命名空间注册多种形态的算子：

// my_ops.cpp #include <torch/extension.h> #include <torch_npu/csrc/framework/utils/OpAdapter.h> // 定义算子 schema TORCH_LIBRARY_FRAGMENT(npu, m) { // 基本算子 m.def("my_add(Tensor x, Tensor y) -> Tensor"); // 带标量参数 m.def("my_mul(Tensor x, Scalar alpha) -> Tensor"); // 多输出 m.def("my_split(Tensor x, int dim) -> (Tensor, Tensor)"); // inplace 算子 m.def("my_relu_(Tensor(a!) self) -> Tensor(a!)"); // 可选参数 m.def("my_conv(Tensor input, Tensor weight, Tensor? bias=None, " "int stride=1, int padding=0) -> Tensor"); }

仓库中的真实示例与上述模式完全一致。在 demos/baseline/add/csrc/host/my_add.cpp 中，仅用 4 行即完成了npu::my_add的 schema 声明：

TORCH_LIBRARY_FRAGMENT(npu, m) { // Declare the custom operator schema m.def("my_add(Tensor x, Tensor y) -> Tensor"); }

之后 Python 端即可通过torch.ops.npu.my_add调用。

步骤 2：实现算子

（1）PTO Kernel 实现（简单写法）

#include <pto/pto-inst.hpp> // PTO Kernel 实现 __global__ __aicore__ void MyAddKernel( __gm__ float* out, __gm__ const float* x, __gm__ const float* y, uint32_t length) { int block_idx = get_block_idx(); int block_num = get_block_num(); int elements_per_block = (length + block_num - 1) / block_num; int start = block_idx * elements_per_block; int end = min(start + elements_per_block, length); using TileT = Tile<TileType::Vec, float, 16, 256>; for (int i = start; i < end; i += 16 * 256) { TileT tile_x, tile_y, tile_out; TLOAD(tile_x, GlobalTensor(x + i)); TLOAD(tile_y, GlobalTensor(y + i)); TADD(tile_out, tile_x, tile_y); TSTORE(GlobalTensor(out + i), tile_out); } }

仓库中demos/baseline/add的 kernel 则展示了生产级的写法：在 demos/baseline/add/csrc/kernel/add_custom.cpp 中，通过TASSIGN为 ping-pong 双缓冲的 Tile 分配 UB 地址，用TLOAD从全局内存加载数据、TADD执行逐元素加法、TSTORE写回，并用set_flag/wait_flag在PIPE_V（向量）、PIPE_MTE2（搬运入）、PIPE_MTE3（搬运出）之间做流水同步。其runTAdd模板函数还做了两层 tile 划分：

• 核间划分：bTileRows = tileRows / BLOCK_ROWS，按BLOCK_DIM（20 个 AIV）切分行；
• 核内划分：tileSCols = bTileCols / tileNum / BUFFER_NUM，配合 ping-pong 双缓冲提升搬运与计算的重叠度。

其中 UB 缓冲布局以常量显式给出（X_PING=0x0、Y_PING=0x10000、Z_PING=0x20000，每块 0x8000 字节），并定义了MAX_TILE_SIZE = 0x10000 - 0x100，配合static_assert在编译期校验 tile 尺寸不超 UB 容量（本示例面向 A2A3 的 192KB UB）。

（2）Host 侧算子实现（简单算子）

// PyTorch 算子实现 at::Tensor my_add_impl(const at::Tensor& x, const at::Tensor& y) { // 检查输入 TORCH_CHECK(x.device() == y.device(), "Inputs must be on same device"); TORCH_CHECK(x.sizes() == y.sizes(), "Inputs must have same shape"); TORCH_CHECK(x.scalar_type() == at::kFloat, "Only float32 supported"); // 分配输出 at::Tensor out = at::empty_like(x); // 获取数据指针 float* out_ptr = out.data_ptr<float>(); const float* x_ptr = x.data_ptr<float>(); const float* y_ptr = y.data_ptr<float>(); uint32_t length = x.numel(); // 启动 kernel int block_num = 24; // A3 核心数 EXEC_KERNEL_CMD(MyAddKernel, block_num, out_ptr, x_ptr, y_ptr, length); return out; }

仓库真实实现见 demos/baseline/add/csrc/host/my_add.cpp：分配输出z = at::empty_like(x)，计算totalLength（各维度乘积），以blockDim = 20启动add_customkernel，并通过EXEC_KERNEL_CMD宏入队执行。

其中EXEC_KERNEL_CMD宏定义于 demos/baseline/add/csrc/host/utils.h，其内部完成三件关键工作：

1. ConvertTypes将at::Tensor转换为裸指针（at_tensor.storage().data()），标量则原样透传；
2. 通过c10_npu::getCurrentNPUStream().stream(false)获取当前 NPU Stream；
3. 调用ACLRT_LAUNCH_KERNEL(kernel_name)(blockdim, acl_stream, params...)真正入队内核，并经at_npu::native::OpCommand::RunOpApi接入 torch_npu 的算子执行框架。

（3）复杂算子实现（带反向传播）

通过继承torch::autograd::Function实现可微算子：

// 前向 class MyConvFunction : public torch::autograd::Function<MyConvFunction> { public: static at::Tensor forward( torch::autograd::AutogradContext* ctx, const at::Tensor& input, const at::Tensor& weight, const at::Tensor& bias, int stride, int padding) { // 保存用于反向传播的张量 ctx->save_for_backward({input, weight, bias}); ctx->saved_data["stride"] = stride; ctx->saved_data["padding"] = padding; // 调用 PTO kernel at::Tensor output = run_conv_forward(input, weight, bias, stride, padding); return output; } static std::vector<at::Tensor> backward( torch::autograd::AutogradContext* ctx, std::vector<at::Tensor> grad_outputs) { // 恢复保存的张量 auto saved = ctx->get_saved_variables(); auto input = saved[0]; auto weight = saved[1]; auto bias = saved[2]; int stride = ctx->saved_data["stride"].toInt(); int padding = ctx->saved_data["padding"].toInt(); auto grad_output = grad_outputs[0]; // 计算梯度 at::Tensor grad_input = run_conv_backward_input( grad_output, weight, stride, padding); at::Tensor grad_weight = run_conv_backward_weight( grad_output, input, stride, padding); at::Tensor grad_bias = run_conv_backward_bias(grad_output); return {grad_input, grad_weight, grad_bias, at::Tensor(), at::Tensor()}; // stride, padding 无梯度 } }; // 包装函数 at::Tensor my_conv( const at::Tensor& input, const at::Tensor& weight, const at::Tensor& bias, int stride, int padding) { return MyConvFunction::apply(input, weight, bias, stride, padding); }

要点：save_for_backward保存反向所需的中间张量，saved_data保存标量参数；backward中非张量参数（如stride、padding）返回空at::Tensor()表示无梯度。

步骤 3：注册实现

通过TORCH_LIBRARY_IMPL将实现绑定到后端。对 NPU 执行而言，torch_npu 使用PrivateUse1dispatch key（关于PrivateUse1的详细机制可参考 PyTorch 官方 privateuseone 文档）：

// 注册到 NPU 后端 TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) { m.impl("my_add", TORCH_FN(my_add_impl)); m.impl("my_mul", TORCH_FN(my_mul_impl)); m.impl("my_conv", TORCH_FN(my_conv)); } // 注册 autograd TORCH_LIBRARY_IMPL(npu, Autograd, m) { m.impl("my_conv", TORCH_FN(my_conv)); }

仓库中的真实注册代码见 demos/baseline/add/csrc/host/my_add.cpp：

TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) { // Register the custom operator implementation function m.impl("my_add", TORCH_FN(ascendc_path::run_add_custom)); }

步骤 4：编译为 Python 扩展

setup.py（通用写法）：

from setuptools import setup from torch.utils.cpp_extension import BuildExtension, CppExtension setup( name='my_pto_ops', ext_modules=[ CppExtension( name='my_pto_ops', sources=['my_ops.cpp'], include_dirs=[ '/path/to/pto-isa/include', '/path/to/torch_npu/include', ], library_dirs=[ '/path/to/pto-isa/lib', ], libraries=['pto'], extra_compile_args=['-std=c++20', '-O3'], ) ], cmdclass={'build_ext': BuildExtension} )

编译：

python setup.py install

仓库中的真实构建方案则更为工程化：demos/baseline/add/setup.py 自定义了CPPLibBuild（build_clib）与Build（build_ext）两个命令类：

• 自动探测cmake/cmake3（要求版本 ≥ 3.18.0）；
• 通过torch.compiled_with_cxx11_abi()自动判断并传入GLIBCXX_USE_CXX11_ABI=1/0，避免 ABI 不兼容；
• 将-DTORCH_PATH、-DTORCH_NPU_PATH传给 CMake，最终产物（.so）被拷贝进op_extension/lib随 wheel 一起打包。

对应 demos/baseline/add/CMakeLists.txt 中：ascendc_library(no_workspace_kernel STATIC csrc/kernel/add_custom.cpp)负责把 PTO kernel 编译为静态库；add_library(op_extension SHARED ...)将 host 侧代码（csrc/host/*.cpp）编译为共享库，并链接no_workspace_kernel、torch_npu、ascendcl、tiling_api、register、platform、ascendalog等依赖；ascendc_include_directories将${PTO_LIB_PATH}/include与${PTO_LIB_PATH}/include/pto/common加入头文件搜索路径。

完整构建运行流程（仓库 demos/baseline/add/run.sh 给出示例脚本，其中PTO_LIB_PATH需指向本仓库路径）：

export ASCEND_HOME_PATH=/usr/local/Ascend/ source /usr/local/Ascend/ascend-toolkit/set_env.sh export PTO_LIB_PATH=[YOUR_PATH]/pto-isa rm -fr build op_extension.egg-info python3 setup.py bdist_wheel cd dist pip uninstall *.whl pip install *.whl

设置目标 SoC：编辑demos/baseline/add/CMakeLists.txt，将SOC_VERSION设置为目标芯片（例如 A2/A3 使用ascend910b1）：

set(SOC_VERSION "ascend910b1" CACHE STRING "system on chip type")

可在目标机器上执行npu_smi info查询芯片名称，并按Ascend<Chip Name>/ascend<chip name>的形式填写。

步骤 5：Python 使用

import torch import torch_npu import my_pto_ops # 创建输入 x = torch.randn(1024, 1024).npu() y = torch.randn(1024, 1024).npu() # 调用自定义算子 z = torch.ops.npu.my_add(x, y) # 验证结果 expected = x + y assert torch.allclose(z, expected, rtol=1e-5) print("✓ Custom op works correctly!")

仓库测试用例 demos/baseline/add/test/test.py 采用了相同模式：构造[20, 2048]的 float16 随机张量，x.npu()/y.npu()后调用torch.ops.npu.my_add(x_npu, y_npu)，与 CPU 上的torch.add(x, y)通过self.assertRtolEqual校验数值一致性。

2.2 通过 torch.library 集成（PyTorch 2.0+）

对于纯 Python 原型或算子主体已在 C++ 侧实现的情形，可使用torch.library.custom_op完成更简洁的注册，并显式提供 fake 实现用于形状推导（compile/fake tensor 模式必需）：

import torch from torch.library import custom_op @custom_op("mylib::my_add", mutates_args=()) def my_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: """自定义加法算子""" return torch.ops.mylib.my_add_impl(x, y) @my_add.register_fake def _(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: """形状推导""" assert x.shape == y.shape return torch.empty_like(x) # 使用 x = torch.randn(10, 10) y = torch.randn(10, 10) z = torch.ops.mylib.my_add(x, y)

2.3 完整示例：Add 算子

基于 PTO 的完整 Add 算子端到端示例（kernel 实现、host 注册、wheel 构建、NPU 测试）请参考仓库内的 demos/baseline/add/README_zh.md，其中详细说明了目录结构、ascendc_library构建配置、PrivateUse1注册原理以及构建运行步骤。

3. TensorFlow 集成

3.1 自定义 Op

步骤 1：定义 Op

// my_ops.cc #include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/shape_inference.h" REGISTER_OP("MyAdd") .Input("x: float") .Input("y: float") .Output("z: float") .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { // 形状推导 c->set_output(0, c->input(0)); return tensorflow::Status::OK(); }) .Doc(R"doc( 自定义加法算子 Args: x: 第一个输入张量 y: 第二个输入张量 Returns: z: x + y )doc");

步骤 2：实现 Kernel

在OpKernel::Compute中获取输入、校验形状、分配输出，然后通过EXEC_KERNEL_CMD启动 PTO kernel：

#include "tensorflow/core/framework/op_kernel.h" #include <pto/pto-inst.hpp> class MyAddOp : public tensorflow::OpKernel { public: explicit MyAddOp(tensorflow::OpKernelConstruction* context) : OpKernel(context) {} void Compute(tensorflow::OpKernelContext* context) override { // 获取输入 const tensorflow::Tensor& x = context->input(0); const tensorflow::Tensor& y = context->input(1); // 检查形状 OP_REQUIRES(context, x.shape() == y.shape(), tensorflow::errors::InvalidArgument( "Inputs must have same shape")); // 分配输出 tensorflow::Tensor* z = nullptr; OP_REQUIRES_OK(context, context->allocate_output(0, x.shape(), &z)); // 调用 PTO kernel const float* x_ptr = x.flat<float>().data(); const float* y_ptr = y.flat<float>().data(); float* z_ptr = z->flat<float>().data(); uint32_t length = x.NumElements(); EXEC_KERNEL_CMD(MyAddKernel, 24, z_ptr, x_ptr, y_ptr, length); } }; // 注册 kernel REGISTER_KERNEL_BUILDER( Name("MyAdd").Device(tensorflow::DEVICE_NPU), MyAddOp);

步骤 3：编译

使用 TensorFlow 提供的编译/链接标志，并链接 PTO 库：

# 使用 TensorFlow 的编译工具 TF_CFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))') ) TF_LFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))') ) g++ -std=c++17 -shared my_ops.cc -o my_ops.so \ ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} \ -I/path/to/pto-isa/include \ -L/path/to/pto-isa/lib -lpto \ -fPIC -O3

步骤 4：Python 使用

import tensorflow as tf # 加载自定义 op my_ops = tf.load_op_library('./my_ops.so') # 使用 x = tf.constant([[1.0, 2.0], [3.0, 4.0]]) y = tf.constant([[5.0, 6.0], [7.0, 8.0]]) z = my_ops.my_add(x, y) print(z.numpy()) # [[6. 8.] # [10. 12.]]

3.2 注册梯度

对需要训练的算子，用tf.RegisterGradient注册梯度函数。对于加法z = x + y，∂z/∂x 与 ∂z/∂y 均为单位映射：

@tf.RegisterGradient("MyAdd") def _my_add_grad(op, grad): """MyAdd 的梯度""" return grad, grad # ∂z/∂x = 1, ∂z/∂y = 1

4. ONNX Runtime 集成

4.1 自定义 Execution Provider

ONNX Runtime 的扩展机制是自定义 Execution Provider（EP），在 EP 内部注册基于OpKernel的算子实现。

步骤 1：定义 Kernel

// my_onnx_ops.cc #include "onnxruntime/core/framework/op_kernel.h" class MyAddKernel : public onnxruntime::OpKernel { public: MyAddKernel(const onnxruntime::OpKernelInfo& info) : OpKernel(info) {} onnxruntime::Status Compute(onnxruntime::OpKernelContext* context) const override { // 获取输入 const onnxruntime::Tensor* X = context->Input<onnxruntime::Tensor>(0); const onnxruntime::Tensor* Y = context->Input<onnxruntime::Tensor>(1); // 分配输出 onnxruntime::Tensor* Z = context->Output(0, X->Shape()); // 调用 PTO kernel const float* x_data = X->Data<float>(); const float* y_data = Y->Data<float>(); float* z_data = Z->MutableData<float>(); size_t length = X->Shape().Size(); EXEC_KERNEL_CMD(MyAddKernel, 24, z_data, x_data, y_data, length); return onnxruntime::Status::OK(); } };

步骤 2：注册 Kernel

通过ONNX_OPERATOR_KERNEL_EX宏将Add算子（kOnnxDomain、opset 7）绑定到kNpuExecutionProvider：

ONNX_OPERATOR_KERNEL_EX( Add, kOnnxDomain, 7, // opset version kNpuExecutionProvider, MyAddKernel);

步骤 3：创建 Execution Provider

自定义 EP 继承onnxruntime::IExecutionProvider，通过GetCapability向运行时声明本 EP 支持的算子子图：

class NpuExecutionProvider : public onnxruntime::IExecutionProvider { public: NpuExecutionProvider() : IExecutionProvider(kNpuExecutionProvider) {} std::vector<std::unique_ptr<onnxruntime::ComputeCapability>> GetCapability(const onnxruntime::GraphViewer& graph, const std::vector<const onnxruntime::KernelRegistry*>& registries) const override { // 返回支持的算子 // ... } };

步骤 4：Python 使用

import onnxruntime as ort # 注册自定义 EP session_options = ort.SessionOptions() session_options.register_custom_ops_library('my_onnx_ops.so') # 创建会话 session = ort.InferenceSession( 'model.onnx', session_options, providers=['NpuExecutionProvider', 'CPUExecutionProvider'] ) # 推理 outputs = session.run(None, {'input': input_data})

注意providers列表中的顺序即优先级：将NpuExecutionProvider放在前面，可使算子优先由 NPU EP 接管；未被 NPU EP 支持的算子会回落到CPUExecutionProvider，从而保证模型整体可运行。

5. 推理框架集成（MindSpore Lite）

MindSpore Lite 通过REGISTER_CUSTOM_KERNEL注册自定义算子内核，Execute中从in_tensors_/out_tensors_取输入输出并调用 PTO kernel：

// 注册自定义算子 #include "include/registry/register_kernel.h" class MyAddKernel : public mindspore::kernel::Kernel { public: int Prepare() override { return RET_OK; } int Execute() override { auto input0 = in_tensors_[0]; auto input1 = in_tensors_[1]; auto output = out_tensors_[0]; // 调用 PTO kernel // ... return RET_OK; } }; // 注册 REGISTER_CUSTOM_KERNEL(NPU, MyProvider, kNumberTypeFloat32, Add, MyAddKernel);

6. 性能优化

6.1 算子融合

减少 kernel 启动次数与中间张量落盘是框架侧最直接有效的优化。PyTorch 示例中，Add + ReLU可融合为一个自定义算子：

# PyTorch 示例：融合 Add + ReLU @torch.jit.script def fused_add_relu(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return torch.relu(x + y) # 使用自定义融合算子替换 torch.ops.npu.fused_add_relu(x, y)

仓库中已有对应的融合算子实战示例：kernels/custom/fused_add_relu_mul 目录包含融合 kernel 的实现（.cpp）、说明文档（.md）与运行脚本（.sh），展示了在 PTO 层面将多个算子合并为单个 kernel 的完整写法，可作为算子融合落地时的参考。更系统的优化方法论见 docs/coding/opt_zh.md 与 docs/coding/performance-best-practices_zh.md。

6.2 内存优化

优先支持 inplace 操作，避免不必要的输出分配与拷贝：

// Inplace 算子 at::Tensor& my_add_inplace(at::Tensor& x, const at::Tensor& y) { // 直接修改 x，避免分配新内存 float* x_ptr = x.data_ptr<float>(); const float* y_ptr = y.data_ptr<float>(); uint32_t length = x.numel(); EXEC_KERNEL_CMD(MyAddInplaceKernel, 24, x_ptr, y_ptr, length); return x; }

在 PTO kernel 内部，仓库示例 demos/baseline/add/csrc/kernel/add_custom.cpp 通过 ping-pong 双缓冲（BUFFER_NUM = 2）复用同一段 UB 地址，配合TASSIGN固定 Tile 的 UB 基址，将内存搬移与向量计算重叠，从而在有限的 UB 空间内最大化吞吐。

6.3 异步执行

在 stream 上异步启动 kernel，避免阻塞主线程等待：

// 使用 CUDA Stream（或 NPU Stream） at::Tensor my_add_async(const at::Tensor& x, const at::Tensor& y) { at::Tensor out = at::empty_like(x); // 获取当前 stream auto stream = at::cuda::getCurrentCUDAStream(); // 异步启动 kernel EXEC_KERNEL_ASYNC(MyAddKernel, 24, stream, out.data_ptr<float>(), x.data_ptr<float>(), y.data_ptr<float>(), x.numel()); return out; }

在 torch_npu 场景下，仓库的EXEC_KERNEL_CMD宏（demos/baseline/add/csrc/host/utils.h）本身就是异步语义：它通过c10_npu::getCurrentNPUStream()取得当前 NPU stream 后将内核入队，配合at_npu::native::OpCommand::RunOpApi统一调度，宿主线程不会被内核执行阻塞。

7. 调试与测试

7.1 单元测试

以 PyTorch 为例，使用unittest校验前向正确性与反向梯度：

import unittest import torch import my_pto_ops class TestMyOps(unittest.TestCase): def test_my_add(self): x = torch.randn(100, 100).npu() y = torch.randn(100, 100).npu() # 自定义算子 z_custom = torch.ops.npu.my_add(x, y) # 参考实现 z_ref = x + y # 验证 self.assertTrue(torch.allclose(z_custom, z_ref, rtol=1e-5)) def test_my_add_backward(self): x = torch.randn(100, 100, requires_grad=True).npu() y = torch.randn(100, 100, requires_grad=True).npu() z = torch.ops.npu.my_add(x, y) loss = z.sum() loss.backward() # 验证梯度 self.assertIsNotNone(x.grad) self.assertIsNotNone(y.grad) self.assertTrue(torch.allclose(x.grad, torch.ones_like(x))) if __name__ == '__main__': unittest.main()

仓库测试用例 demos/baseline/add/test/test.py 采用相同的"自定义算子 vs 参考实现"对比策略，但通过torch_npu.testing.testcase.TestCase.assertRtolEqual提供容差比较。此外，PTO 提供 CPU 模拟执行能力（include/pto/pto-inst.hpp中__CPU_SIM分支），可将同一 kernel 编译到 CPU 上验证逻辑正确性，是调试数值问题的有力手段；完整的调试方法与工具链见 docs/coding/debug_zh.md。

7.2 性能基准测试

基准测试必须包含预热（warmup）与同步（synchronize）步骤，避免首次调用开销与异步执行造成测量偏差：

import torch import time def benchmark(func, *args, warmup=10, iterations=100): # 预热 for _ in range(warmup): func(*args) # 同步 torch.npu.synchronize() # 测量 start = time.time() for _ in range(iterations): func(*args) torch.npu.synchronize() end = time.time() avg_time = (end - start) / iterations * 1000 # ms return avg_time # 对比性能 x = torch.randn(1024, 1024).npu() y = torch.randn(1024, 1024).npu() time_custom = benchmark(lambda: torch.ops.npu.my_add(x, y)) time_builtin = benchmark(lambda: x + y) print(f"Custom op: {time_custom:.3f} ms") print(f"Built-in op: {time_builtin:.3f} ms") print(f"Speedup: {time_builtin / time_custom:.2f}x")

提示：性能数据应基于目标硬件实测获得，不同芯片（A2/A3 等）的 AIV 数量、UB 容量不同，kernel 的blockDim与 tile 尺寸需按目标 SoC 调优，仓库 demos/baseline/add/CMakeLists.txt 中的SOC_VERSION即用于控制编译目标。

8. 最佳实践

8.1 设计原则

✅DO：

• 保持算子接口简单清晰
• 提供完整的类型支持（float32, float16, int32 等）
• 实现形状推导和类型推导
• 提供详细的文档和示例
• 编写完整的单元测试

❌DON'T：

• 不要在算子内部分配大量临时内存
• 不要假设输入总是连续的（使用contiguous()）
• 不要忽略边界情况（空张量、单元素张量）
• 不要在算子内部使用全局状态

8.2 性能检查清单

• 算子是否支持 inplace 操作
• 是否实现了算子融合
• 是否使用了异步执行
• 是否避免了不必要的内存拷贝
• 是否支持多种数据类型
• 是否进行了性能基准测试

8.3 兼容性检查清单

• 是否支持动态形状
• 是否支持广播语义
• 是否支持梯度计算（训练）
• 是否支持 JIT 编译
• 是否支持导出为 ONNX
• 是否提供 CPU fallback

参考资源

• PTO Add 算子示例（PyTorch 集成完整教程）
• PTO Add 算子 Host 侧注册实现
• EXEC_KERNEL_CMD 启动宏与工具函数
• PTO Add 算子 Kernel（Tile 指令与流水同步）
• PTO 内核统一头文件（多后端指令分发）
• 算子调试指南
• 性能优化指南
• 性能最佳实践
• 融合算子实战示例（fused_add_relu_mul）
• PTO 虚拟指令集手册

【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa