CANN PTO-ISA TREDUCE 指令详解多 NPU 跨卡归约Reduce的原理、约束与工程实践【免费下载链接】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-isaTREDUCE 是 CANN PTO-ISA 通信指令集中的归约Reduce指令由根rootNPU 从多个远端 rank 收集数据并在本地执行逐元素归约常用于分布式训练中的梯度求和、最大值/最小值聚合等场景。本文基于 docs/isa/comm/TREDUCE.md 展开并结合include/pto/comm/下的指令声明与实现源码、tests/npu/下的 ST 测试用例深入讲解 TREDUCE 的数学语义、汇编与 C 两种编程形态、AIV/CCU 双引擎路径、大 Tensor 自动分块2D sliding chunked机制及全部使用约束读完即可在自己的 PTO kernel 中正确编写与验证 TREDUCE 归约逻辑。一、指令语义根节点本地归约TREDUCE 的语义是根节点从多个远端 NPU 拉取数据在本地完成逐元素归约结果写入根节点本地 GM。用一句话概括dst[local] reduce(src[0], src[1], ..., src[N-1])。1.1 数学定义对有效区域valid region内的每个元素(i, j)TREDUCE 执行如下归约$$ \mathrm{dst}^{\mathrm{local}}{i,j} \bigoplus{r0}^{N-1} \mathrm{src}^{(r)}_{i,j} $$其中 $N$ 是 rank 数量即ParallelGroup的nranks$\oplus$ 是归约操作求和 Sum、取最大 Max、取最小 Min 等。1.2 执行角色划分只有根节点执行指令TREDUCE 与 TGATHER、TBROADCAST、TSCATTER 等 PTO 集合通信指令一致遵循根节点驱动模型只有根root需要执行TREDUCE根的身份由parallelGroup.GetRootIdx()标识非根 rank只需保证其源缓冲区source buffer在操作期间已就绪ready且保持有效valid不需要执行任何指令在非根 rank 上调用TREDUCE属于未定义行为undefined behavior开发者应通过if (my_rank root) { ... }之类的判断避免。从源码看这一约定在 include/pto/comm/pto_comm_inst.hpp 的注释中有明确说明Only the root needs to execute. Non-root ranks ensure source buffers are ready.只有根需要执行非根 rank 只需确保源缓冲区就绪。1.3 支持的操作类型ReduceOp归约操作类型定义在 include/pto/comm/comm_types.hpp 中enum class ReduceOp : uint8_t { Sum 0, // Element-wise sum Max 1, // Element-wise maximum Min 2, // Element-wise minimum };在 AIV 路径下三种操作分别落到底层的TADD求和、TMAX取最大、TMIN取最小向量指令具体映射见 include/pto/comm/a2a3/TReduce.hpp 中的ReduceTiles()case ReduceOp::Sum: TADD(acc, acc, recv); break; case ReduceOp::Max: TMAX(acc, acc, recv); break; case ReduceOp::Min: TMIN(acc, acc, recv); break; default: PTO_ASSERT(false, TREDUCE: unknown ReduceOp); break;这也解释了归约发生在本地每个远端 rank 的数据块先被TLOAD到接收 tile再通过一条向量运算指令累加进累加 tile全程只在根节点的 UBUnified Buffer中操作。二、两种编程形态汇编语法与 C Intrinsic2.1 汇编语法同步形式PTO ISA 层面的同步归约指令写法如下treduce %group, %dst {op #pto.reduce_opSum} : (!pto.group..., !pto.memref...) treduce %group, %dst {op #pto.reduce_opMax} : (!pto.group..., !pto.memref...)%group由ParallelGroup描述的 rank 组封装了tensors数组、nranks、rootIdx%dst根节点本地的目标 GMmemrefop归约操作属性支持#pto.reduce_opSum、#pto.reduce_opMax、#pto.reduce_opMin。需要注意的是汇编级指令在 lowering降级过程中会引入内部的累加 tile 与接收 tile 来构建归约流水线而 C intrinsic 需要显式提供这些 tile 操作数。2.2 C Intrinsic 声明TREDUCE 的 C 接口声明于 include/pto/comm/pto_comm_inst.hpp共有两个重载基础形态累加 tile 接收 tiletemplate CollEngine engine CollEngine::AIV, typename ParallelGroupType, typename GlobalDstData, typename TileData, typename... Args PTO_INST RecordEvent TREDUCE(ParallelGroupType parallelGroup, GlobalDstData dstGlobalData, TileData accTileData, TileData recvTileData, ReduceOp op, Args... args);乒乓双缓冲形态累加 tile ping tile pong tiletemplate CollEngine engine CollEngine::AIV, typename ParallelGroupType, typename GlobalDstData, typename TileData, typename... Args PTO_INST RecordEvent TREDUCE(ParallelGroupType parallelGroup, GlobalDstData dstGlobalData, TileData accTileData, TileData pingTileData, TileData pongTileData, ReduceOp op, Args... args);两个重载均返回RecordEvent便于与 PTO 的事件同步机制配合。第一个重载使用单接收 tile 串行处理各个远端 rank第二个重载使用 ping-pong 双 tile让下一个远端 rank 的TLOAD与当前 tile 的归约计算重叠掩盖远程数据搬运的延迟。模板参数engine选择后端引擎详见第四节取值为 include/pto/comm/comm_types.hpp 中定义的enum class CollEngine : uint8_t { AIV 0, CCU 1, };CollEngine::AIV默认基于 tile 的路径即TLOAD 向量归约 TSTORECollEngine::CCU仅 Ascend950NPU_ARCH 3510可用AIV 触发 CKE gate归约数据通路在 CCU 硬件上执行。当engine CollEngine::CCU时第一个变参Args...必须是CcuTriggerContext其中包含 CKE slot VA 与 gate mask详见第五节 CCU 路径。三、使用约束ConstraintsTREDUCE 对类型、内存与ParallelGroup有严格约束违反任一条件都会导致编译期static_assert失败或运行期PTO_ASSERT断言。3.1 类型约束ParallelGroup::value_type::RawDType必须等于GlobalDstData::RawDType即参与归约的每个 rank 源张量与目标张量的原始数据类型一致TileData::DType必须等于GlobalDstData::RawDType即 tile 元素类型与全局目标数据类型一致。这两条在源码中以static_assert硬编码实现见 include/pto/comm/a2a3/TReduce.hppstatic_assert(std::is_same_vT, typename GlobalDstData::RawDType, TREDUCE: GlobalData type mismatch!); static_assert(std::is_same_vT, typename TileData::DType, TREDUCE: TileData element type must match GlobalData element type);3.2 内存约束dstGlobalData必须指向本地地址当前 NPU 上的 GM归约结果只落在根节点accTileData、recvTileData或accTileData、pingTileData、pongTileData必须是预先分配的 UB tile不能使用未分配的内存远程 rank 的源缓冲区parallelGroup.tensors[r]指向的远端 GM需要在操作期间保持有效。3.3 ParallelGroup 约束parallelGroup.tensors[r]必须指向 rankr的源缓冲区对根节点而言是远端 GMparallelGroup.GetRootIdx()标识当前调用 NPU 为归约根节点所有源张量被假定具有相同的 shape 与 stridesTREDUCE 实现不会做逐 rank 的形状校验或对齐仅以根节点张量为参照refTensor parallelGroup[rootRank]驱动整个归约与分块过程ParallelGroup的大小必须大于 0rootIdx必须落在[0, nranks)区间内否则触发断言PTO_ASSERT见 include/pto/comm/a2a3/TReduce.hpp。3.4 分块模式约束Large Tile Support当 GlobalTensor 在行或列上超过单个 UB tile 容量时TREDUCE 会自动进入 2D sliding 分块模式此时额外施加以下约束若TileData具有静态ValidRow则GetShape(DIM_3)行维度必须能被ValidRow整除若需要支持行上的部分块partial row必须改用DYNAMICValidRow 的 Tile若TileData具有静态ValidCol则GetShape(DIM_4)列维度必须能被ValidCol整除若需要支持列上的部分块partial col必须改用DYNAMICValidCol 的 Tile。源码中的对应断言位于 include/pto/comm/a2a3/TReduce.hppif constexpr (!isDynamicRow) { PTO_ASSERT(gShape3 % chunkRows 0, TREDUCE chunked: shape3 must be divisible by tile ValidRow when ValidRow is static. Use a Tile with DYNAMIC ValidRow for partial row chunk support.); } if constexpr (!isDynamicCol) { PTO_ASSERT(gShape4 % chunkCols 0, TREDUCE chunked: shape4 must be divisible by tile ValidCol when ValidCol is static. Use a Tile with DYNAMIC ValidCol for partial column chunk support.); }四、大 Tensor 自动分块2D Sliding 归约流水TREDUCE 的 AIV 实现中内置了数据放得下就走简单路径放不下就自动分块的分派逻辑入口在 include/pto/comm/a2a3/TReduce.hpp 的TREDUCE_IMPL。4.1 路径分派规则实现先计算张量总行数totalRows gShape0 * gShape1 * gShape2 * gShape3并读取 tile 的有效行/列容量chunkRows accTileData.GetValidRow()、chunkCols accTileData.GetValidCol()若totalRows 0 || gShape4 0直接返回空数据短路对应测试中的 Empty Rows 用例若totalRows chunkRows gShape4 chunkCols进入TreduceSimple简单路径整块数据一次性TLOAD根节点数据、逐个TLOAD远端数据并归约、最后TSTORE否则进入TreduceChunkedSingle分块路径ping-pong 版本对应TreduceChunkedPingPong。4.2 分块遍历策略分块路径对五维 GlobalTensor 做如下切分源码注释与实现一致外层维度DIM_0、DIM_1、DIM_2显式迭代DIM_3行按tileValidRow大小切块DIM_4列按tileValidCol大小切块行、列切块时依据TileData的ValidRow/ValidCol是否为DYNAMIC设置当前块的RowMaskInternal/ColMaskInternal从而支持非整除的部分块见TreduceProcessChunkSingle中isDynamicRow/isDynamicCol分支。每个分块执行完整的归约流水TreduceProcessChunkSingleTLOAD根节点该块数据到accTileData对每个远端 rankTLOAD该块到recvTileData→ 用ReduceTiles归约进accTileData将归约完成的块TSTORE到dstGlobalData对应位置。4.3 简单路径的流水细节以TreduceSimpleinclude/pto/comm/a2a3/TReduce.hpp为例归约流水严格使用 PTO 事件同步set_flag/wait_flag协调 MTE2加载、向量单元 PIPE_V归约、MTE3存储三个流水级TLOAD(accTileData, parallelGroup[rootIdx]); set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); for (int r 0; r nranks; r) { if (r rootIdx) continue; TLOAD(recvTileData, parallelGroup[r]); set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); detail::ReduceTiles(accTileData, recvTileData, op); set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); } set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); TSTORE(dstGlobalData, accTileData);另外还有一个值得注意的特例nranks 1时 TREDUCE 退化为纯搬运——只有根节点自身的数据直接TLOAD后TSTORE不再经过向量归约见TreduceSimple开头的if (nranks 1)分支。4.4 乒乓双缓冲用三块 tile 隐藏远程加载延迟TREDUCE_IMPL的 ping-pong 版本include/pto/comm/a2a3/TReduce.hpp将远端归约改为双缓冲流水时间线如下源码注释中的示意图// 无乒乓 // [TLOAD remote0] - [Reduce] - [TLOAD remote1] - [Reduce] - ... // 有乒乓TLOAD[i1] 与 Reduce[i] 重叠 // [TLOAD remote0] - [Reduce remote0 | TLOAD remote1] - [Reduce remote1 | TLOAD remote2] - ...实现通过EVENT_ID1/EVENT_ID2两个事件交替标记 ping/pong tile 的加载完成状态在归约当前 tile 的同时预先下发下一个远端 rank 的TLOAD。块间chunk 间的存储同样用PIPE_V → PIPE_MTE3事件同步收尾。代价是需要在 UB 中额外多分配一块 tileacc ping pong 共三块且三块 tile 的 UB 地址必须不重叠。五、AIV 与 CCU 双引擎路径TREDUCE 通过模板参数engine选择实现路径两条路径在 include/pto/comm/pto_comm_inst.hpp 的TREDUCE外层模板中用if constexpr分派到TREDUCE_IMPL或TREDUCE_CCU_IMPL。5.1 AIV 路径默认AIV 路径即第四节的 tile 化实现数据搬运TLOAD/TSTORE与向量归约TADD/TMAX/TMIN全部在 AIVAI Vector 核心上完成。该路径跨 A2/A3 与 A5 平台共享——include/pto/comm/a5/TReduce.hpp 直接#include pto/comm/a2a3/TReduce.hpp复用 AIV 实现避免代码重复。5.2 CCU 路径仅 Ascend950CCU 路径的语义与 AIV 有本质差异AIV 路径只有根节点调用TREDUCECCU 路径所有 rank 都必须通过宿主侧HcclCcuKernelRegister/HcclCcuKernelLaunch注册并启动 CCU kernelAIV 侧的 trigger kernel 只负责敲门doorbell CKE gate真正的归约数据通路运行在 CCU 硬件引擎上。A5 的 CCU 实现位于 include/pto/comm/a5/TReduce.hpp核心只有一行CcuStoreTriggerSelf(parallelGroup, accTileData, ctx, events...)即把本 rank 的部分和partial result写入 CCU 输入 HBM 并触发 CKE gatedstGlobalData、recvTileData、op在 trigger kernel 侧均被忽略。CCU 路径的变参必须携带CcuTriggerContext其字段包括 CKE slot VA、gate mask、selfIdx以及输入来源CcuInputSource定义见 include/pto/comm/comm_types.hppenum class CcuInputSource : uint8_t { HostManaged 0, // 宿主已填充 parallelGroup[selfIdx]AIV 仅触发 CKE gate AivStored 1, // AIV 将本 rank 部分和放在 accTileDataCCU IMPL 先 TSTORE 再触发 };HostManaged宿主如aclrtMemcpy或前序算子已把输入 HBM 填好AIV trigger kernel 只 doorbell CKE gateaccTileData/recvTileData不被 CCU 实现消费AivStoredAIV trigger kernel 把本 rank 的部分和放在accTileData中CCU IMPL 将其TSTORE到parallelGroup[selfIdx]即宿主注册的输入 VA执行pipe_barrier(PIPE_MTE3)后再 doorbell——这使 AIV 前序计算阶段 → CCU 归约 → 输出可以全在设备侧完成无需宿主 memcpy。另外A2/A3 平台上TREDUCECollEngine::CCU是非法的a2a3 侧保留了一个带依赖static_assert的TREDUCE_CCU_IMPL桩函数一旦实例化即报错CCU engine is not available on A2/A3见 include/pto/comm/a2a3/TReduce.hpp。A5 侧则通过宏PTO_COMM_A5_TREDUCE_PROVIDED屏蔽该桩启用真正的 CCU 重载避免编译器在重载决议时误选通用变参版本。六、完整代码示例以下两个示例均来自原指令参考文档可直接复制到 PTO kernel 中使用。二者都通过__gm__ T* group_addrs[NRANKS]接收各 rank 的 GM 地址数组用comm::ParallelGroup封装后调用comm::TREDUCE。6.1 基本归约求和Reduce Sum#include pto/comm/pto_comm_inst.hpp using namespace pto; template typename T, int SIZE, int NRANKS void reduce_sum(__gm__ T* group_addrs[NRANKS], __gm__ T* result, int my_rank) { using TileT TileTileType::Vec, T, 1, SIZE; using GTensor GlobalTensorT, Shape1,1,1,1,SIZE, BaseShape2DT, 1, SIZE, Layout::ND, Layout::ND; // Stack-allocated tensors GTensor tensors[NRANKS]; for (int i 0; i NRANKS; i) { tensors[i] GTensor(group_addrs[i]); } comm::ParallelGroupGTensor group(tensors, NRANKS, my_rank); GTensor dstG(result); TileT accTile, recvTile; comm::TREDUCE(group, dstG, accTile, recvTile, comm::ReduceOp::Sum); }6.2 取最大归约Reduce Max#include pto/comm/pto_comm_inst.hpp using namespace pto; template typename T, int SIZE, int NRANKS void reduce_max(__gm__ T* group_addrs[NRANKS], __gm__ T* result, int my_rank) { using TileT TileTileType::Vec, T, 1, SIZE; using GTensor GlobalTensorT, Shape1,1,1,1,SIZE, BaseShape2DT, 1, SIZE, Layout::ND, Layout::ND; GTensor tensors[NRANKS]; for (int i 0; i NRANKS; i) { tensors[i] GTensor(group_addrs[i]); } comm::ParallelGroupGTensor group(tensors, NRANKS, my_rank); GTensor dstG(result); TileT accTile, recvTile; comm::TREDUCE(group, dstG, accTile, recvTile, comm::ReduceOp::Max); }代码要点TileT使用TileType::Vec向量 tileT, 1, SIZE表示单行SIZE列的 tile当SIZE超出 UB 容量时TREDUCE_IMPL会按第四节所述自动分块GTensor使用五维Shape1,1,1,1,SIZE与BaseShape2DT, 1, SIZE, Layout::ND即一个1×SIZE的 ND 布局全局张量ParallelGroup的构造参数依次为 tensor 数组、rank 数量、本 rank 的根索引my_rank。若要构造非根 rank 的 ParallelGroup通常将 rootIdx 显式传为根节点编号三个 tile或三块 tile 的 ping-pong 形态必须在调用前分配好TREDUCE内部不会替你分配 UB 资源。关于ParallelGroup的完整定义tensors、nranks、rootIdx字段及Create工厂函数可参见 docs/isa/comm/README.md其GetRootIdx()、GetSize()、下标访问operator[]被 TREDUCE 实现用于取根节点张量与遍历远端 rank。七、源码级验证ST 测试用例如何印证 TREDUCE 行为仓库在 CPU 模拟与 NPU 两个层级都提供了 TREDUCE 的 ST 测试可作为理解指令行为与验证实现正确性的第一手材料。7.1 A5 NPU 测试基础归约与期望值校验npu/a5/comm/st/testcase/treduce/treduce_kernel.cpp 中的TReduceKernelImpl展示了真实 kernel 的完整写法用pt::Shape/pt::Stride定义五维动态 shape 的GlobalTensor每个 rank 的输入地址通过CommRemotePtr(hcclCtx, input, i)构造远端指针填入ParallelGroup只有my_rank root的核执行pto::comm::TREDUCE(pg, outputG, accTile, recvTile, op)测试前先用TASSIGN(accTile, 0x0)初始化累加 tile用TASSIGN(recvTile, 0x10000)初始化接收 tile不同的初始值有助于暴露未覆盖元素被误用的问题。期望值在主机侧由ReduceExpected计算rankr的输入为base r * 100Sum 为逐项累加、Max/Min 为逐项取极值仅 root 校验结果。测试显式实例化了 float/int32_t 的 Sum、Max、Min 组合例如template bool RunReducefloat, 256, pto::comm::ReduceOp::Sum(int, int, int, int); template bool RunReduceint32_t, 256, pto::comm::ReduceOp::Max(int, int, int, int); template bool RunReduceint32_t, 256, pto::comm::ReduceOp::Min(int, int, int, int);7.2 大 shape 分块与乒乓测试同一文件中还包含三类针对性用例直接印证第四节的实现机制Large Shape Chunked 测试TReduceLargeShapeKernelImpltotal_rows tile_rows且要求total_rows % tile_rows 0静态 tile 的整除约束通过static_assert硬性保证实例化如int32_t, 128 行 × 32 列, tile 16 行即 8 个分块Ping-Pong 测试TReducePingPongKernelImpl使用三块 UB tileacc/ping/pong每块按((tile_rows * cols * sizeof(T) 1023) / 1024) * 1024字节对齐到 1KB 边界放置确保 UB 地址不重叠调用 6 参数的乒乓重载TREDUCE(pg, outputG, accTile, pingTile, pongTile, op)Empty Rows 测试TReduceEmptyKernelImplshape 行数为 0验证实现中的空数据短路路径不会写坏输出缓冲区输出保持-1哨兵值。7.3 CCU 路径测试npu/a5/comm/st/testcase/treduce_ccu/treduce_ccu_kernel.cpp 提供了 CCU 引擎的完整范例覆盖CcuInputSource的两种输入来源treduce_ccu_trigger_kernelHostManaged 路径宿主已填充输入 HBMAIV 只构造CcuTriggerContext{ckeVA, mask, selfIdx, CcuInputSource::HostManaged}并调用TREDUCECollEngine::CCU(...)触发 CKE gatetreduce_ccu_fused_kernelAivStored 路径AIV 先用TEXPANDS(accTile, fillValue)生成部分和再以CcuInputSource::AivStored调用 TREDUCE实现前序 AIV 计算 → CCU 归约 → 输出的设备侧融合两个 kernel 均通过extern C导出*_launch函数配合宿主侧HcclCcuKernelRegister/HcclCcuKernelLaunch完成 CCU 集合的注册与启动。7.4 CPU 模拟测试此外 tests/cpu/st/testcase/treduce/treduce_kernel.cpp 提供了 CPU 模拟环境下的 TREDUCE 测试配合 tests/cpu/st/testcase/treduce/main.cpp 与 tests/cpu/st/testcase/treduce/gen_data.py可在无 NPU 硬件的环境下先行验证归约逻辑是开发期快速迭代 TREDUCE 相关 kernel 的有效手段。八、常见误区与最佳实践小结综合指令文档与源码实现编写 TREDUCE kernel 时建议重点检查以下几点角色判断务必用if (my_rank root)包裹TREDUCE调用非根 rank 调用属于未定义行为CCU 路径则相反所有 rank 都要参与 CCU kernel 的注册与启动tile 资源accTileData/recvTileData或 acc/ping/pong 三块必须预分配且 UB 地址不重叠ping-pong 形态记得按块大小对齐分配避免乒乓 tile 相互覆盖类型一致ParallelGroup 的源张量、目标张量、tile 三者数据类型必须一致否则编译期static_assert直接报错分块整除性使用静态ValidRow/ValidCol的 tile 时DIM_3/DIM_4必须能整除 tile 容量遇到不整除的尾块需求改用DYNAMICValidRow/ValidCol 的 tile实现会通过RowMaskInternal/ColMaskInternal自动处理部分块形状一致性ParallelGroup 中所有 rank 的张量 shape 与 strides 必须一致实现以根节点张量为参照不会逐 rank 校验空数据行数或列数为 0 时 TREDUCE 是安全的空操作直接返回测试中也专门覆盖了这一场景。TREDUCE 是 PTO-ISA 集合通信指令集中多卡数据聚合的关键一环与 TGATHER根收集不归约、TSCATTER根分发、TBROADCAST根广播共同构成四类基本集合原语可参阅 docs/isa/comm/README.md 了解完整的 PTO 通信指令家族。理解其根节点驱动 本地归约 自动分块 双引擎的设计是编写高性能跨卡算子、乃至基于 PTO-Virtual-ISA-Manual 进行分布式算子开发的基础。【免费下载链接】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），仅供参考