CANN PTO 演示示例完全指南：从 CPU 模拟到 NPU 生产级算子

CANN PTO 演示示例完全指南从 CPU 模拟到 NPU 生产级算子【免费下载链接】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 仓库中 demos 目录下的全部演示示例覆盖生产级 PyTorch 算子baseline、跨平台 CPU 模拟cpu与即时编译torch_jit三大场景。读者可以据此掌握 PTO Tile Library 的三种典型用法如何在 Ascend NPU 上将自定义 PTO 内核封装为可通过torch.ops调用的 PyTorch 算子、如何在无 NPU 硬件的机器上完成算法原型设计与 CI 测试以及如何跳过 wheel 构建、以 JIT 方式快速验证内核正确性。演示全景三类场景定位demos 目录的核心价值在于用最小可运行示例覆盖 PTO Tile Library 的完整使用链路——从纯 C 内核编写到 Host 侧算子注册再到 Python 侧集成与验证。整个目录按用途而非算子类型划分为三个子目录demos/ ├── baseline/ # 生产级 PyTorch 算子示例NPU │ ├── add/ # 基础逐元素加法 │ ├── gemm_basic/ # 带流水线优化的 GEMM │ └── flash_atten/ # 带动态分块的 Flash Attention ├── cpu/ # CPU 模拟演示跨平台 │ ├── gemm_demo/ │ ├── flash_attention_demo/ │ └── mla_attention_demo/ └── torch_jit/ # PyTorch JIT 编译示例 ├── add/ ├── gemm/ └── flash_atten/此外仓库实际布局中 baseline/allgather_async 还提供了基于 PTO 通信指令的异步 AllGather 内核示例含 URMA 变体可作为 baseline 家族在集合通信方向的扩展参考。三类示例的选择逻辑如下表所示场景是否需要 NPU 硬件典型用途产物形态baseline是A2/A3/A5生产级算子开发、上线交付wheel 安装包 PyTorch 算子cpu否x86_64/AArch64算法原型、学习编程模型、CI/CD独立可执行程序torch_jit是快速原型、内核正确性验证、性能基准运行时编译的.soBaseline从 PTO 内核到 PyTorch 算子的完整流水线baseline 下的示例是生产级模板它们展示如何实现自定义 PTO 内核并通过torch_npu将其作为 PyTorch 算子公开包含从内核实现到 Python 集成的完整工作流并自带 CMake 构建系统和 wheel 打包setup.py。支持平台为 A2/A3/A5。add逐元素加法内核逐行拆解add 是最基础也最适合入门的示例。其目录结构本身就构成了一套算子工程模板demos/baseline/add/ ├── op_extension/ # Python 包入口模块加载 ├── csrc/ │ ├── kernel/ # PTO kernel 实现 │ └── host/ # Host 侧 PyTorch 算子注册 ├── test/ # 最小化 Python 测试 ├── CMakeLists.txt # 构建配置 ├── setup.py # Wheel 构建脚本 └── requirements.txt第 1 步实现 PTO kernel内核位于 csrc/kernel/add_custom.cpp。它使用#if __CCE_AICORE__ 220 defined(__DAV_C220_VEC__)限定 A2/A3 平台的向量核编译路径核心逻辑runTAdd展示了 PTO 编程模型的关键要素UB 双缓冲BUFFER_NUM 2定义 ping-pong buffer并显式为 x/y/z 三个张量在 UB 上分配 ping/pong 地址如X_PING 0x0、X_PONG 0x8000 0x100UB_SIZE 0x30000192KB对应 A2/A3 的 UB 容量GlobalTensor Tile 视图用pto::GlobalTensorT, ShapeDim5, StridDim5描述全局内存张量用TileTileType::Vec, T, tileSRows, tileSCols, BLayout::RowMajor, -1, -1定义 UB 上的 Tile后两个-1表示动态 mask多核分块BLOCK_DIM 20对应向量核数先按BLOCK_ROWS × BLOCK_COLS做核间分块再在每个核内做tileNum × BUFFER_NUM的核内分块static_assert保证分块不溢出 UB硬件流水线同步通过set_flag/wait_flag显式编排PIPE_MTE2GM→UB 搬入、PIPE_V向量计算与PIPE_MTE3UB→GM 搬出之间的依赖指令级编程循环体内部依次执行TASSIGN更新 GlobalTensor 地址偏移、TLOAD搬入、TADD向量加法、TSTORE搬出pingpong_flag在 0/1 之间翻转以复用 buffer。内核入口add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength)以extern C __global__ AICORE导出并以half类型、tileRows20, tileCols2048的固定 tile 形状调用runTAdd。要让内核进入构建需要在 CMakeLists.txt 中添加ascendc_library(no_workspace_kernel STATIC csrc/kernel/add_custom.cpp )第 2 步Host 侧注册 PyTorch 算子Host 侧实现在 csrc/host/my_add.cpp 与 csrc/host/utils.h分三步完成集成定义算子 schemaPyTorch 通过TORCH_LIBRARY_FRAGMENT(npu, m)在npu命名空间声明my_add随后 Python 侧即可用torch.ops.npu.my_add调用TORCH_LIBRARY_FRAGMENT(npu, m) { m.def(my_add(Tensor x, Tensor y) - Tensor); }实现算子run_add_custom引入构建系统自动生成的 kernel launch 头文件aclrtlaunch_add_custom.h分配输出张量后通过EXEC_KERNEL_CMD内部封装ACLRT_LAUNCH_KERNEL将内核入队到当前 NPU 流执行#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; }utils.h中EXEC_KERNEL_CMD宏的实现细节值得注意它用ConvertTypes将at::Tensor转为裸数据指针、通过c10_npu::getCurrentNPUStream()获取当前流再经at_npu::native::OpCommand::RunOpApi入队从而保证算子与 PyTorch 的流语义一致。注册实现NPU 执行走torch_npu的PrivateUse1dispatch keyTORCH_LIBRARY_IMPL(npu, PrivateUse1, m) { m.impl(my_add, TORCH_FN(ascendc_path::run_add_custom)); }第 3 步Python 加载与测试op_extension/__init__.py会在导入时调用 op_extension/_load.py 中的_load_opextension_so通过torch.ops.load_library加载 wheel 内置的lib/libop_extension.so。测试用例 test/test.py 构造[20, 2048]的 float16 张量将torch.ops.npu.my_add(x_npu, y_npu)的结果与torch.add的 CPU 参考结果做assertRtolEqual对比——这一NPU 结果 vs CPU 参考的验证模式贯穿仓库所有 baseline 示例。gemm_basicTiling 策略与双缓冲流水线gemm_basic 实现固定维度[m, k, n] [512, 2048, 1536]的 GEMMC A × BA 为 ND 布局 float16B 为 DN 布局 float16输出为 float ND内核名为gemm_basic_custom。其 README 提供了完整的 Tiling 参数表是学习多维分块的范本参数值说明m/k/n512 / 2048 / 1536全局问题规模singleCoreM/K/N128 / 2048 / 256单核分块24 核按4 × 6分组切分 m 与 nbaseM/baseK/baseN128 / 64 / 256每核 tile 超出 L0 容量后按 k64 的 base block 进一步切分在代码层面GEMM 示例通过TileShape2D/BaseShape2D/GlobalTensor的组合为 GM 上的 AND、BDN、CND建立带形状与步长的视图流水线方面它在 L1 与 L0 上使用双缓冲来重叠数据搬运与矩阵计算同步点包括正向依赖MTE2 - MTE1、MTE1 - MMAD、MMAD - FIXPIPE与反向依赖MTE1 - MTE2、MMAD - MTE1整体调度结构如下这一示例说明当单核 tile 超过 L0 容量时需要像再按 baseK 切分这样引入二级 tiling并通过显式流水线同步换取搬运与计算的重叠——这是 PTO 高性能 GEMM 内核的通用方法论。flash_atten动态分块flash_atten 展示带自动 tile 大小选择的 Flash Attention 实现同样以torch_npu算子形式交付工程结构、构建与测试方式与 add 示例一致可用于学习不规则/动态形状算子在 PTO 下的处理方式。通用构建流程SOC 与 wheel所有 baseline 示例共享同一套构建流程。要点包括设置目标 SoC编辑示例目录下的CMakeLists.txt将SOC_VERSION设为目标芯片例如set(SOC_VERSION ascend910b1 CACHE STRING system on chip type)可在目标机器上执行npu_smi info查询芯片名称并按AscendChip Name的形式填写A2/A3 平台对应Ascend910B1一类取值。配置环境并构建 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_wheel其中PTO_LIB_PATH指向本仓库根目录setup.py依赖torch_npu.utils.cpp_extension.NpuExtension会以 CMake 驱动内核编译并要求 CMake 3.18脚本优先探测cmake3同时依据torch.compiled_with_cxx11_abi()自动传递GLIBCXX_USE_CXX11_ABI开关。CMakeLists.txt会从ASCEND_CANN_PACKAGE_PATH或ASCEND_HOME_PATH定位ascendc_kernel_cmake并 includeascendc.cmake。安装并测试pip install dist/*.whl cd test python3 test.py依赖清单见各示例的 requirements.txt如torch2.8.0、torch-npu2.8.0.post2、numpy1.26.0、setuptools80.9.0等。CPU 模拟无硬件环境下的跨平台开发cpu 目录提供在 CPUx86_64/AArch64上运行的跨平台示例不依赖 Ascend 硬件适合算法原型设计、学习 PTO 编程模型以及 CI/CD 测试。包含三个演示基础 GEMMgemm_demo、Flash Attentionflash_attention_demo与多潜在注意力 MLAmla_attention_demo。每个 demo 是独立的 CMake 工程例如 gemm_demo/CMakeLists.txt 通过get_filename_component(PTO_TILE_LIB_REPO_ROOT ...)定位仓库根目录、以__CPU_SIM __PTO_AUTO__编译宏切换到 CPU 模拟实现并把 include 加入头文件搜索路径。编译标准方面GCC 14 使用 C23否则回退 C20。推荐通过仓库统一的驱动脚本 tests/run_cpu.py 一键构建并运行python3 tests/run_cpu.py --demo gemm --verbose python3 tests/run_cpu.py --demo flash_attn --verbose python3 tests/run_cpu.py --demo mla python3 tests/run_cpu.py --demo all # 依次运行全部三个 demo该脚本会自动完成编译器探测优先clang 15其次g 13也可用--cxx/--cc显式指定、CMake 工具链检查缺失时通过 pip 安装cmake3.16、demo 源码的 configure/build/run 全流程并打印perf:开头的性能输出行。CPU 模拟也因此天然适配无 Ascend 硬件的 CI 流水线。PyTorch JIT免 wheel 的即时编译原型torch_jit 展示将 C 内核即时编译并与 PyTorch 张量直接集成的用法适合快速原型设计无需预先构建 wheel。包含 JIT 加法add、JIT GEMMgemm与带基准测试套件的 JIT Flash Attentionflash_atten。以 add 为例运行方式为export PTO_LIB_PATH[YOUR_PATH]/pto-isa cd demos/torch_jit/add python add_compile_and_run.py其底层机制jit_util_add.py值得说明即时编译调用bisheng编译器以-xcce --npu-archdav-2201 -O2 -stdc17 -I${PTO_LIB_PATH}/include将add_custom.cpp编译为共享库默认 120 秒超时ctypes 直连通过ctypes.CDLL加载.so按内核签名声明argtypesblockDim、stream、x/y/z 数据指针、元素数 N并将torch.npu.current_stream()作为 stream 传入add_func最终以call_kernel(blockDim, stream, ...)形式执行张量级验证add_compile_and_run.py 在 NPU 上构造[20, 2048]的 float16 张量调用 JIT 内核后与x y做torch.testing.assert_close校验。JIT 方式把改内核 → 重编 → 验证的迭代周期压缩到一次 Python 执行内非常适合内核开发早期阶段的正确性调试需要性能基准时可参考 torch_jit/flash_atten 附带的 benchmark 脚本fa_benchmark.py与配套说明。前置要求Baseline 和 JITNPUAscend AI 处理器 A2/A3/A5910B/910C/950CANN Toolkit 8.5.0带torch_npu的 PyTorchPython 3.9.x、CMake 3.16CPU 演示支持 C23 的 C 编译器CMake 3.16Python 3.9.x可选备注Python 已宣布 3.7.x/3.8.x 进入 EOL 阶段CANN 即将停止对这两个版本的支持请升级到 Python 3.9.x 的版本。快速开始按推荐顺序第 1 步CPU 模拟——零硬件成本验证 PTO 编程模型与算法正确性python3 tests/run_cpu.py --demo gemm --verbose python3 tests/run_cpu.py --demo flash_attn --verbose第 2 步NPU JIT 示例——有 NPU 环境时快速验证内核export PTO_LIB_PATH[YOUR_PATH]/pto-isa cd demos/torch_jit/add python add_compile_and_run.py第 3 步NPU Baseline 生产级算子——完整走一遍内核 → 算子 → wheel → 测试cd demos/baseline/add python -m venv virEnv source virEnv/bin/activate pip install -r requirements.txt export PTO_LIB_PATH[YOUR_PATH]/pto-isa python3 setup.py bdist_wheel pip install dist/*.whl cd test python3 test.py延伸阅读入门指南docs/getting-started_zh.md编程教程docs/coding/tutorial_zh.mdISA 参考docs/isa/README_zh.md手工内核kernels/manual/README_zh.md自定义算子kernels/custom/README_zh.md测试用例tests/README_zh.md在此基础上可进一步阅读 docs/coding/ProgrammingModel_zh.md编程模型与 docs/coding/compilation-process_zh.md编译流程理解 PTO 的底层执行模型或参考 kernels/manual 与 tests/cpu 中更完整的工业级内核与回归测试用例。【免费下载链接】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），仅供参考