CANN PTO-ISA TPREFETCH_ASYNC 指令详解基于 SDMA CMO 的 L2 Cache 异步预取【免费下载链接】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-isaTPREFETCH_ASYNC是 CANN PTO-ISA 提供的一条面向计算侧的缓存预取指令它通过 SDMA CMOCache Maintenance Operationopcode6将 Global MemoryGM/HBM中的数据异步预热到 NPU 片上 L2 Cache使后续TLOAD直接命中 L2 而非回 GM 取数同时不占用数据对应的 UB 空间。本文基于 TPREFETCH_ASYNC 指令文档 为主线结合仓库中公开头文件、SDMA 后端实现与 NPU 单卡测试用例完整讲解该指令的接口语义、内部实现链路、约束条件与实战写法读者可据此在自己的 PTO kernel 中正确接入 L2 预取以优化访存性能。指令概述与定位TPREFETCH_ASYNC是一条逻辑上的内存访问/缓存提示类指令源数据必须位于 Global MemoryGM/HBM 地址空间预取目的地是片上 L2 Cache数据本身不进入 UB因此不消耗数据规模的 UB 空间。它内部经由 SDMA CMO 路径提交硬件请求但公开 API 位于pto命名空间与同步预取指令pto::TPREFETCH并列。与TPUT_ASYNC/TGET_ASYNC类似TPREFETCH_ASYNC依赖 SDMA 异步基础设施workspace、AsyncSession、AsyncEvent其数据流可表示为GM / HBM --(SDMA CMO prefetch)-- L2 Cache | -- 后续 TLOAD 命中 L2快速路径从实现上看见 TPrefetchAsyncImpl.hpp 头部注释该指令恰好以 AI Core 提交 SDMA CMO SQE 的方式实现自身因此依赖支撑TPUT_ASYNC/TGET_ASYNC的同一套 SDMA 基础设施实现本身是架构中立的A2A3 与 A5 的差异全部封装在 SDMA 后端头文件中共享实现只定义一份由薄封装头include/pto/npu/a2a3/TPrefetchAsync.hpp与include/pto/npu/a5/TPrefetchAsync.hpp引入用户看到的 API 形态即pto::下的内存访问指令。C 内建接口指令的公开声明位于 include/pto/common/pto_instr.hpp公共包含头为pto/pto-inst.hppnamespace pto { template typename GlobalData, typename... WaitEvents PTO_INST comm::AsyncEvent TPREFETCH_ASYNC(GlobalData src, PrefetchAsyncContext ctx, WaitEvents ... events); } // namespace pto公开包装器在pto_instr.hpp中先执行detail::PtoWaitEvents(events...)等待所有可选事件再转发到TPREFETCH_ASYNC_IMPL。该包装器受编译开关保护#if (defined(__CCE_AICORE__) || defined(__CPU_SIM)) !defined(__COSTMODEL) !defined(PTO_COMM_NOT_SUPPORTED)即仅在 AI Core 设备编译或 CPU 模拟__CPU_SIM场景下可见成本模型__COSTMODEL或显式关闭通信指令PTO_COMM_NOT_SUPPORTED时该 API 不参与编译。PrefetchAsyncContext上下文与 Session 语义PrefetchAsyncContext保存由 Host 侧SdmaWorkspaceManager::Init初始化后的 SDMA workspace 指针。其基类定义见 TPrefetchAsyncImpl.hppstruct PrefetchAsyncContextBase { __gm__ uint8_t* workspace{nullptr}; comm::AsyncSession session; comm::AsyncSession* externalSession{nullptr}; // 构造仅 workspace / workspace 外部 session AICORE comm::AsyncSession GetSession() { return externalSession ! nullptr ? *externalSession : session; } };要点如下未指定外部 Session 时Context 持有自身内部 Session。首次调用TPREFETCH_ASYNC时会懒初始化该 Session见下文实现链路后续复用。外部 Session 模式当TPREFETCH_ASYNC与TGET_ASYNC或TPUT_ASYNC共用同一 Channel Group 时必须复用它们的 Session通过PrefetchAsyncContext(workspace, sharedSession)构造以保证 SQE 提交队列一致。Event 等待语义等待返回的预取 Eventevt.Wait(ctx.GetSession())也会完成该共享 Session 中此前所有 SDMA 操作。生命周期外部 Session 的生命周期必须覆盖 Context 及相关异步 Event 的使用阶段并发使用的独立 Context 必须使用不同的 workspace或保证串行执行。内部开销手动模式非__PTO_AUTO__下PrefetchAsyncContext还持有一个 256 字节的 UB scratch tileTileTileType::Vec, uint8_t, 1, UB_ALIGN_SIZE用于构造 SDMA SQE 元数据这与数据规模无关。参数参数类型说明srcGlobalData需要预取到 L2 的 GlobalTensor 区域须位于 GM/HBMctxPrefetchAsyncContext计算侧预取上下文包含 workspace以及内部 Session 或共享外部 Sessionevents...WaitEvents...可选同步事件调用前先等待返回值返回comm::AsyncEvent用于跟踪异步预取完成状态。后续TLOAD依赖预取结果时调用evt.Wait(ctx.GetSession())等待完成在 CPU 模拟后端返回的空事件Wait/Test恒为true见 include/pto/cpu/TPrefetchAsync.hpp。底层实现链路源码级从公开包装器到硬件 SQE 提交指令的调用链为pto::TPREFETCH_ASYNC (pto_instr.hpp) └─ TPREFETCH_ASYNC_IMPL (TPrefetchAsyncImpl.hpp) ├─ 懒初始化 SessionTASSIGN_IMPL(scratchTile, 0) InitPrefetchAsyncSession └─ TPrefetchAsyncSdmaImpl └─ __sdma_cmo_prefetch (sdma_cmo_intrin.hpp) └─ SdmaCmoPrefetch → SubmitCmoPrefetchSqes → AddOneCmoSqe (opcode6)核心实现逻辑TPREFETCH_ASYNC_IMPLTPrefetchAsyncImpl.hpp的核心逻辑取ctx.GetSession()若 Session 无效则先对 scratch tile 执行TASSIGN_IMPL(ctx.scratchTile, 0x0)再调用InitPrefetchAsyncSession构建内部 Session——该函数以单条 SQE 块大小kSingleSqeBlockBytes 64MB、syncId0、自动 Channel Group 索引kAutoChannelGroupIdx为参数调用BuildSdmaSessionSession 引擎非 SDMA 时返回空事件否则进入TPrefetchAsyncSdmaImpl。TPrefetchAsyncSdmaImplTPrefetchAsyncImpl.hpp包含三层守卫srcGlobalData.data() nullptr→ 返回空事件非平坦连续一维布局 → 返回空事件。平坦连续判定TPrefetchAsyncIsFlatContiguous1D要求p41 p3dim4 p2dim3*p3 p1dim2*p2 p0dim1*p1且dim0..dim3均为 1即单行紧凑排布总字节数为 0 → 返回空事件。总字节数由五维 shape 乘积乘以元素类型大小计算。通过守卫后__sdma_cmo_prefetch(src, totalBytes, sdmaSession)提交硬件请求并把sdmaSession.runtimeCtx回写到 Session 以便后续等待。SQE 构造opcode6 的 CMO 请求sdma_cmo_intrin.hpp 定义constexpr uint32_t kCmoPrefetchOpcode 6U;。AddOneCmoSqeL35-L87向 STARS 提交队列写入一条 SDMA 类型的 SQE其关键字段通用字段type RT_STARS_SQE_TYPE_SDMA、opcode 6、sssv/dssv/sns/dns 1、源地址高低 32 位写入srcAddrLow/High目的地址为 0CMO 无写目的地A5 分支wrCqe 1、numBlocks 0、rtStreamId channelInfo-stream_id、taskId、kernelCredit K_CREDIT_TIME_DEFAULT非 A5 分支A2A3 等blockDim 0、kernel_credit、ptr_mode 0、ie2 0、qos 6、partid 63U、linkType 0等字段。提交过程SubmitCmoPrefetchSqesL90-L110按config.iter_num迭代将总字节数按block_bytes分块轮转分布到queue_num个队列逐条写入并推进sqTailSdmaCmoPrefetch则在BeginSdmaPost中完成地址换算与状态准备提交后由FinishSdmaPost产出AsyncEvent。值得注意的是头文件注释说明该模板包装是为了延迟代码生成避免经pto-inst.hpp传递包含时在每个翻译单元膨胀 IR 并触发 Bisheng 优化器问题——这也解释了为何指令实现采用模板化封装。Host 侧 workspace 初始化SDMA workspace 必须在 kernel 启动前由 Host 侧初始化。SdmaWorkspaceManager::Initinclude/pto/comm/async/sdma/sdma_workspace_manager.hpp的流程LoadDynamicSymbols从libruntime.so动态解析rtStreamGetSqid、rtStreamGetCqid、rtGetDeviceInfo从libopapi.so解析aclnnShmemSdmaStarsQuery(GetWorkspaceSize)CreateStarsStreams(kSdmaMaxChan)通过aclrtCreateStreamWithConfig(..., ACL_STREAM_DEVICE_USE_ONLY)创建 STARS 流逐个采集stream_id/sq_id/cq_id/logic_cq_id等 64 字节的HostStreamInfoMallocWorkspaceaclrtMalloc分配并清零设备侧 workspaceCopyOpResToDevice将流信息表拷到设备侧LaunchAicpuKernel在 AICPU 流上执行aclnnShmemSdmaStarsQuery把硬件 SQ 地址、寄存器基址、队列深度等写入 workspace。初始化完成后GetWorkspaceAddr()L149返回的指针即可作为__gm__ uint8_t *workspace传入 kernel最终转发给BuildSdmaSession。该头文件是 Host-only 头包含#error守卫禁止在设备代码__CCE_KT_TEST__中引入。约束与使用注意源数据必须位于 Global MemoryGM/HBM 地址空间GlobalTensor 必须是平坦连续的一维布局TPrefetchAsyncIsFlatContiguous1D守卫会静默返回空事件SDMA workspace 需要在 kernel 启动前由 Host 侧初始化并传入 kernelPrefetchAsyncContext内部持有 256 字节 UB scratch tile 和AsyncSession用于构造 SDMA 元数据并等待事件完成与TGET_ASYNC或TPUT_ASYNC共用 Channel Group 时必须复用其 Session外部 Session 的生命周期必须覆盖 Context 及相关异步 Event 的使用阶段并发使用的独立 Context 必须采用不同的 workspace或保证串行执行SDMA CMO 按 cache line 粒度工作非对齐范围由硬件处理CPU 模拟后端中该指令为空操作返回空AsyncEventWait/Test恒真。此外从 TPrefetchAsyncImpl.hpp 可知自动模式__PTO_AUTO__下该指令同样是 no-op 桩CCE 的 tile_size 类型系统无法从Tile::data()提取裸__ubuf__指针且自动模式调度使手动预取无必要因此提供可编译但不做任何事的实现以保持 API 兼容。与 TPREFETCH 的对比TPREFETCH是同步预取指令详见 TPREFETCH 指令文档两者的差异决定了适用场景维度pto::TPREFETCHpto::TPREFETCH_ASYNC数据流GM → UBGM → L2 Cache硬件路径MTEcopy_gm_to_ubufSDMA CMOopcode6UB 占用需要目标 Tile数据不占用 UB仅内部使用 256B scratch同步方式同步流水线屏障异步AsyncEvent典型用途小数据预取到 UB大数据或跨阶段数据预热到 L2简言之TPREFETCH把数据真正搬进 UB 供后续计算直接使用适合小数据块的同步预载TPREFETCH_ASYNC只做 L2 预热适合大数据量、跨阶段的数据提前加载配合TLOAD命中 L2 获得收益且不挤占宝贵的 UB 容量。编程示例基本用法以下示例预取一段 16384 个half元素的 GM 数据到 L2等待完成后执行TLOAD#include pto/pto-inst.hpp using namespace pto; __global__ AICORE void my_kernel(__gm__ half *src, __gm__ half *dst, __gm__ uint8_t *workspace) { using GShape Shape1, 1, 1, 1, 16384; using GStride Stride1, 1, 1, 1, 1; GlobalTensorhalf, GShape, GStride srcGlobal(src); PrefetchAsyncContext ctx(workspace); auto evt TPREFETCH_ASYNC(srcGlobal, ctx); evt.Wait(ctx.GetSession()); using TileData TileTileType::Vec, half, 128, 128, BLayout::RowMajor; TileData tile; TASSIGN(tile, 0x100); TLOAD(tile, srcGlobal); }注意workspace需在 Host 侧通过SdmaWorkspaceManager::Init()初始化后以参数传入 kernel等待完成后再执行依赖预取结果的TLOAD才能保证命中 L2。复用外部 Session当预取与TGET_ASYNC/TPUT_ASYNC共用 Channel Group 时复用外部构造的sharedSessionPrefetchAsyncContext ctx(workspace, sharedSession); auto getEvt TGET_ASYNC(dstGlobal, srcGlobal, sharedSession); auto prefetchEvt TPREFETCH_ASYNC(prefetchGlobal, ctx); auto putEvt TPUT_ASYNC(remoteGlobal, localGlobal, sharedSession); (void)putEvt.Wait(ctx.GetSession());等待putEvt共享 Session 中最后的操作即可顺带完成此前所有 SDMA 操作包括预取。测试用例佐证仓库在 A5 与 A2A3 上均有对应的单卡 ST 用例tests/npu/a5/src/st/testcase/tprefetch_async/tprefetch_async_kernel.cppA2A3 版本见 tests/npu/a2a3/src/st/testcase/tprefetch_async/tprefetch_async_kernel.cpp。正确性 kernelL64-L98以postCount次循环连续发起预取支持waitEachEvent逐次等待或仅等待最后事件两种模式useExternalSession分支演示了先用TASSIGN(ctx.scratchTile, 0x0)初始化 scratch、再BuildAsyncSession(ctx.scratchTile, sdmaWorkspace, sharedSession)构建外部 Session 的完整流程等待完成后经CopyViaTile用TLOAD/TSTORE分块搬运并校验输出。L2 收益 benchmark kernelPTO_TPREFETCH_ASYNC_L2_BENEFIT_ST宏控制L128-L159通过SYS_CNT系统计数器分别测量 4096 个float冷读直接TLOAD与预取后读的周期数成对比较并输出cold_avg_us与prefetched_avg_us用真实硬件计时验证 L2 预热收益。平台行为差异小结场景行为A5 / A2A3__CCE_AICORE__手动模式完整提交 SDMA CMO SQE异步返回AsyncEvent自动模式__PTO_AUTO__no-op 桩返回空事件CPU 模拟__CPU_SIMno-op 桩AsyncEvent::Wait/Test恒真成本模型 /PTO_COMM_NOT_SUPPORTEDAPI 不参与编译总结TPREFETCH_ASYNC以零 UB 数据开销 异步等待的方式把 GM/HBM 数据预热进 L2 Cache是 PTO-ISA 中面向大块数据跨阶段访存优化的关键指令。理解其PrefetchAsyncContext的 Session 复用规则、平坦连续一维布局约束、Host 侧SdmaWorkspaceManager初始化前提以及等待共享 Session 完成即完成全部 SDMA 操作的事件语义是在真实 kernel 中正确使用它的前提。需要深入了解实现细节的读者可继续阅读 TPrefetchAsyncImpl.hpp、sdma_cmo_intrin.hpp、sdma_workspace_manager.hpp 及对应平台的 ST 测试用例。【免费下载链接】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），仅供参考