CANN/catlass TLA张量详解

TLA Tensors【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass本文介绍 TLA 中的Tensor。如果说Layout负责描述“逻辑坐标如何映射到内存”那么Tensor就是在Layout的基础上再绑定具体数据、当前视图起点和存储层级后的可访问对象。在本文中Tensor一律指逻辑视图MakeTensor创建的是视图不发生数据拷贝。operator()的切片结果是子视图不发生数据拷贝。GetTile与TileView返回的是 tile 视图不发生数据拷贝。MakeTensorLike只是把一块已有存储绑定成“与参考 Tensor 逻辑尺寸一致”的新视图本身不执行数据搬运。真正的数据移动应由显式的搬运或计算接口完成而不是由这些视图构造接口隐式完成。关于Layout的基础定义请先参考 Layout。先分清四个组成部分Tensor的模板参数是BuiltinTensor、Layout、Coord、Position。第一次接触时建议先把这四部分分开理解。BuiltinTensorBuiltinTensor是 AscendC 提供的底层张量对象例如GlobalTensor或LocalTensor。它表示“底层存储对象本身”。LayoutLayout描述逻辑坐标如何映射到内存以及逻辑有效范围如何表达。CoordCoord是当前Tensor视图在BuiltinTensor所表达的父逻辑空间中的起点坐标。这里需要特别强调两点coord的单位是元素不是字节。coord表示“这个视图从BuiltinTensor所表达的父逻辑空间的哪里开始看”不是 tile 编号。例如一个逻辑大小为(8, 16)的矩阵中如果某个子 Tensor 的coord()是(2, 4)它表示“这个视图的左上角对应父逻辑矩阵的第 2 行、第 4 列”。PositionPosition是 AscendC 中的位置标签例如Arch::PositionGM{}、Arch::PositionL1{}。它用于区分数据位于 GM、L1、L0 等哪一层存储。Tensor 构造当前使用MakeTensor构造Tensor。using namespace tla; GlobalTensorfloat A ...; auto layout tla::MakeLayoutfloat, Catlass::layout::RowMajor(8, 16); // 1. 默认从逻辑坐标 (0, 0) 开始 auto tensorA MakeTensor(A, layout, Arch::PositionGM{}); // 2. 显式指定当前视图起点 auto tensorA_sub MakeTensor(A, layout, tla::MakeCoord(1, 5), Arch::PositionGM{});可以按下面的方式理解layout决定“如何解释这块内存”。coord决定“当前视图从BuiltinTensor所表达的父逻辑空间的哪里开始”。Tensor 的常用接口TLATensor提供以下常用接口.data()返回底层内存对象。.layout()返回布局。.coord()返回当前视图起点。.shape()返回layout.shape()。.stride()返回layout.stride()。.originShape()返回layout.originShape()。(coord0, coord1, ...)按坐标索引或切片。统一理解三类“坐标”TLA 文档中最容易混淆的是几类不同的“坐标”。下面给出统一约定。元素坐标 element coord元素坐标表示“按元素计数的逻辑位置”例如(row, col)。GetTile、crd2offset、普通索引访问等接口使用的都是这种坐标。tile 坐标 tile coordtile 坐标表示“第几个 tile”不是第几个元素。例如在tileShape (64, 128)时tileCoord (1, 2)表示第 1 个行 tile、第 2 个列 tile。它对应的元素起点是(1 * 64, 2 * 128)。视图起点 view coordtensor.coord()表示当前Tensor视图在BuiltinTensor所表达的父逻辑空间中的起点。它由创建这个视图的操作决定例如MakeTensor、GetTile、TileView或切片操作。可以用一句话概括element coord是元素位置。tile coord是 tile 编号。tensor.coord()是当前视图的起点。用一个完整示例理解coord()下面用同一个矩阵串联MakeTensor、GetTile几种情形。using namespace tla; GlobalTensorfloat A ...; auto layout tla::MakeLayoutfloat, Catlass::layout::RowMajor(8, 16); auto tensorA MakeTensor(A, layout, Arch::PositionGM{}); // tensorA.coord() (0, 0) auto tensorA_sub MakeTensor(A, layout, MakeCoord(1, 5), Arch::PositionGM{}); // tensorA_sub.coord() (1, 5) auto tileA GetTile(tensorA_sub, MakeCoord(2, 4), MakeShape(4, 8)); // tileA.coord() (3, 9)上面分别表示tensorA直接观察整块逻辑矩阵因此起点是(0, 0)。tensorA_sub从BuiltinTensor所表达的父逻辑空间的(1, 5)开始观察因此起点变为(1, 5)。tileA在tensorA_sub的基础上再取一个起点为(2, 4)的 tile因此新视图起点是(1, 5) (2, 4) (3, 9)。使用operator()进行索引与切片TLATensor支持使用operator()做索引也支持使用tla::_表达整维切片返回子 Tensor 视图。基本规则不带tla::_时tensor(i, j, ...)返回一个底层BuiltinTensor访问结果本质上对应tensor.data()[offset]。带tla::_时tensor(..., tla::_, ...)返回子 Tensor 视图被索引的维度会被固定保留tla::_所在维度。这里使用的坐标参数必须是一层 tuple即每个维度都是标量或tla::_不支持嵌套 tuple。等价语义可写为tensor.data()[tensor.layout()(tensor.coord() coord_arg)]输出 Tensor 的维度设输入张量 rank 为 $R$coord中出现tla::_的维度索引集合为 ${d_0, d_1, ..., d_{k-1}}$则输出 Tensor 的 rank 为 $k$。输出 Tensor 的layout.shape()、layout.stride()、layout.originShape()是输入布局在这些维度上的投影。输出 Tensor 的coord()会重新从全 0 开始因为它已经成为新的局部视图。例如对 3D 张量A(B, M, K)auto A2 A3(b, tla::_, tla::_); // 3D - 2D得到 (M, K) 视图 auto A1 A2(r, tla::_) // 2D - 1D得到 (K)视图获取 TileTensorGetTileGetTile用于从父 Tensor 上切出一个 tile 视图不拷贝数据。template class Tensor, class Coord, class Shape auto GetTile(Tensor const tensor, Coord const coord, Shape const shape);参数语义如下coord元素坐标表示 tile 左上角在父 Tensor 逻辑空间中的起点。shapetile 的期望尺寸单位是元素。using namespace tla; auto layout tla::MakeLayoutfloat, Catlass::layout::RowMajor(8, 16); auto tensor MakeTensor(A, layout, Arch::PositionGM{}); // 从逻辑坐标 (2, 4) 开始取一个 4 x 8 的 tile auto tile GetTile(tensor, tla::MakeCoord(2, 4), MakeShape(4, 8));返回结果可理解为tile.coord()tensor.coord()(2, 4)。tile.layout().shape()表示期望 tile 尺寸或其与父布局结构一致的表达形式。tile.layout().originShape()表示该 tile 真实有效的逻辑范围触边时会自动裁剪。使用约束支持tensor.layout().depth 1。若tensor.layout().depth 1即分形或嵌套布局当前GetTileLayout仅支持rank 2。coord与shape都必须为一层 tuple并满足rank(coord) rank(shape) Tensor::rank。边界行为例如父 Tensor 的逻辑尺寸是(8, 16)执行auto tail GetTile(tensor, tla::MakeCoord(6, 10), MakeShape(4, 8));那么期望尺寸仍然是(4, 8)。但逻辑上只剩下 2 行、6 列有效数据。因此tail.layout().originShape()会变成(2, 6)。TileViewTileView与GetTile的行为等价区别只在于输入坐标的单位不同GetTile接收元素坐标。TileView接收 tile 坐标。template class TensorT, class TileCoord, class TileShape auto TileView(TensorT const tensor, TileCoord const tileCoord, TileShape const tileShape);例如auto tensorTileA tla::TileView( tensorA, tla::MakeCoord(0u, kLoopIdx), tla::MakeShape(IntL1_TILE_M{}, IntL1_TILE_K{}) );等价关系TileView与GetTile可以直接按下面的等式理解TileView(t, tileCoord, tileShape) GetTile(t, tileCoord ⊙ tileShape, tileShape)这里的⊙表示逐维相乘例如(1, 2) ⊙ (64, 128) (64, 256)这条等式表示TileView先把 tile 坐标转换为元素坐标。然后按GetTile的规则创建同一个 tile 视图。因此两者的差别只在于调用者提供的是哪一种坐标单位而不是返回结果的逻辑语义。为什么TileView更适合分块循环在实际 kernel 或 block 循环中循环变量通常就是 tile 编号而不是元素坐标。因此TileView往往更直接。下面用同一个按 K 维分块的例子做对比。写法一使用GetTileconstexpr uint32_t tileM 64; constexpr uint32_t tileK 128; for (uint32_t kTile 0; kTile kTiles; kTile) { auto coord tla::MakeCoord(0u, kTile * tileK); auto shape tla::MakeShape(tileM, tileK); auto tensorTileA tla::GetTile(tensorA, coord, shape); // use tensorTileA }写法二使用TileViewconstexpr uint32_t tileM 64; constexpr uint32_t tileK 128; for (uint32_t kTile 0; kTile kTiles; kTile) { auto tensorTileA tla::TileView( tensorA, tla::MakeCoord(0u, kTile), tla::MakeShape(tileM, tileK) ); // use tensorTileA }这两段代码的逻辑结果相同但第二种写法直接使用 tile 坐标更贴近分块循环本身的语义也更不容易把“tile 坐标”和“元素坐标”混淆。创建类似的 TensorMakeTensorLikeMakeTensorLike用于创建一个“逻辑尺寸与likeTensor一致”的新 Tensor。最常见的用途是从一个已有 tile 视图出发在另一层内存中构造对应 Tensor并自动继承其originShape()。在未指定layoutBase时行为为根据 LayoutTagDst 决定布局从 LikeTensor::Element 推断 ElementDst从 likeTensor 的 originShape 提取尺寸。调用MakeLayoutElementDst, LayoutTagDst(originShape())构造目标 layout可能会因分型布局合法要求对shape进行以分型为粒度的向上取整。指定layoutBase时使用MakeLayout(layoutBase.shape(), layoutBase.stride(), likeTensor.originShape())构造目标layout。这里仍然需要强调MakeTensorLike构造的是新视图不执行数据搬运。它只是把用户传入的builtinTensor绑定成一个新的 TLATensor并让这个新视图复用likeTensor的逻辑尺寸语义。当前MakeTensorLike仅支持likeTensor.rank 2。接口分为三类典型场景。// 1) 从 LikeTensor::Element 推断 ElementDst template class LayoutTagDst, class BuiltinTensor, class LikeTensor, class PositionType auto MakeTensorLike(BuiltinTensor const builtinTensor, LikeTensor const likeTensor, PositionType); // 2) 显式指定 ElementDst template class LayoutTagDst, class ElementDst, class BuiltinTensor, class LikeTensor, class PositionType auto MakeTensorLike(BuiltinTensor const builtinTensor, LikeTensor const likeTensor, PositionType); // 3) 提供 layoutBase template class LayoutTagDst, class BuiltinTensor, class LikeTensor, class PositionType, class LayoutBase auto MakeTensorLike(BuiltinTensor const builtinTensor, LikeTensor const likeTensor, PositionType, LayoutBase const layoutBase); template class LayoutTagDst, class ElementDst, class BuiltinTensor, class LikeTensor, class PositionType, class LayoutBase auto MakeTensorLike(BuiltinTensor const builtinTensor, LikeTensor const likeTensor, PositionType, LayoutBase const layoutBase);场景一源和目标元素类型相同这是最常见的场景。例如从 GM 中的一个halftile 创建对应的 L1 Tensor元素类型不变只是存储层级改变。auto tensorTileA tla::TileView( tensorA, tla::MakeCoord(blockM, kTile), tla::MakeShape(L1_TILE_M, L1_TILE_K) ); auto tensorL1A tla::MakeTensorLikeLayoutTagL1A( l1ATensorList[l1ListId], tensorTileA, Arch::PositionL1{} ); // 结果 // 1. tensorL1A 使用 L1 目标布局 // 2. tensorL1A 的 originShape 与 tensorTileA 相同 // 3. 元素类型从 likeTensor 自动推断场景二目标元素类型不同当目标 Tensor 的元素类型与源 Tensor 不一致时需要显式指定ElementDst。例如L0C 中使用 accumulator 类型。需要从half输入生成float累加视图。目标内存对象的PrimType与LikeTensor::Element不同。auto tensorL0C tla::MakeTensorLikeLayoutTagL0C, float( l0cTensor, tensorTileC, Arch::PositionL0C{} ); // 结果 // 1. tensorL0C 的逻辑尺寸继承自 tensorTileC // 2. 目标元素类型显式为 float // 3. 适用于 accumulator 或类型提升场景场景三目标布局需要额外控制有些场景下仅指定LayoutTagDst还不够因为目标布局的基础形状或步长需要用户显式给出。例如目标 Tensor 采用特定分形布局。需要固定某个 L1 的物理排布。注意L0的排布由originShape唯一确定因此定制L0上的非预期排布为不合法行为。需要预先给出特殊的shape/stride结构但逻辑有效范围仍要继承自likeTensor。auto layoutBaseL1A tla::MakeLayouthalf, LayoutTagL1A(L1_TILE_M, L1_TILE_K); auto tensorL1A tla::MakeTensorLikeLayoutTagL1A( l1ATensor, tensorTileA, Arch::PositionL1A{}, layoutBaseL1A ); // 结果 // 1. tensorL1A 的 shape/stride 来自 layoutBaseL1A // 2. tensorL1A 的 originShape 继承自GM上的 tensorTileA // 3. 即使当前 tile 是尾块逻辑有效范围也不会丢失如果既要控制目标布局又要显式指定目标元素类型可以使用同时带layoutBase和ElementDst的重载。实际使用模式在 block 层和 kernel 层常见写法通常是两步用TileView从父 Tensor 得到 tile 视图自动处理边界。用MakeTensorLike在目标内存层级构造对应 Tensor自动继承originShape()。这套模式的价值在于主流程始终围绕 tile 编程。尾块逻辑通过originShape自动传递。数据搬运和计算阶段都能复用同一套逻辑尺寸语义减少边界分支和歧义。【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考