CANN/pto-isa Auto模式加法算子示例

自定义 PyTorch 算子KERNEL_LAUNCH示例【免费下载链接】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本示例展示如何实现一个基于 PTO 的自定义的 auto模式的 kernel并通过torch_npu将其暴露为 PyTorch 算子。目录结构demos/baseline/add/ ├── op_extension/ # Python 包入口模块加载 ├── csrc/ │ ├── kernel/ # PTO kernel 实现auto模式 │ └── host/ # Host 侧 PyTorch 算子注册 ├── test/ # 最小化 Python 测试 ├── CMakeLists.txt # 构建配置 ├── setup.py # Wheel 构建脚本 └── README.md # 本文档1. 实现 kernel在auto_mode/demos/baseline/add/csrc/kernel/下新增 kernel 源码并将其加入构建。例如要构建add_custom.cpp需要在auto_mode/demos/baseline/add/CMakeLists.txt中添加ascendc_library(no_workspace_kernel STATIC csrc/kernel/add_custom.cpp ) ascendc_compile_options(no_workspace_kernel PRIVATE --cce-enable-pto-passes -O2)不像manual模式你不需要手动调用TASSIGN和同步指令编译器会替你自动完成。注意:你需要使用--cce-enable-pto-passes编译选项来使能编译器的auto模式编译kernel时必须使用-O2目前这个用例没有使用双缓冲而且非常不建议在目前的auto模式下使用双缓冲因为并没有完全支持。构建选项与细节请参考昇腾社区文档https://www.hiascend.com/ascend-c2. 与 PyTorch 集成torch_npuHost 侧实现位于auto_mode/demos/baseline/add/csrc/host/。2.1 定义算子 schemaAten IRPyTorch 使用TORCH_LIBRARY/TORCH_LIBRARY_FRAGMENT声明算子 schema使其可从 Python 通过torch.ops.namespace.op_name调用。示例在npu命名空间注册一个自定义my_add算子TORCH_LIBRARY_FRAGMENT(npu, m) { m.def(my_add(Tensor x, Tensor y) - Tensor); }之后 Python 可通过torch.ops.npu.my_add调用。2.2 实现算子引入由构建系统生成的 kernel launch 头文件aclrtlaunch_kernel_name.h。按需分配输出张量/工作区workspace。通过ACLRT_LAUNCH_KERNEL在本示例中由EXEC_KERNEL_CMD封装将 kernel 入队执行。#include utils.h #include aclrtlaunch_add_custom.h at::Tensor run_add_custom(const at::Tensor x, const at::Tensor y) { at::Tensor z at::empty_like(x); uint32_t blockDim 20; uint32_t totalLength 1; for (uint32_t size : x.sizes()) { totalLength * size; } EXEC_KERNEL_CMD(add_custom, blockDim, x, y, z, totalLength); return z; }2.3 注册实现使用TORCH_LIBRARY_IMPL注册实现。对 NPU 执行而言torch_npu使用PrivateUse1dispatch key关于PrivateUse1的详细介绍请参考 PyTorch 官方文档 https://docs.pytorch.org/tutorials/advanced/privateuseone.htmlTORCH_LIBRARY_IMPL(npu, PrivateUse1, m) { m.impl(my_add, TORCH_FN(run_add_custom)); }3. 构建与运行本示例依赖 PTO Tile Lib、PyTorch、torch_npu与 CANN。请参考torch_npu官方安装指南https://gitcode.com/ascend/pytorch#%E5%AE%89%E8%A3%85或执行python3 -m pip install -r requirements.txt3.1 设置目标 SoC编辑auto_mode/demos/baseline/add/CMakeLists.txt把SOC_VERSION设置为目标芯片例如 A2/A3 使用Ascend910B1set(SOC_VERSION Ascendxxxyy CACHE STRING system on chip type)可在目标机器上执行npu_smi info查询芯片名称并按AscendChip Name的形式填写。3.2 构建 wheel设置 PTO Tile Lib 路径并构建 wheelexport ASCEND_HOME_PATH/usr/local/Ascend/ source /usr/local/Ascend/ascend-toolkit/set_env.sh export PTO_LIB_PATH[YOUR_PATH]/pto-isa rm -rf build op_extension.egg-info python3 setup.py bdist_wheel3.3 安装 wheelcd dist pip uninstall *.whl pip install *.whl3.4 运行测试cd test python3 test.py【免费下载链接】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），仅供参考