PTO 实现 GEMMAllReduce 融合算子A2/A3 双流重叠实战指南【免费下载链接】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 仓库中的 kernels/manual/a2a3/gemm_ar 示例为主线讲解如何在 A2/A3 芯片上使用 PTOParallel Tile Operation指令集实现多卡GEMM AllReduce融合算子。该示例通过计算流 通信流双流设计与逐 tile 级别的信号量握手让 GEMM 计算与跨卡 AllReduce 通信并发执行在参考平台上获得约1.443x端到端加速。读完本文你将掌握该算子的整体架构、计算/通信内核的关键实现、无锁就绪队列与 RDMA 窗口编程模型以及构建、运行、调参与性能解读的完整实战方法。功能概述与算子规格示例实现的多卡融合语义为$$ C_{final} \sum_{i0}^{nranks-1} A_i \times B $$其中A_i为各 rank 私有的M x K矩阵B为所有 rank 共享的K x N矩阵C_i是本地 GEMM 结果C_final是 AllReduce 之后的最终M x N输出。默认矩阵配置定义在 gemm_ar_config.h 中M5416、K6144、N1408。算子规格如下ItemValueOpTypeGEMM AllReduceInputA_i:M x K,float16,ND每 rank 私有B:K x N,float16,DN共享OutputC_final:M x N,float16,NDAllReduce 结果Compute kernel nameGemmComputeKernelCube 架构dav-c220-cubeComm kernel nameGemmCommAllKernelVector 架构dav-c220-vec目录布局kernels/manual/a2a3/gemm_ar/ ├── CMakeLists.txt # 构建配置3 个 targetcube 内核、vec 内核、host 可执行程序 ├── run.sh # 一键构建运行脚本自动计算 HCCL_BUFFSIZE 并定位 MPI ├── gemm_ar_config.h # 全局配置矩阵形状、tile 尺寸、block 数量 ├── main.cpp # 入口MPI 初始化、数据生成、HCCL 初始化、窗口分配、性能测量、验证 ├── gemm_compute_kernel.cpp # GEMM 计算内核Cube 侧L0C FP32 - GM FP16 自动转换 ├── comm_kernel.cpp # 通信内核Vector 侧单内核内重叠 RS/AG 的 AllReduce ├── kernel_launchers.h # host 侧内核 launcher 声明 ├── common.hpp # 设备侧 HcclRemotePtr 封装RDMA 窗口地址转换 ├── comm_context.h # HcclDeviceContext 结构每个 rank 的 RDMA 窗口地址 ├── ready_queue.hpp # 多 block 无锁 tile 队列计算 - 通信信令 └── comm_mpi.h # MPI 动态加载封装dlopen/dlsym无硬链接依赖说明README 目录清单中标注的hccl_context.h在当前仓库中实际命名为 comm_context.h其中定义了CommDeviceContext内含rankId、rankNum、winSize以及每个 rank 的windowsIn[HCCL_MAX_RANK_NUM]/windowsOut[]RDMA 窗口地址数组。总体架构双流通信-计算重叠示例采用双流设计Compute Stream 上运行GemmComputeKernel24 个 AIC 块Comm Stream 上运行GemmCommAllKernel24 个 AIV 块。PTO 通信指令直接操作 HCCL RDMA 窗口完成 AllReduce计算与通信之间通过逐 tile 信号量进行流水线交接而非设备级全局屏障。┌──────────────────────────────────────────────────────────────────────────────┐ │ Compute Stream (24 AIC) Comm Stream (24 AIV) │ │ │ │ GemmComputeKernel: GemmCommAllKernel: │ │ ┌─────────────────────────┐ ┌──────────────────────────────┐ │ │ │ for each tile: │ │ RS/AG overlap loop │ │ │ │ K-loop (L1 - L0 - Cube) │ poll Ready Queue │ │ │ │ TSTORE - gemm_output │──Ready──→ │ TLOAD tile from gemm_output│ │ │ │ pipe_barrier(ALL) │ Queue │ TSTOREAtomicAdd - owner │ │ │ │ Enqueue tile_idx │ │ subtile-ready / summary │ │ │ └─────────────────────────┘ │ drain ready subtiles for AG│ │ │ │ TLOAD - TSTORE to remote │ │ │ │ ready-driven AG handoff │ │ │ │ subtile-level overlap │ │ │ │ │ │ │ └──────────────────────────────┘ │ └──────────────────────────────────────────────────────────────────────────────┘从源码角度看main.cpp中分别通过aclrtCreateStream创建computeStream与commStreamhost 侧LaunchCompute调用launchGemmCompute见 gemm_compute_kernel.cpp 底部LaunchComm调用launchGemmCommAll见 comm_kernel.cpp 底部launcher 声明统一放在 kernel_launchers.h 中。计算内核实现细节计算内核的时间流水线可概括为Time - L1 (MTE2): [TLOAD A0,B0] [TLOAD A1,B1] ... L0 (MTE1): [TEXTRACT k0] [k1] [k2] [k3] [TEXTRACT k0] ... Cube (M): [TMATMUL k0] [ACC k1] [ACC k2] [ACC k3] [TMATMUL k0] ... ^ full three-stage overlap ^每个 AIC 按block_idx x tiles_per_block方式领取 tile 子集对每个 tile 依次执行Block Swizzle 映射SwizzleTileIndex()把线性 tile 索引映射为之字形遍历顺序奇数行反转使相邻 tile 在 L1 中复用矩阵B的列K 循环每stepKa4次迭代做一次批量TLOAD将 4 个 K 片拉入 L1之后每次用TEXTRACT取出一个 K 片到 L0再做TMATMUL/TMATMUL_ACC累加见ProcessKIteration()TSTOREL0C 中的 FP32 结果经 FixPipe 自动转成 FP16 后写回gemm_outputpipe_barrier(PIPE_ALL)确保 GM 写完成MultiBlockEnqueueFast将tile_idx入队通知通信内核。源码中相关常量与断言位于 gemm_compute_kernel.cpp 的G_K_LOOP、G_STEP_KA/G_STEP_KB 4处static_assert强制G_K_LOOP可被stepK整除、stepKa stepKb且G_BASE_K固定为 64。tile 类型别名TileMatAData、TileMatBData、ResTileT分别对应 L1 面板、L0 切片与 L0C 累加器。通信内核实现细节通信内核GemmCommAllImpl()见 comm_kernel.cpp把 RSReduceScatter生产与 AGAllGather消费交织在同一个循环里同步点是逐 subtile 计数器而不是设备级屏障。RS 生产者路径每个通信 block 认领的队列子集为queues(block b) { b, b num_comm_blocks, b 2*num_comm_blocks, ... }默认COMPUTE_BLOCK_NUM COMM_BLOCK_NUM 24时退化为 1:1通信 block 更少时一个 block 通过RsPollQueues()/RsWaitOnQueue()轮询多个计算队列。对每个出队的 tile沿M方向把 tile 拆成G_COMM_SUBTILES_PER_TILE G_BASE_M / G_COMM_SUB_M个定高 subtileRsPipelineStep()使用 ping/pong 两块 UB tile使当前 subtile 的TLOAD与上一 subtile 的TSTOREAtomicAdd重叠RS 目的地是 owner rankowner tile_idx % nranks归约直接写进该 rank 的reduced_output。RS/AG 重叠同步重叠协议使用 owner rank 的signal_matrix中的两类计数器subtile-ready[local_subtile_id]统计有多少 rank 已对某个 owner-local subtile 完成 RSag-summary[summary_block]面向负责该 subtile 的 AG block 的粗粒度唤醒门铃。发布侧遵循RsPublishSubtileReady()先pipe_barrier(PIPE_ALL)冲刷本地流水线再dsb(DSB_DDR)保证reduced_output写全局可见随后RsNotifySubtileReady()递增 owner-local 就绪计数、RsNotifyAgSummary()递增由AgSummaryBlockForSubtile()选出的 AG 唤醒计数两者底层均为pto::comm::TNOTIFY(..., NotifyOp::AtomicAdd)。消费侧遵循AgDrainReadyAssignedSubtiles()用TTEST(..., nranks, GE)探测本 block 认领的subtile-ready计数一次 drain pass 首次命中时只执行一次 acquire fencepipe_barrier dsb即可消费该批全部就绪 subtile将每个就绪 subtile 搬运到所有远端 rank若无法推进则AgWaitAssignedSummary()阻塞在summary_ack_count 1等待下一次分配的唤醒。AgSummaryBlockForSubtile()采用反转条带reversed-stripe映射让 AG 重的 block 落在 RS 轻的 block 上从而拉平rs_work ag_work的组合负载。AG 执行器路径AG 工作在 owner-local subtile 空间分配total_local_subtiles my_tile_count * G_COMM_SUBTILES_PER_TILE assigned_ids(block b) { num_comm_blocks - 1 - b k*num_comm_blocks }对每个就绪的 assigned subtileAgDecodeLocalSubtile()把 owner-local subtile id 映射回reduced_output的全局行偏移AgTransferSubtileToAll()恰好向每个远端 rank 广播G_COMM_SUB_M行首个远端 peer 按local_subtile_id % (nranks - 1)轮转避免所有 block 同时锤击同一目的地。该设计使 AG 可以在某个 owner-local subtile 被所有 rank 完全归约后立即开始无需等待整卡 RS 结束。就绪队列机制┌─────────────┐ ┌─────────────┐ │ AIC 0 │ │ AIV 0 │ │ (Compute) │──Queue──│ (Comm) │ │ block_idx0│ 0 │ block_idx0│ └─────────────┘ └─────────────┘ ┌─────────────┐ ┌─────────────┐ │ AIC 1 │ │ AIV 1 │ │ (Compute) │──Queue──│ (Comm) │ │ block_idx1│ 1 │ block_idx1│ └─────────────┘ └─────────────┘ ... ... ┌─────────────┐ ┌─────────────┐ │ AIC 23 │ │ AIV 23 │ │ (Compute) │──Queue──│ (Comm) │ │ block_idx23│ 23 │ block_idx23│ └─────────────┘ └─────────────┘实现位于 ready_queue.hpp每个队列是 64 字节对齐的PerBlockQueue含count生产者侧单调递增计数与data[]tile 索引数组槽位通过GetQueueSlot()辅助函数寻址不依赖隐式的data[idx]布局假设生产者AICPerBlockQueueEnqueueFast通过GetQueueSlot()写目标槽再递增count并用dcci刷新缓存状态使条目对 AIV 可见由调用方跟踪槽位可将 dcci 次数从 5 次降到 2 次消费者AIVPerBlockQueueTryDequeue用硬件TTEST检查count head1经GetQueueSlot()刷新目标槽后返回 tile id无就绪数据返回-1长时间空闲后退化为硬件TWAIT当COMM_BLOCK_NUM COMPUTE_BLOCK_NUM时一个通信 block 按轮询顺序排空多个队列队列分配是静态的因此无需跨 block 原子仲裁设计为单生产者-单消费者队列内部无需任何原子操作。内存布局与 HCCL 窗口只有被远端TPUT或TNOTIFY写入的缓冲区才必须放在 HCCL RDMA 窗口内仅本地读写使用的缓冲区可用普通aclrtMalloc分配BufferSizeLocationWhyreduced_outputM x N x 2BHCCL windowRSAtomicAdd与 AG 远端TPUT写入FP16signal_matrixG_SIGNAL_TOTAL_SLOTS x 4B按 64B 对齐HCCL windowsubtile-ready 与 AG-summary 计数外加保留的旧版 barrier 槽位gemm_outputM x N x 2BaclrtMalloc仅本地读写FP16src0_dev,src1_dev输入矩阵FP16aclrtMalloc仅本地读写main.cpp中AllocDeviceBuffers通过WindowAlloc(windowBase, ...)从hostCtx.windowsIn[rankId]起步分配reduced_output与signal_matrix并校验winOffset winSize否则报 HCCL window too small。窗口大小由环境变量HCCL_BUFFSIZE控制run.sh按 padded 后的reduced_output占用自动计算并附加较大余量pad(M, G_BASE_M) x pad(N, G_BASE_N) x 2 / 1MB 64MBsignal_matrix位于同一窗口内相对64MB余量可忽略。关于跨 rank 寻址common.hpp 的CommRemotePtr(ctx, localPtr, pe)以本 rank 窗口基址计算偏移再叠加到目标 rank 的windowsIn[pe]上得到远端地址设备上下文由 comm_context.h 的CommDeviceContext承载。拓扑适配方面main.cpp在 MESH 拓扑A2下直接复用 HCCL 返回的上下文结构在 RING 拓扑A3下通过hccl_compat兼容结构从CommOpResParam的remoteRes提取远端 RDMA 地址并手工组装 host 上下文。性能数据参考值以下数据来自当前subtile-ready / AG-summary overlap实现采集环境为 8 卡 Ascend 910BM5416, K6144, N1408pad 为5504 x 1536258 tiles (43 x 6)compute_blocks24、comm_blocks24。每个 rank 计算完整 GEMMC_i A_i x BAllReduce 对 8 个C_i求和。MetricValueCompute-only368.1 us254546 GFLOPSSequential808.9 uscompute371.6 us comm437.3 us 63.6 GB/sPipelined560.6 uscompute done367.2 uscomm done560.6 us 49.7 GB/sSpeedup1.443xTime saved248.4 us30.7%Overlap eff66.8%Throughput1337307 GFLOPStotalmain.cpp中PrintTimingDetails给出了上述指标的完整计算口径Speedup Sequential / PipelinedTime saved seq_avg - pipe_avgOverlap eff (seq_comp seq_comm - pipe) / min(seq_comp, seq_comm)。指标解读Compute-only纯 GEMM 无通信的执行时间反映单卡 Cube 利用率上限Sequential先计算后通信、无重叠的路径808.9 us371.6 us计算 437.3 us通信Pipelined双流并发路径560.6 us相对 Sequential 为1.443x加速重叠效率66.8%Overlap efficiency较短阶段被重叠隐藏的比例66.8%表示约三分之二的短阶段被隐藏。优化历史下表为历史优化检查点最后一行是当前subtile-ready / AG-summary overlap路径的最新端到端结果。较早行应视为背景参考而非当前路径的字面分解。OptimizationPipelined (us)GainConclusionBaseline808--Block Swizzle793-1.8%KeptRSAtomicAdd移除独立 Reduce 阶段736-6.6%KeptAG 行级扁平化调度623-15.4%历史检查点48 AIVRSskip AGparticipate639RS 仅 24 AIVAG 48 AIVRevertedAIC干扰48 AIV 双队列1 AIC : 2 AIV667RS 与 AG 均在 48 AIVRevertedAIC干扰当前subtile-ready / AG-summary overlap路径560.6较623 us历史检查点约-10.0%当前结果性能调优指南1. 优先做多核划分每个 AIC 按block_idx x tiles_per_block领取 tile 子集块间互不干扰。检查清单调节COMPUTE_BLOCK_NUM使每个 block 分到相近数量的 tile不同矩阵形状下按G_NUM_TILES (M_padded/128) x (N_padded/256)重算总 tile 数。配套源码gemm_ar_config.h 中的GEMM_AR_BLOCK_TILE_COUNT/GEMM_AR_BLOCK_START_TILE宏实现了余数摊给前remainder个 block的均衡划分避免经典 ceil 划分产生1 个矮块 N-1 个满块的负载不均。2. 选择合适的 base tileL0A 与 L0B 采用 ping/pong 双缓冲每个缓冲限制为 32 KiB。对于 FP16 输入2 bytes/elemL0A tile 字节 ~baseM x baseK x 2128 x 64 x 2 16 KiBL0B tile 字节 ~baseK x baseN x 264 x 256 x 2 32 KiB通信 tile 尺寸为baseM x baseN x sizeof(FP16) 128 x 256 x 2 64 KB3. 用 L1stepK缓存提升复用stepKastepKb4时一次TLOAD将 4 个 K 片带入 L1之后逐个TEXTRACT拉入 L0。L1 占用2 x 64KB (A) 2 x 128KB (B) 384KB 1024KB增大stepK可减少 DMA 启动开销但总量必须仍能装进 L1。4. 保住流水线重叠性能关键来自两件事的组合计算内核内部双缓冲L1/L0A/L0B计算与通信之间的双流重叠。当观察到communication time compute time计算侧已高效应聚焦改善通信或提升重叠度compute time communication time通信已被完全隐藏应聚焦计算侧。5. 调节通信 block 数量COMM_BLOCK_NUM控制通信内核的 AIV 并行度可通过--comm-blocks调整。在 Ascend 910B 上将COMM_BLOCK_NUM从 24 提升到 48 会使 AIC 计算时间显著增加约24%原因是 HBM 带宽争用与 TSCH 调度开销因此默认稳定取 24。6. 约束条件K必须能被G_BASE_K x G_STEP_KA整除默认64 x 4 256M自动 pad 到 128 的倍数N自动 pad 到 256 的倍数所有 HCCL 窗口缓冲区必须在每个 rank 上以相同偏移分配每次迭代前必须用aclrtMemset重置signal_matrix见main.cpp中ResetDeviceState。构建与运行配置 Ascend CANN 环境export ASCEND_CANN_PATH/usr/local/Ascend/cann-version/set_env.sh source ${ASCEND_CANN_PATH}激活提供 Python 与 NumPy 的 conda 环境conda activate your-conda-env以 8 rank 运行示例cd ${git_clone_path}/kernels/manual/a2a3/gemm_ar ./run.sh --nranks 8 --soc-version Ascend910B1指定起始设备编号FIRST_DEVICE0 ./run.sh --nranks 8 --soc-version Ascend910B1使用自定义计算/通信 block 数量./run.sh --nranks 8 --compute-blocks 20 --comm-blocks 4运行成功时程序打印GEMM AllReduce demo completed successfully.脚本 run.sh 内部会按getopt解析-r/--run-mode、-v/--soc-version、-n/--nranks、-d/--ndevices、--compute-blocks、--comm-blocks等参数先清理遗留 HCCL 共享内存状态rm -rf /dev/shm/sem.hccl*; ipcrm -a按PAD_M * PAD_N * 2 / 1024 / 1024 64MB自动抬升HCCL_BUFFSIZE再调用cmake与make最后执行mpirun -n ${NRANKS} ./gemm_allreduce --first-device ${FIRST_DEVICE}。环境变量Environment VariablePurposeDefault BehaviorASCEND_CANN_PATHCANNset_env.sh脚本的完整路径自动 glob/usr/local/Ascend/cann-*/set_env.sh并取最新MPI_SEARCH_DIRSMPIbin/目录搜索路径空格分隔搜索/usr/local/mpich/bin、/home/mpich/bin等常见位置ASCEND_DRIVER_PATHCMake 使用的 Ascend 驱动路径默认/usr/local/Ascend/driverMPI_LIB_PATH运行时动态加载的libmpi.so绝对路径由run.sh按发现的 MPI 安装自动设置HCCL_BUFFSIZEHCCL RDMA 窗口大小MB由run.sh根据 paddedM/N占用自动计算FIRST_DEVICE起始 NPU 设备编号默认0修改矩阵维度更新 gemm_ar_config.h 中的CONFIG_G_M、CONFIG_G_K、CONFIG_G_N注意宏带#ifndef保护源码通过 include 共享配置。也可以直接从 CMake 传入cmake -DCONFIG_G_M8192 -DCONFIG_G_K8192 -DCONFIG_G_N2048 ..约束K必须能被G_BASE_K x G_STEP_KA整除默认64 x 4 256HCCL_BUFFSIZE由run.sh自动计算。此外 gemm_ar_config.h 还支持CONFIG_G_BASE_M、CONFIG_G_BASE_K、CONFIG_G_BASE_N、CONFIG_COMPUTE_BLOCK_NUM、CONFIG_COMM_BLOCK_NUM、CONFIG_COMM_SUB_M等编译期参数并带有static_assert如G_COMM_SUB_M必须整除G_BASE_M、AG summary 槽至少占一个 cache line 等作为配置安全网。FAQProblemCause and FixHCCL window too small窗口必须覆盖 paddedreduced_output占用加上signal_matrix。检查是否手动覆盖了HCCL_BUFFSIZErun.sh会自动按pad(M) x pad(N) x 2 / 1MB 64MB抬升HcclGetRootInfo failed: 7上次运行的遗留脏状态。执行rm -rf /dev/shm/sem.hccl*; ipcrm -a或等待约 30 秒重试HCCL 初始化后挂起通常是 rank 同步问题。检查所有 rank 是否都到达了CommMpiBarrier通信内核段错误通常是窗口地址非法。验证windowsIn[]条目非零信号等待死锁或 AG 停滞signal_matrix未在迭代间清零或 subtile-ready / AG-summary 归属映射错误。检查resetState是否对signal_matrix调用了memset校验时max_diff过大FP16 精度有限。校验容差为atol1.0, rtol0.01见main.cpp的VerifyOutput。若 diff 异常大检查 subtile-ready / AG-summary 同步与 owner 映射aclInit repeat init100002无害。代码已对单进程重复aclInit做了防护--allow-run-as-root失败本项目使用 MPICH该选项是 OpenMPI 专有构建系统CMakeLists.txt 定义了三个 targetCompilerbishengCANN 内置 clang 15.0.5Cube 内核 flags--cce-aicore-archdav-c220-cube -DMEMORY_BASEgemm_compute_kernel链接runtime--cce-fatobj-linkVector 内核 flags--cce-aicore-archdav-c220-vec -DMEMORY_BASEcomm_kernelHost 可执行程序标准-xc编译gemm_allreduce链接gemm_compute_kernel、comm_kernel、runtime/ascendcl/hcomm/m/tiling_api等pto-comm-isa的 include 路径必须放在最前以覆盖 CANN 自带的pto_tile.hppinclude_directories中${PROJECT_SOURCE_DIR}/../../../../include即仓库根include/目录CMAKE_CCE_COMPILE_OPTIONS中开启--cce-pto-enable并设置了 aicore 栈大小、溢出记录、地址变换等 LLVM 选项RUN_MODE为sim时链接runtime_camodel以支持模拟器。变更日志DateChange2025-12-15初始版本双流 GEMM AllReduce 融合2026-04-01适配 CANN 9.0.0移除废弃的hccl/hccl.h依赖2026-04-02RSAtomicAdd移除独立 Reduce 阶段AG 扁平化改善负载均衡2026-04-21通信模式由RS - DeviceBarrier - AG改为subtile-ready / AG-summary overlap延伸阅读该示例依赖的 PTO 通信指令说明可参考 docs/isa/comm 目录及 docs/PTO-Virtual-ISA-Manual_zh.md本仓库更多手动算子示例位于 kernels/manual/a2a3 与 kernels/manual/a5算子融合与多核流水线的基础知识见 docs/coding/operator-fusion_zh.md 与 docs/coding/multi-core-programming_zh.md。【免费下载链接】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），仅供参考