GlobalTensor 全局内存张量编程模型PTO 5 维 shape/stride 元数据与 TLOAD/TSTORE 数据搬运实战【免费下载链接】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-isapto::GlobalTensor是 CANN PTOParallel Tile Operation虚拟指令集中描述全局内存GM张量的核心抽象一个轻量的元数据包装由一个__gm__指针加上一组5 维 shape 与 stride描述构成被TLOAD、TSTORE、MGATHER、MSCATTER等内存类指令直接消费。本文以 docs/coding/GlobalTensor.md 为骨架结合仓库源码include/pto/common/pto_tile.hpp、include/pto/common/type.hpp、include/pto/common/pto_instr.hpp等与真实算子示例demos/baseline/add、tests/cpu/st讲解 GlobalTensor 的类型定义、静态/动态维度建模、布局 hint、2D 便捷 helper、地址绑定TASSIGN以及数据搬运指令的底层实现使读者能够独立编写、阅读和调试基于 GlobalTensor 的 PTO 算子代码。一、GlobalTensor 是什么GM 张量的轻量视图在 PTO 编程模型中计算单元是片上的Tile固定容量的 2 维片上缓冲区而数据源与数据归宿则是位于全局内存GM中的张量。pto::GlobalTensor就是后者的抽象它本身不持有数据只持有一个__gm__指针指向 GM 中数据的起始地址一个5 维的 shape 与 stride 描述。这份元数据会被内存类指令消费典型指令包括TLOADGM → 片上 Tile 的加载TSTORE片上 Tile → GM 的存储MGATHER按索引从 GM 表张量中收集数据到 TileMSCATTER按索引将 Tile 数据散射写回 GM。从仓库源码看这些指令的入口统一封装在 include/pto/common/pto_instr.hpp例如TLOAD(TileData dst, GlobalData src, WaitEvents... events)第 257-268 行与TSTORE(GlobalData dst, TileData src, WaitEvents... events)第 378-386 行它们的GlobalData模板形参正是GlobalTensor类型在 include/pto/common/pto_tile.hpp 中is_globalGlobalTensor...被特化为std::true_type用于在编译期区分全局张量与 Tile。本文涉及的标识符GlobalTensor、Shape、Stride、Layout、TileShape2D、BaseShape2D等默认定义于 include/pto/common/pto_tile.hpp部分辅助枚举Layout、GlobalTensorDim定义于 include/pto/common/type.hpp。二、GlobalTensor 类型定义四个模板参数源码中的类型定义为template typename Element_, typename Shape_, typename Stride_, Layout Layout_ Layout::ND struct GlobalTensor { using Shape Shape_; using Stride Stride_; using RawDType remove_gm_tElement_; using DType __gm__ RawDType; static constexpr Layout layout Layout_; ... };见 include/pto/common/pto_tile.hpp各模板参数的含义模板参数含义说明Element_GM 中存放的标量元素类型如float、int32_t、halfGlobalTensor::DType即__gm__ Element_指针类型随宏__CCE_AICORE__下定义于 pto_tile.hpp 的remove_gm特化剥离__gm__修饰后重新加上Shape_一个pto::Shape...最多 5 维的 shape 描述Stride_一个pto::Stride...最多 5 个 stride以元素数计而非字节Layout_布局hintND、DN、NZ等默认Layout::ND用于指导 lowering 与目标相关 fast pathGlobalTensor内部还暴露了编译期可见的元数据staticShape[GlobalTensorDim::TOTAL_DIM]从Shape::staticShape拷贝而来pto_tile.hppstaticStride[GlobalTensorDim::TOTAL_DIM]从Stride::staticStride拷贝而来pto_tile.hpp静态成员defaultShape/defaultStride均为全 1pto_tile.hpp用作构造函数的默认参数。三、5 维 Shape 与 Stride静态维度与动态维度PTO 将全局内存张量统一建模为5 维对象。多数 2 维用法会将前三个维度设为1用最后两个维度表示(rows, cols)。维度索引由命名空间GlobalTensorDim提供include/pto/common/type.hppnamespace GlobalTensorDim { constexpr int DIM_0 0; constexpr int DIM_1 1; constexpr int DIM_2 2; constexpr int DIM_3 3; constexpr int DIM_4 4; constexpr int TOTAL_DIM 5; }3.1pto::Shapetemplate int64_t N1 DYNAMIC, int64_t N2 DYNAMIC, int64_t N3 DYNAMIC, int64_t N4 DYNAMIC, int64_t N5 DYNAMIC struct Shape { static constexpr int64_t staticShape[5] {N1, N2, N3, N4, N5}; int64_t shape[GlobalTensorDim::TOTAL_DIM] {1}; ... };见 pto_tile.hpp每个模板参数可以是编译期常量也可以是pto::DYNAMIC值为-1定义于 pto_tile.hpp静态维度固化在类型中通过Shape::staticShape[dim]访问编译器可据此生成更优的代码动态维度存放在运行时的Shape::shape[dim]数组中由Shape(...)构造函数填充。Shape提供 1 到 5 个参数的构造函数以及默认构造函数并通过static_assert强制**运行时参数个数必须等于动态维度个数例如 2 参数构造函数断言动态维度总数必须等于GlobalTensorDim::DIM_2pto_tile.hpp。这意味着如果构造参数与动态维度数量不匹配会在编译期**直接失败从根上避免了运行时越界或漏填维度。参数会按动态维度出现的顺序依次填充到对应shape[dim]槽位默认构造则把所有动态维度置为1。3.2pto::Stridetemplate int64_t SN1 DYNAMIC, int64_t SN2 DYNAMIC, int64_t SN3 DYNAMIC, int64_t SN4 DYNAMIC, int64_t SN5 DYNAMIC struct Stride { static constexpr int64_t staticStride[GlobalTensorDim::TOTAL_DIM] {SN1, SN2, SN3, SN4, SN5}; int64_t stride[GlobalTensorDim::TOTAL_DIM] {1}; ... };见 pto_tile.hppStride与Shape遵循同一套静态/动态 参数数量 static_assert模式但语义是描述跨步stride 以元素数计不是字节数某个维度的 stride 表示该维索引加 1 时指针应前进多少个元素例如行主序row-major的(rows, cols)张量DIM_4最内维stride 为1DIM_3stride 为cols即ldleading dimension。3.3 为什么是 5 维而非更少5 维统一建模的价值在于无论是普通的 2 维(rows, cols)、带 batch 的多维张量还是 NZ / MX 等需要拆分内维的打包布局都能用同一份Shape...Stride...描述而GlobalTensor与内存指令的接口保持稳定这正是 PTO 跨平台可移植性的基础。四、GlobalTensor 的构造与访问接口GlobalTensor存放指针以及动态维度的 shape/stride 运行时值。核心构造与访问接口pto_tile.hppusing GT pto::GlobalTensorfloat, pto::Shape1,1,1,-1,-1, pto::Stride1,1,1,-1,1, pto::Layout::ND; GT t(ptr, /*shape*/{rows, cols}, /*stride*/{ld}); auto* p t.data(); // 返回 DType*__gm__ float* int cols t.GetShape(pto::GlobalTensorDim::DIM_4); int ld t.GetStride(pto::GlobalTensorDim::DIM_3);4.1 构造函数GlobalTensor(DType* data, const Shape shape defaultShape, const Stride stride defaultStride)pto_tile.hpp核心构造函数。内部用if constexpr逐维判断只有staticShape[dim] DYNAMIC的维度才从传入的shape拷贝运行时值静态维度由类型本身携带、不需要也不允许传入stride 同理。GlobalTensor()pto_tile.hpp默认构造data_置nullptr。4.2 访问接口接口语义源码位置data()返回底层DType*指针pto_tile.hppGetShape(int dim)运行时查询第dim维 shape静态维返回编译期值动态维返回运行时值pto_tile.hppGetStride(int dim)运行时查询第dim维 stride语义同上pto_tile.hppGetShapedim()编译期查询仅适用于静态维度对动态维度static_assert报错pto_tile.hppGetStridedim()编译期查询 stride仅适用于静态维度pto_tile.hppSetShapedim(s)等运行时修改动态维 shapestatic_assert要求目标维必须为DYNAMICpto_tile.hppSetStridedim(s)等运行时修改动态维 stride同上约束pto_tile.hpp对完全静态的张量可以完全在编译期查询constexpr int cols GT::GetShapepto::GlobalTensorDim::DIM_4();注意编译期接口对动态维度会触发static_assert如 dim 4 is dynamic, cannot be obtained using the template interface.这正是 PTO 静态/动态维度设计的体现能在编译期确定的绝不放到运行时。另外在__CPU_SIMCPU 仿真模式下GlobalTensor还提供GetElement(i0..i4)/SetElement(...)/AddToElement(...)等逐元素访问辅助接口pto_tile.hpp内部按offset Σ i_dim * GetStride(dim)计算元素偏移方便在 CPU 仿真环境中直接验证 GlobalTensor 的寻址语义。五、布局 hintpto::Layout为什么与 Tile 布局不同GlobalTensor的第 4 个模板参数是布局枚举pto::Layout完整定义位于 include/pto/common/type.hppenum class Layout { ND, // ND RowMajor DN, // DN ColMajor NZ, // NZ for cube SCALE, MX_A_ND, MX_A_DN, MX_A_ZZ, MX_B_ND, MX_B_DN, MX_B_NN, HIF4_A_ZZ, HIF4_B_NN, NC1HWC0, GNC1HWC0, NCHW, GNCHW, NHWC, NDC1HWC0, NCDHW, FRACTAL_Z, FRACTAL_Z_S16S8, FRACTAL_Z_3D, MAX, };其中注释明确给出了基础语义ND是 RowMajor行主序、DN是 ColMajor列主序、NZ是面向 Cube矩阵乘单元的分形打包。其余为SCALE量化 scale、MX 系列MX_A_*/MX_B_*微缩放浮点格式的 A/B 矩阵布局、HIF4 系列、卷积相关的NC1HWC0/NCHW/NHWC等。5.1 为什么 GlobalTensor 的 Layout 不等于 Tile 布局原文档明确指出两者并非一一对应原因有二Tile 布局BLayout/SLayout是 2 维概念外层布局加可选的内层盒化boxed layout是外层内层的二元结构GlobalTensor 是 5 维对象单一的外层内层二元对无法覆盖所有 5 维场景。因此Layout采用标签化hint方式表达常见存储模式例如NDvsDN表达 minor 2-D 的排列顺序NZ表达立方友好的打包方式。这个 hint 会在 lowering 与指令实现中用于选择目标相关 fast path同时也会触发编译期约束检查见下文 TLOAD 实现中的布局匹配static_assert。六、2 维便捷 helperTileShape2D与BaseShape2D2 维张量是算子开发中最常见的形态PTO 为此提供了两个 helper 家族均在 pto_tile.hpp 中6.1pto::TileShape2DT, rows, cols, layout生成 5 维Shapelayout Layout::NDShape1, 1, 1, rows, colspto_tile.hpplayout Layout::DNshape 与 ND 相同Shape1,1,1,rows,colspto_tile.hpp行列顺序的差异体现在 stride 上layout Layout::NZShape1, cols/C0, rows/16, 16, C0其中C0Size由 GetNZC0Size 按平台/类型推导A5/A6 下对float4_e2m1x2_t等类型有专门分支并带static_assertrows必须能被FRACTAL_NZ_ROW16整除、cols必须能被C0Size整除pto_tile.hpp。6.2pto::BaseShape2DT, rows, cols, layout尽管名称含 Shape它实际是stride helper继承自pto::Stride。典型特化Layout::NDStriderows*cols, rows*cols, rows*cols, cols, 1——DIM_4stride 为1DIM_3stride 为cols行主序行间跨步即一行元素数pto_tile.hppLayout::DNStriderows*cols, rows*cols, rows*cols, 1, rows—— 列主序DIM_3stride 为1DIM_4stride 为rowspto_tile.hppLayout::NZStride..., 16*C0, C0, 1将 NZ 分形块内的 C0 连续元素与 16 行块折叠进最内两维 stridepto_tile.hpp。此外TileShape2D与BaseShape2D还支持MX_A_ZZ、MX_A_ND、MX_A_DN、MX_B_NN、MX_B_ND、MX_B_DN、HIF4_A_ZZ、HIF4_B_NN等布局特化pto_tile.hpp分别按照MX_ROW_LEN/MX_COL_LEN/MX_BLOCK_SIZE、HIF4_ROW_LEN/HIF4_COL_LEN/HIF4_BLOCK_SIZE等常量拆解 shape/stride并附带对应的可整除性static_assert。这些 helper 让 2 维代码无需手写 5 维 shape/stride同时把布局约束显式化。七、地址绑定TASSIGNGlobalTensor构造时绑定指针但算子执行中经常需要循环内更新地址如分块搬运时每次推进一个 tile 的偏移。此时使用TASSIGNTASSIGN(globalTensor, ptr);指令入口位于 include/pto/common/pto_instr.hpptemplate typename T, typename AddrType PTO_INST void TASSIGN(T obj, AddrType addr) { MAP_INSTR_IMPL(TASSIGN, obj, addr); }不同后端有对应实现以 CPU 仿真实现 include/pto/cpu/TAssign.hpp 为例逻辑清晰展示了 GlobalTensor 与 Tile 的两类绑定template typename T, typename AddrType PTO_INTERNAL void TASSIGN_IMPL(T obj, AddrType addr) { if constexpr (is_tile_data_vT || is_conv_tile_vT) { static_assert(std::is_integral_vAddrType, Tile can only be assigned with address of int type.); obj.assignData(NPUMemoryModel::Instance().ResolveAssignedAddressT(...)); } else { static_assert(is_global_data_vT, Only Tile and GlobalTensor data types are supported.); static_assert(std::is_pointer_vAddrType, GlobalTensor can only be assigned with address of pointer type.); static_assert( std::is_same_vstd::remove_cv_tstd::remove_pointer_tAddrType, typename T::DType, GlobalTensor can only be assigned with pointer of same data type.); obj.SetAddr(addr); } }要点Tile 用整数地址UB 内地址GlobalTensor 用指针static_assert强制指针的去 cv、去指针后的元素类型必须与GlobalTensor::DType一致类型不匹配在编译期即失败SetAddr是GlobalTensor的私有成员TASSIGN_IMPL通过友元声明pto_tile.hpp访问外部代码无法绕过类型检查直接改写指针。此外TASSIGNAddr(tile)还提供编译期地址重载pto_instr.hpp可对 Tile/ConvTile 执行静态的边界与对齐检查但该重载仅对 Tile 类型启用enable_if_tis_tile_data_vT || is_conv_tile_vTGlobalTensor 仍走运行时指针绑定路径。八、内存指令如何消费 GlobalTensor8.1 TLOADGM → TileTLOAD的指令入口pto_instr.hpp支持带WaitEvents的同步语义最终路由到平台实现。以 include/pto/common/arch/memory/tload_common.hpp 的TLOAD_TILE_IMPL为例可以看到 GlobalTensor 元数据的完整消费方式通过src.GetShape(GlobalTensorDim::DIM_0..DIM_4)与src.GetStride(...)逐维取出 shape/stride 传给底层TLoadGm2ub/TLoadGm2L1系列函数布局 hint 驱动 fast path 选择Vec Tile 只支持ND2ND/DN2DN/NZ2NZ不匹配直接static_assertMat Tile 额外支持ND2NZ、DN2ZN的跨布局加载64 位类型int64_t/uint64_t只允许ND2ND/DN2DN运行前还会做PTO_ASSERT校验源 shape 各维大于 0。8.2 TSTORETile → GMTSTORE入口pto_instr.hpp同样接受(GlobalData dst, TileData src)并提供多个重载支持AtomicType原子存储模式、STPhaseunit-flag 阶段选择、ReluPreModeReLU 前置融合等参数化能力默认AtomicType::AtomicNone。8.3 MGATHER / MSCATTER索引收集与散射MGATHER/MSCATTER同样以 GlobalTensor 作为 GM 侧操作数仓库测试 tests/cpu/st/testcase/mgather/mgather_kernel.cpp 展示了典型用法用GlobalTensorT, TableShape, TableStride描述索引表table、GlobalTensorTIdx, ...描述索引、GlobalTensorT, OutShape, OutStride描述输出再配合对应 Tile 完成 gather 计算。九、最小示例与仓库实战9.1 最小示例原文档#include pto/pto-inst.hpp using namespace pto; void example(__gm__ float* in, __gm__ float* out) { using TileT TileTileType::Vec, float, 16, 16; using GShape Shape1, 1, 1, 16, 16; using GStride BaseShape2Dfloat, 16, 16, Layout::ND; using GT GlobalTensorfloat, GShape, GStride, Layout::ND; GT gin(in); GT gout(out); TileT t; TLOAD(t, gin); TSTORE(gout, t); }要点拆解GShape用Shape1,1,1,16,16显式给出 5 维前三维为1GStride直接用BaseShape2Dfloat, 16, 16, Layout::ND生成行主序 stride等价于Stride256,256,256,16,1GlobalTensor构造即完成指针绑定TLOAD/TSTORE完成一次 16×16 的 GM↔UB 搬运全程无显式地址计算。9.2 仓库实战一add 算子TASSIGN 循环分块真实算子 demos/baseline/add/csrc/kernel/add_custom.cpp 展示了 GlobalTensor 在分块流水中的标准姿势using ShapeDim5 pto::Shape1, 1, 1, tileSRows, tileSCols; using StridDim5 pto::Stride1, 1, 1, tileCols, 1; using GlobalData pto::GlobalTensorT, ShapeDim5, StridDim5; GlobalData xGlobal(x); GlobalData yGlobal(y); GlobalData zGlobal(z);随后在主循环内每个迭代通过TASSIGN(xGlobal, x iterOffset)更新 GM 指针add_custom.cpp再执行TLOAD(xTiles[pingpong_flag], xGlobal)与TSTORE(zGlobal, zTiles[pingpong_flag])add_custom.cpp配合PIPE_V/PIPE_MTE2/PIPE_MTE3的事件同步实现 ping-pong 双缓冲流水。这里 GlobalTensor 的 shape/stride 在一次构造后保持固定只需循环中换指针正是轻量包装设计的高效体现。9.3 仓库实战二hashfind 测试静态 shape 的 GlobalTensortests/cpu/st/testcase/hashfind/hashfind_kernel.cpp 使用完全静态的 shape/stride 构造表张量using TableShape Shape1, 1, 1, 1, kCap; using TableStride Stride1, 1, 1, kCap, 1; using TableGT GlobalTensorint32_t, TableShape, TableStride;这类用法中所有维度都是编译期常量stride 也全部静态编译器可以完全展开寻址计算适用于容量固定如哈希表桶的 GM 数据结构。十、总结与最佳实践围绕pto::GlobalTensor可以提炼出以下实战要点GlobalTensor 不拥有数据它只是GM 指针 5 维 shape/stride 布局 hint的轻量元数据真正的数据搬运由TLOAD/TSTORE/MGATHER/MSCATTER完成。优先静态维度能写成编译期常量的维度写成模板参数可获得编译期寻址优化与static_assert保护动态维度数量必须与构造函数参数个数一致。stride 用元素数、别用字节Stride..., cols, 1表达行主序 2-D 张量的典型跨步BaseShape2D会自动生成正确 stride推荐直接使用。布局 hint 认真选Layout::ND/DN/NZ决定 TLOAD/TSTORE 的 fast path 与布局匹配约束ND2ND、DN2DN、NZ2NZMat Tile 还支持ND2NZ/DN2ZN跨布局搬运MX/HIF4 等特殊布局需保证行列可整除性约束。循环内用TASSIGN换指针分块流水如 add 算子中 shape/stride 保持固定仅更新__gm__指针即可避免重复构造开销。编译期与运行时接口分工明确GetShapedim()/GetStridedim()只适用于静态维度动态维度必须用运行时GetShape(dim)/GetStride(dim)。想进一步深入可继续阅读仓库内相关材料Tile 编程模型、PTO 编程模型总览、向量加教程、TLOAD 指令说明、TSTORE 指令说明以及 CPU 仿真测试入口 tests/run_cpu_tests.sh。【免费下载链接】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），仅供参考