TGET¶
简介¶
远程读操作:将远端 NPU 的数据读取到本地内存。数据通过 UB Tile 作为中间暂存缓冲区进行传输。
当 GlobalTensor 超出 UB Tile 容量时,TGET 将自动执行二维滑动——沿行(DIM_3)和列(DIM_4)分块以适配 Tile,并遍历所有外层维度(DIM_0、DIM_1、DIM_2)。
数学语义¶
对有效区域内每个元素 (i, j):
\[\mathrm{dst}^{\mathrm{local}}_{i,j} = \mathrm{src}^{\mathrm{remote}}_{i,j}\]
数据流:srcGlobalData(远端 GM) → stagingTileData(UB) → dstGlobalData(本地 GM)
汇编语法¶
PTO-AS 形式:参见 PTO-AS 规范。
同步形式:
tget %dst_local, %src_remote : (!pto.memref<...>, !pto.memref<...>)
降级时会为 GM→UB→GM 数据路径引入 UB 暂存 Tile;C++ 内建接口需要显式传入 stagingTileData(或 pingTile / pongTile)操作数。
C++ 内建接口¶
声明于 include/pto/comm/pto_comm_inst.hpp
单 Tile(自动分块)¶
template <typename GlobalDstData, typename GlobalSrcData, typename TileData, typename... WaitEvents>
PTO_INST RecordEvent TGET(GlobalDstData &dstGlobalData, GlobalSrcData &srcGlobalData,
TileData &stagingTileData, WaitEvents&... events);
乒乓双缓冲¶
使用两个暂存 Tile,将相邻块的 TLOAD 与 TSTORE 重叠执行,隐藏 DMA 传输延迟。
template <typename GlobalDstData, typename GlobalSrcData, typename TileData, typename... WaitEvents>
PTO_INST RecordEvent TGET(GlobalDstData &dstGlobalData, GlobalSrcData &srcGlobalData,
TileData &pingTile, TileData &pongTile, WaitEvents&... events);
约束¶
- 类型约束:
GlobalSrcData::RawDType必须等于GlobalDstData::RawDType。TileData::DType必须等于GlobalSrcData::RawDType。GlobalSrcData::layout必须等于GlobalDstData::layout。
- 内存约束:
srcGlobalData必须指向远端地址(源 NPU)。dstGlobalData必须指向本地地址(当前 NPU)。stagingTileData/pingTile/pongTile必须预先在统一缓冲区中分配。
- 有效区域:
- 传输大小由
GlobalTensor的形状决定(自动分块以适配 Tile)。
- 传输大小由
- 乒乓约束:
pingTile和pongTile必须具有相同的类型和维度。- 必须位于不重叠的 UB 偏移处。
示例¶
基础用法¶
#include <pto/comm/pto_comm_inst.hpp>
#include <pto/pto-inst.hpp>
using namespace pto;
template <typename T>
void example_tget(__gm__ T* local_data, __gm__ T* remote_addr) {
using TileT = Tile<TileType::Vec, T, 16, 16>;
using GShape = Shape<1, 1, 1, 16, 16>;
using GStride = BaseShape2D<T, 16, 16, Layout::ND>;
using GTensor = GlobalTensor<T, GShape, GStride, Layout::ND>;
GTensor srcG(remote_addr);
GTensor dstG(local_data);
TileT stagingTile;
TASSIGN(stagingTile, 0);
// 基础远程读
comm::TGET(dstG, srcG, stagingTile);
}
乒乓双缓冲¶
constexpr size_t tileUBBytes = ((64 * 64 * sizeof(float) + 1023) / 1024) * 1024;
TileT pingTile(64, 64);
TileT pongTile(64, 64);
TASSIGN(pingTile, 0);
TASSIGN(pongTile, tileUBBytes); // 不重叠的 UB 区域
// 将 TLOAD[i+1] 与 TSTORE[i] 重叠执行以提升流水线利用率
comm::TGET(dstG, srcG, pingTile, pongTile);