将 PTO 内核集成到深度学习框架运行时torch_npu、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 仓库的 framework-integration 文档 展开系统讲解如何把基于 PTOParallel Tile Operation虚拟指令集实现的 AI Core 内核暴露给 PyTorchtorch_npu、torch.library、TensorFlow、ONNX Runtime、MindSpore Lite 等框架运行时并配套算子融合、内存优化、异步执行、单元测试与性能基准等工程实践。读完本文你将掌握从写好一个 PTO kernel到在训练/推理框架中一键调用的完整集成链路并能对照仓库中的可运行示例demos/baseline/add亲手复现。1. 集成概述1.1 集成架构PTO 内核本身运行在 Ascend AI Core矢量/矩阵核上而深度学习框架运行在 Host 侧二者之间需要算子注册层作为桥梁。整体分层如下应用层 ↓ 框架运行时与算子注册层 ↓ 基于 PTO 的内核实现 ↓ 目标 Ascend AI Core 执行其中框架运行时与算子注册层负责声明算子 Schema参数、输出、属性、把 Host 侧算子实现绑定到 NPU 后端、提供类型/形状推导基于 PTO 的内核实现则包含 kernel 本体使用 pto-inst.hpp 提供的 TLOAD/TADD/TSTORE 等 Tile 指令与 Host 侧启动逻辑如EXEC_KERNEL_CMD。1.2 集成方式对比集成方式优点缺点适用场景Python 扩展开发快速、易调试Python/C 边界存在额外开销原型开发、快速验证C 扩展性能好、类型安全构建和注册流程更复杂生产环境、性能关键框架插件 / 自定义后端路径更贴近部署形态维护成本更高稳定产品集成注意不同框架版本、torch_npu集成方式以及产品发布分支的注册接口可能不同下文示例应视为实现思路而非可直接照搬的模板代码。1.3 集成工作流无论面向哪个框架完整集成都遵循同一套五步流程1. 定义算子接口 ├─ 输入/输出 Tensor 规格 ├─ 参数类型和默认值 └─ 算子属性inplace、deterministic 2. 实现算子逻辑 ├─ 前向计算 ├─ 反向传播训练 └─ 形状推导 3. 注册算子 ├─ 框架算子注册 ├─ 后端绑定 └─ 类型推导 4. 测试验证 ├─ 单元测试 ├─ 数值正确性 └─ 性能基准测试 5. 文档和示例 ├─ API 文档 ├─ 使用示例 └─ 性能报告2. PyTorch 集成PyTorch 集成是仓库中落地最完整的一条链路。仓库 demos/baseline/add 提供了可直接构建运行的完整示例其目录结构如下demos/baseline/add/ ├── op_extension/ # Python 包入口模块加载 ├── csrc/ │ ├── kernel/ # PTO kernel 实现 │ └── host/ # Host 侧 PyTorch 算子注册 ├── test/ # 最小化 Python 测试 ├── CMakeLists.txt # 构建配置 ├── setup.py # Wheel 构建脚本 └── README.md # 示例文档2.1 通过 torch_npu 集成步骤 1定义算子 SchemaPyTorch 使用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? biasNone, int stride1, int padding0) - Tensor); }仓库示例 demos/baseline/add/csrc/host/my_add.cpp 中的真实写法与此一致——它在匿名命名空间中注册了npu::my_addTORCH_LIBRARY_FRAGMENT(npu, m) { // Declare the custom operator schema m.def(my_add(Tensor x, Tensor y) - Tensor); }步骤 2实现算子简单算子实现先写 PTO kernel再用at::Tensor包装。PTO kernel 使用__global__ __aicore__声明__gm__标记全局内存指针通过get_block_idx()/get_block_num()做多核切分配合Tile与TLOAD/TADD/TSTORE完成 Tile 级计算#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 TileTileType::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); } } // 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_ptrfloat(); const float* x_ptr x.data_ptrfloat(); const float* y_ptr y.data_ptrfloat(); 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; }仓库中的真实 kernel demos/baseline/add/csrc/kernel/add_custom.cpp 更完整它使用half类型、20 个 AIVBLOCK_DIM 20、192KB UBUB_SIZE 0x30000、双缓冲 ping-pong 流水BUFFER_NUM 2并通过set_flag/wait_flag在 MTE2搬入、V矢量计算、MTE3搬出三条硬件流水线之间做同步TASSIGN为每个 Tile 显式指定 UB 地址set_mask_norm()与set_vector_mask(-1, -1)设置全量 mask。这个实现演示了在框架集成场景下 kernel 侧多核分块 乒乓流水的典型写法。Host 侧实现位于 demos/baseline/add/csrc/host/my_add.cpp它引入构建系统生成的aclrtlaunch_add_custom.h计算blockDim 20与totalLength后调用EXEC_KERNEL_CMD(add_custom, blockDim, x, y, z, totalLength)。复杂算子实现带反向传播需要训练支持时通过torch::autograd::Function自定义反向用ctx-save_for_backward保存前向张量、ctx-saved_data保存标量参数在backward中恢复并计算三个梯度// 前向 class MyConvFunction : public torch::autograd::FunctionMyConvFunction { 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::vectorat::Tensor backward( torch::autograd::AutogradContext* ctx, std::vectorat::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); }步骤 3注册实现对 NPU 执行而言torch_npu使用PrivateUse1dispatch key。将实现注册到PrivateUse1后torch.ops.npu.my_add在 NPU 张量上就会路由到我们的 PTO 实现带自动求导的算子还需额外注册到Autogradkey// 注册到 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 的注册代码与之一致m.impl(my_add, TORCH_FN(ascendc_path::run_add_custom))。值得一提的是EXEC_KERNEL_CMD宏定义在 demos/baseline/add/csrc/host/utils.h它通过ConvertTypes将at::Tensor统一转换为裸指针获取当前 NPU streamc10_npu::getCurrentNPUStream()最终经ACLRT_LAUNCH_KERNEL入队并用at_npu::native::OpCommand::RunOpApi保证与 torch_npu 的算子执行模型对齐。步骤 4编译为 Python 扩展setup.py基于torch.utils.cpp_extension的最小写法from setuptools import setup from torch.utils.cpp_extension import BuildExtension, CppExtension setup( namemy_pto_ops, ext_modules[ CppExtension( namemy_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[-stdc20, -O3], ) ], cmdclass{build_ext: BuildExtension} )仓库示例的 setup.py 则采用CMake 编译 kernel 静态库 NpuExtension 打包的方式CPPLibBuild命令会调用 CMake要求 3.18自动探测TORCH_PATH/TORCH_NPU_PATH并依据torch.compiled_with_cxx11_abi()设置-DGLIBCXX_USE_CXX11_ABI将编译产物libop_extension.so拷贝进op_extension/libop_extension/_load.py 在包导入时通过torch.ops.load_library加载该 so。对应的 CMakeLists.txt 关键点# 将 kernel 编译为静态库使用 CANN 的 ascendc cmake 工具链 ascendc_library(no_workspace_kernel STATIC csrc/kernel/add_custom.cpp ) # 注入 PTO 头文件搜索路径需要预先 export PTO_LIB_PATH ascendc_include_directories(no_workspace_kernel PRIVATE ${PTO_LIB_PATH}/include ${PTO_LIB_PATH}/include/pto/common )编译python setup.py install步骤 5Python 使用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, rtol1e-5) print(✓ Custom op works correctly!)仓库测试 demos/baseline/add/test/test.py 验证了同样的调用路径它创建[20, 2048]的 float16 随机张量调用torch.ops.npu.my_add(x_npu, y_npu)再与 CPU 上的torch.add结果用assertRtolEqual比对。示例的运行环境约束可从 requirements.txt 看到torch2.8.0、torch-npu2.8.0.post2。2.2 通过 torch.library 集成PyTorch 2.0PyTorch 2.0 提供更简洁的 Python 侧注册方式torch.library.custom_op配合register_fake完成形状/类型推导meta kernel无需编写 C schemaimport 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 算子仓库 demos/baseline/add/README_zh.md 给出了从 kernel 实现、PyTorch 集成、Wheel 构建到测试运行的全流程教程。构建时需注意以下几点设置目标 SoC编辑demos/baseline/add/CMakeLists.txt将SOC_VERSION设置为目标芯片例如 A2/A3 使用Ascend910B1可在目标机器上执行npu_smi info查询芯片名称并按AscendChip Name的形式填写设置 CANN 与 PTO 环境后构建 Wheelexport ASCEND_HOME_PATH/usr/local/Ascend/ source /usr/local/Ascend/ascend-toolkit/set_env.sh export PTO_LIB_PATH[YOUR_PATH]/pto-isa python3 setup.py bdist_wheel安装并运行测试cd dist pip uninstall *.whl pip install *.whl cd ../test python3 test.py3. TensorFlow 集成3.1 自定义 Op步骤 1定义 Op使用REGISTER_OP声明输入/输出类型并在SetShapeFn中实现形状推导// 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(Rdoc( 自定义加法算子 Args: x: 第一个输入张量 y: 第二个输入张量 Returns: z: x y )doc);步骤 2实现 Kernel继承tensorflow::OpKernel在Compute中取输入、校验形状、分配输出然后调用 PTO kernel。与 PyTorch 路径类似这里同样通过EXEC_KERNEL_CMD完成启动24为核数参数#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.flatfloat().data(); const float* y_ptr y.flatfloat().data(); float* z_ptr z-flatfloat().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 头文件与库文件直接以g编译成动态库# 使用 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 -stdc17 -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步骤 4Python 使用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 ∂z/∂y 1因此梯度直接把上游 grad 原样返回tf.RegisterGradient(MyAdd) def _my_add_grad(op, grad): MyAdd 的梯度 return grad, grad # ∂z/∂x 1, ∂z/∂y 14. ONNX Runtime 集成4.1 自定义 Execution Provider步骤 1定义 Kernel继承onnxruntime::OpKernel在Compute中从OpKernelContext取输入、分配输出并调用 PTO 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-Inputonnxruntime::Tensor(0); const onnxruntime::Tensor* Y context-Inputonnxruntime::Tensor(1); // 分配输出 onnxruntime::Tensor* Z context-Output(0, X-Shape()); // 调用 PTO kernel const float* x_data X-Datafloat(); const float* y_data Y-Datafloat(); float* z_data Z-MutableDatafloat(); 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把 ONNX 标准Add算子opset 7绑定到自定义的 NPU Execution ProviderONNX_OPERATOR_KERNEL_EX( Add, kOnnxDomain, 7, // opset version kNpuExecutionProvider, MyAddKernel);步骤 3创建 Execution Provider实现IExecutionProvider接口在GetCapability中声明该 EP 支持的算子集合决定图切分时哪些节点落到 NPU 上class NpuExecutionProvider : public onnxruntime::IExecutionProvider { public: NpuExecutionProvider() : IExecutionProvider(kNpuExecutionProvider) {} std::vectorstd::unique_ptronnxruntime::ComputeCapability GetCapability(const onnxruntime::GraphViewer graph, const std::vectorconst onnxruntime::KernelRegistry* registries) const override { // 返回支持的算子 // ... } };步骤 4Python 使用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})5. 推理框架集成MindSpore Lite对移动/端侧推理场景MindSpore Lite 提供 registry 机制注册自定义算子在Execute中拿到输入输出张量后调用 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 启动意味着多次调度与访存开销将Add ReLU这类组合融合为单个算子可显著减少中间张量。先用torch.jit.script描述融合语义作为正确性参照再用自定义融合算子替换# 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 目录下提供了融合算子fused add relu mul的完整实现、编译脚本与文档可作为算子融合实战的对照参考。6.2 内存优化尽量使用 inplace 语义避免中间张量分配——在 schema 中声明Tensor(a!)并在实现中直接写回输入// Inplace 算子 at::Tensor my_add_inplace(at::Tensor x, const at::Tensor y) { // 直接修改 x避免分配新内存 float* x_ptr x.data_ptrfloat(); const float* y_ptr y.data_ptrfloat(); uint32_t length x.numel(); EXEC_KERNEL_CMD(MyAddInplaceKernel, 24, x_ptr, y_ptr, length); return x; }6.3 异步执行将 kernel 启动挂到当前 stream 上让 Host 侧不被同步等待阻塞从而实现与计算的重叠// 使用 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_ptrfloat(), x.data_ptrfloat(), y.data_ptrfloat(), x.numel()); return out; }在仓库的 torch_npu 路径中这一异步模型由 demos/baseline/add/csrc/host/utils.h 的EXEC_KERNEL_CMD落地它获取c10_npu::getCurrentNPUStream()作为 ACL stream 传入ACLRT_LAUNCH_KERNEL配合OpCommand::RunOpApi进入 torch_npu 的算子执行队列。7. 调试与测试7.1 单元测试用unittest编写前向与反向两组用例前向以 PyTorch 内建算子结果为参照做数值比对反向对requires_gradTrue的张量求梯度并核对解析值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, rtol1e-5)) def test_my_add_backward(self): x torch.randn(100, 100, requires_gradTrue).npu() y torch.randn(100, 100, requires_gradTrue).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()7.2 性能基准测试基准测试务必先预热、后计时并在计时前后torch.npu.synchronize()否则测得的是异步队列的入队时间而非真实执行时间import torch import time def benchmark(func, *args, warmup10, iterations100): # 预热 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(fCustom op: {time_custom:.3f} ms) print(fBuilt-in op: {time_builtin:.3f} ms) print(fSpeedup: {time_builtin / time_custom:.2f}x)8. 最佳实践8.1 设计原则✅DO保持算子接口简单清晰提供完整的类型支持float32, float16, int32 等实现形状推导和类型推导提供详细的文档和示例编写完整的单元测试❌DONT不要在算子内部分配大量临时内存不要假设输入总是连续的使用contiguous()不要忽略边界情况空张量、单元素张量不要在算子内部使用全局状态8.2 性能检查清单算子是否支持 inplace 操作是否实现了算子融合是否使用了异步执行是否避免了不必要的内存拷贝是否支持多种数据类型是否进行了性能基准测试8.3 兼容性检查清单是否支持动态形状是否支持广播语义是否支持梯度计算训练是否支持 JIT 编译是否支持导出为 ONNX是否提供 CPU fallback9. 仓库配套资源围绕框架集成仓库还提供了以下可直接查阅的配套材料demos/baseline/add/README_zh.mdAdd 算子集成 PyTorch 的完整分步教程含构建与运行demos/baseline/add/csrc/kernel/add_custom.cpp带乒乓流水与跨核分块的 PTO kernel 实现demos/baseline/add/csrc/host/my_add.cpp 与 demos/baseline/add/csrc/host/utils.hschema 声明、实现注册与EXEC_KERNEL_CMD启动宏demos/baseline/add/setup.py 与 demos/baseline/add/CMakeLists.txt基于 CMake NpuExtension 的 Wheel 构建链路demos/baseline/add/test/test.py最小化数值正确性测试kernels/custom/fused_add_relu_mul算子融合参考实现docs/coding/debug_zh.md算子调试指南打印、dump、逐指令排错docs/coding/opt_zh.md性能优化指南访存、流水、融合等更深层优化手段。此外仓库 demos/baseline/add 目录下还提供了add、allgather_async、flash_atten、gemm_basic等多个 baseline 示例以及 demos/auto_mode 下的自动模式Auto Mode示例可用于对比手写注册与自动生成两条集成路径的差异。若需在本地搭建环境可通过git clone获取仓库后参照 demos/baseline/add/README_zh.md 的步骤准备 CANN、PyTorch、torch_npu 与 PTO Tile Lib即可复现从 kernel 到torch.ops.npu.my_add的完整集成闭环。【免费下载链接】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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考