ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

CANN PTO 异步 L2 预取指令 TPREFETCH_ASYNC 深度解析:基于 SDMA CMO 的缓存预热机制

CANN PTO 异步 L2 预取指令 TPREFETCH_ASYNC 深度解析:基于 SDMA CMO 的缓存预热机制 CANN PTO 异步 L2 预取指令 TPREFETCH_ASYNC 深度解析基于 SDMA CMO 的缓存预热机制【免费下载链接】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本文聚焦 Ascend CANN 的 Parallel Tile OperationPTO虚拟指令集中的TPREFETCH_ASYNC它如何通过 SDMA CMOCache Maintenance Operation, opcode6将 Global Memory 数据异步预热进 L2 Cache从而让后续TLOAD命中热缓存。文中将结合 docs/isa/TPREFETCH_ASYNC_zh.md 的接口规范与 include/pto/comm/async/sdma/TPrefetchAsyncImpl.hpp、include/pto/comm/async/sdma/sdma_cmo_intrin.hpp 等源码实现完整覆盖接口签名、上下文与 Session 管理、约束条件、与同步TPREFETCH的差异以及可复现的示例代码。设计动机为什么需要只预热 L2的异步预取在 NPU 上Global MemoryGM/HBM的访问延迟远高于片上存储。典型的高性能算子会先用 MTE 将数据搬入片上缓冲区UB再计算但存在两类场景让搬到 UB显得不划算数据体量远超 UB 容量一个大张量无法一次性进 UB若逐块搬运每次TLOAD都要经历一次完整的内存访问延迟跨阶段数据复用同一份 Global 数据会被多个阶段、多个核反复读取每次从 GM 重新拉取会重复支付延迟。TPREFETCH_ASYNC针对这两类场景提供了第三条路径异步地把数据从 GM/HBM预热到 L2 Cache但不占任何 UB 空间。后续依赖该数据的TLOAD可以直接从 L2 命中命中后 MTE 的搬运延迟显著低于冷读 GM。它本质上是一个面向计算侧AI Core 上执行的 kernel 内的内存访问 / 缓存提示接口这一点从仓库源码的注释可以直接印证This instruction is logically amemory accessinstruction (it stages data from GM/HBM into the on-chip L2 cache so that subsequent TLOADs hit warm lines). —— include/pto/comm/async/sdma/TPrefetchAsyncImpl.hpp实现上它恰好通过向 SDMA 引擎提交一条 CMO SQE 完成因此它复用了支撑TPUT_ASYNC/TGET_ASYNC的整套 SDMA 基础设施但公开 API 仍位于pto命名空间对用户呈现为一条普通的内存访问指令。数据流与工作原理TPREFETCH_ASYNC的数据流如下GM / HBM --(SDMA CMO prefetch)-- L2 Cache | -- 后续 TLOAD 命中 L2SDMA CMO 的 opcode 为6在源码中以常量形式定义constexpr uint32_t kCmoPrefetchOpcode 6U;见 include/pto/comm/async/sdma/sdma_cmo_intrin.hpp。CMO 按cache line 粒度工作非对齐范围由硬件处理因此调用方无需做手工对齐但也不应依赖它做小于一个 cache line 的精细控制。底层的 SQE 构造A5 与 A2A3 两套字段AddOneCmoSqe负责把一条预取请求编码成 SDMA SQE并按架构分支填写不同字段sdma_cmo_intrin.hppA5 路径sqe-opcode kCmoPrefetchOpcode同时设置sssv/dssv/sns/dns为 1、wrCqe 1写入lengthMove与源地址低/高位srcAddrLow/srcAddrHigh目的地址字段清零A2A3 路径同样opcode kCmoPrefetchOpcode但使用kernel_credit、length、qos 6、partid 63U等字段。两条路径的dstAddrLow/dstAddrHigh均为 0因为目标不是某段显存而是片上 L2 Cache 本身。提交完成后会执行pipe_barrier(PIPE_ALL)保证后续操作与 SQE 提交之间的流水线顺序。大数据的自动分块提交SubmitCmoPrefetchSqes会把一次预取按config.block_bytes拆成多条 SQE循环写入多个 channelqueueIdx idx % config.queue_num最后一条按余量修正传输字节数sdma_cmo_intrin.hpp。这正是它在源码中被称为面向大数据量跨阶段预热的原因单条 CMO SQE 无法承载任意大的范围分块 多队列轮转让大区域预取也能高效执行。完整调用链为TPREFETCH_ASYNC (公共 API) - TPREFETCH_ASYNC_IMPL // pto/comm/async/sdma/TPrefetchAsyncImpl.hpp - detail::TPrefetchAsyncSdmaImpl // 校验数据指针、平坦性、总字节数 - __sdma_cmo_prefetch // pto/comm/async/sdma/sdma_cmo_intrin.hpp - detail::SdmaCmoPrefetch // BeginSdmaPost - SubmitCmoPrefetchSqes - FinishSdmaPostC 内建接口TPREFETCH_ASYNC的公开声明位于 include/pto/common/pto_instr.hpp公共包含头为pto/pto-inst.hpp内部实现头为pto/common/pto_instr.hpp。接口签名如下namespace pto { template typename GlobalData, typename... WaitEvents, std::enable_if_tall_events_vWaitEvents..., int 0 PTO_INST comm::AsyncEvent TPREFETCH_ASYNC(GlobalData srcGlobalData, PrefetchAsyncContext ctx, WaitEvents... events); } // namespace pto入口处先执行detail::PtoWaitEvents(events...)等待可选同步事件再转发到TPREFETCH_ASYNC_IMPL。该接口仅在(__CCE_AICORE__ || __CPU_SIM) !__COSTMODEL !PTO_COMM_NOT_SUPPORTED条件下编译可见。参数说明参数类型说明srcGlobalData需要预取到 L2 的 GlobalTensor 区域ctxPrefetchAsyncContext计算侧预取上下文包含 workspace 以及内部 Session 或共享外部 Sessionevents...WaitEvents...可选同步事件需满足all_events_vWaitEvents...约束返回值返回comm::AsyncEvent用于跟踪异步预取完成状态。后续TLOAD依赖预取结果时调用evt.Wait(ctx.GetSession())等待完成也可以先批量发出多次预取最后统一等待测试用例正是这样做的见下文。PrefetchAsyncContextworkspace、Session 与 256-Byte scratchPrefetchAsyncContext是本次预取的计算侧上下文其结构在 include/pto/comm/async/sdma/TPrefetchAsyncImpl.hpp 中定义。它由两部分拼装基类PrefetchAsyncContextBase持有三个核心成员__gm__ uint8_t* workspace由 Host 侧SdmaWorkspaceManager::Init初始化后传入 kernel 的 SDMA workspace 指针comm::AsyncSession sessionContext 内部自持的 Sessioncomm::AsyncSession* externalSession可选的外部共享 Session。GetSession()的实现为externalSession ! nullptr ? *externalSession : session即指定外部 Session 时优先使用外部 Session否则使用内部 Session。派生类PrefetchAsyncContext额外持有ScratchTile其类型为using ScratchTile pto::Tilepto::TileType::Vec, uint8_t, 1, comm::sdma::UB_ALIGN_SIZE;即一个 256-Byte 的 UB scratch tile用于构造 SDMA 元数据描述符/SQE这也是文档约束中内部持有 256-Byte UB scratch tile的出处。Context 持有自身 Session 时的懒初始化当未指定外部 Session 时Context 的session.valid初始为 false。TPREFETCH_ASYNC_IMPL在第一次调用时完成如下初始化TPrefetchAsyncImpl.hpp对ctx.scratchTile执行TASSIGN_IMPL(ctx.scratchTile, 0x0)把 scratch tile 绑定到 UB 地址调用detail::InitPrefetchAsyncSession以kSingleSqeBlockBytes 64 * 1024 * 102464 MB作为单块大小、SdmaBaseConfig{64MB, 0, 1}、自动 Channel Group 索引kAutoChannelGroupIdx调用BuildSdmaSession构建 SDMA Session若 workspace 为空或构建失败返回valid false的 Session 与空AsyncEvent后续调用直接复用已缓存的 Session。SdmaWorkspaceManager位于 include/pto/comm/async/sdma/sdma_workspace_manager.hpp提供 Host 侧的Init()/Finalize()生命周期管理且禁止拷贝构造/赋值单例式用法。SDMA workspace 必须在 kernel 启动前由 Host 侧初始化并传入 kernel这是使用本指令的前置条件。平坦连续一维布局校验在真正提交预取之前TPrefetchAsyncSdmaImpl会做三重快速校验TPrefetchAsyncImpl.hppsrcGlobalData.data() nullptr—— 空指针直接返回空事件TPrefetchAsyncIsFlatContiguous1D—— 校验布局。判断逻辑为p4 1 p3 dim4 p2 dim3 * p3 p1 dim2 * p2 p0 dim1 * p1紧凑打包布局且dim0..dim3 1单行两者同时满足才认为是一段平坦连续的一维区域TPrefetchAsyncGetTotalBytes计算的总字节数为 0。任一不满足即返回空AsyncEvent不提交任何 SQE而不是报错这保证了接口的容错性。与 TPREFETCH 的对比维度pto::TPREFETCHpto::TPREFETCH_ASYNC数据流GM 到 UBGM 到 L2 Cache硬件路径MTEcopy_gm_to_ubufSDMA CMOopcode6UB 占用需要目标 Tile数据不占用 UB仅内部使用 scratch同步方式同步异步AsyncEvent典型用途小数据预取到 UB大数据或跨阶段数据预热到 L2两者在公共头中的声明相邻TPREFETCH见 pto_instr.hpp但语义完全不同TPREFETCH是同步的 MTE 搬运必须显式提供目标 TileTPREFETCH_ASYNC是异步的 L2 预热不占用数据对应的 UB 空间只消耗 256-Byte scratch。实践中可将二者组合少量热点数据用TPREFETCH提前进 UB大块数据用TPREFETCH_ASYNC预热 L2。约束与使用注意事项基于 docs/isa/TPREFETCH_ASYNC_zh.md 与源码实现使用时必须遵守以下约束源数据必须位于 Global MemoryGM/HBM且GlobalTensor必须是平坦连续的一维布局源码会对打包布局与单行条件做运行时校验SDMA workspace 需要在 kernel 启动前由 Host 侧初始化SdmaWorkspaceManager::Init并作为__gm__指针传入 kernelPrefetchAsyncContext内部持有 256-Byte UB scratch tile 和AsyncSession用于构造 SDMA 元数据并等待事件完成与TGET_ASYNC或TPUT_ASYNC共用 Channel Group 时必须复用其 Session。等待返回的预取 Event也会完成该共享 Session 中此前所有 SDMA 操作——这是由 Session 的串行提交语义决定的外部 Session 的生命周期必须覆盖 Context 及相关异步 Event 的使用阶段并发使用的独立 Context 必须采用不同的 workspace或保证串行执行SDMA CMO 按 cache line 粒度工作非对齐范围由硬件处理CPU simulation 后端中该指令为空操作no-op返回空AsyncEvent见 include/pto/cpu/TPrefetchAsync.hppAuto 模式__PTO_AUTO__下同样提供 no-op stub因 CCE tile_size 类型系统无法从Tile::data()提取裸__ubuf__指针且自动调度使手动预取失去意义见 TPrefetchAsyncImpl.hpp预取本身不保证数据驻留L2 是容量有限的缓存是否命中取决于后续访问时机与缓存替换策略因此evt.Wait(ctx.GetSession())之后仍需通过正常的TLOAD取数。示例基本用法以下示例来自 docs/isa/TPREFETCH_ASYNC_zh.md演示一次典型的异步预热 同步装载流程#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); }要点PrefetchAsyncContext ctx(workspace)使用 Context 自持 Session 的最简形式evt.Wait(ctx.GetSession())保证预取 SQE 全部完成后再发TLOAD从而确保TLOAD命中 L2TLOAD仍走正常的 MTE 取数路径预取只是把冷读 GM变成命中 L2。示例复用外部 Session与 TGET_ASYNC / TPUT_ASYNC 混用当同一 Channel Group 上既有异步通信TGET_ASYNC/TPUT_ASYNC又有异步预取时必须复用 Session否则会破坏通道上的串行顺序。以下示例假设sharedSession已针对所选 Channel Group 在外部完成构造PrefetchAsyncContext 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 中此前排队的getEvt与prefetchEvt因此只需等待最后一个事件即可。测试验证源码中的实证仓库为TPREFETCH_ASYNC提供了完整的 ST系统测试用例位于tests/npu/a5/src/st/testcase/tprefetch_async/tprefetch_async_kernel.cpptests/npu/a2a3/src/st/testcase/tprefetch_async/tprefetch_async_kernel.cpp测试 kernelTPrefetchAsyncCorrectnessKernel覆盖了关键行为使用动态 Shape/StrideShapeDYNAMIC,...构造GlobalTensor通过useExternalSession开关验证内部 Session 与外部 Session 两种模式外部模式先TASSIGN(ctx.scratchTile, 0x0)再BuildAsyncSession(ctx.scratchTile, sdmaWorkspace, sharedSession)完成 Session 构建通过postCount/waitEachEvent开关验证批量发事件 最后统一等待的语义循环发TPREFETCH_ASYNC结束后lastEvent.Wait(ctx.GetSession())预取等待完成后执行CopyViaTile逐块TLOAD/TSTORE校验数据正确性并在越界时pipe_barrier(PIPE_ALL)提前结束文件还定义了PTO_TPREFETCH_ASYNC_L2_BENEFIT_ST对应 L2 收益类用例对比预取与不预取的耗时表现。测试覆盖了文档约束中外部 Session 复用批量等待workspace 传入等全部关键路径是理解本指令语义的最佳参考实现。平台支持与适用前提后端行为A5NPU完整实现走 SDMA CMOopcode6见 include/pto/npu/a5/TPrefetchAsync.hpp薄封装复用共享实现A2A3NPU完整实现同样走共享实现见 include/pto/npu/a2a3/TPrefetchAsync.hppCPU 仿真no-opAsyncEvent::Wait/Test恒为 true见 include/pto/cpu/TPrefetchAsync.hppAuto 模式__PTO_AUTO__no-op stub编译通过但不执行实际预取Costmodel__COSTMODEL公共 API 不参与编译!defined(__COSTMODEL)条件以当前仓库为准该指令面向 AI Core 侧__CCE_AICORE__手动模式 kernel 使用。将TPREFETCH_ASYNC与合理分块、批量事件等待、外部 Session 复用相结合可以让大张量跨阶段计算从反复冷读 GM变为一次异步预热、多次 L2 命中是 PTO 手动编程中值得优先尝试的缓存优化手段。【免费下载链接】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),仅供参考
返回列表