TLA Tensors

这篇文档描述了CATLASS的TLA(Tensor Layout Abstraction)下的Tensor

本质上,Tensor (张量)表示一个多维数组。Tensor 抽象了数组元素在内存中的组织方式与存储方式的细节。这使得用户能够编写通用地访问多维数组的算法,并可根据张量的特性(traits)对算法进行特化。如张量的depth、rank、Layout、数据的类型、位置等。

Tensor 包含4个模板参数: BuiltinTensorLayoutCoordPosition。 关于 Layout 的描述, 请参考 Layout

BuiltinTensor 和 Position

BuiltinTensor 为AscendC内的 GlobalTensor 或者 LocalTensorPosition 为AscendC定义的各层级位置。相关使用参考AscendC文档。

Tensor 构造

当前提供 MakeTensor 接口构造Tensor, 包含四个模板参数: BuiltinTensorLayoutCoordPosition

有如下两种方式构造:

GlobalTensor<float> A = ...;

// 缺省coord, 默认为(0, 0)
Layout w8xh16 = MakeLayout(MakeShape(8, Int<16>{}), MakeStride(Int<16>{},Int<1>{}));
Tensor tensor_8x16 = MakeTensor(A, w8xh16, Arch::PositionGM{});

// 用户指定coord
Tensor tensor_8x16 = MakeTensor(A, w8xh16, tla::MakeCoord(1, 5), Arch::PositionGM{});

Tensors 接口

TLA Tensor 提供获取相应特性的接口:

  • .data(). 返回 Tensor 的内存。

  • .layout(). 返回 Tensorlayout

  • .coord(). 返回 Tensorcoord

  • .shape(). 返回 Tensorshape

  • .stride(). 返回 Tensorstride

获取 TileTensor

提供一个 GetTile 接口获取 Tensor 的一片子tensor,会根据输入坐标对coord进行更新,根据新的Tile的shape变换Layout(只是逻辑层面的数据组织形式),底层的数据实体不变更。

Layout w8xh16 = MakeLayout(MakeShape(8, Int<16>{}), MakeStride(Int<16>{},Int< 1>{}));
Tensor tensor_8x16 = MakeTensor(A, w8xh16, Arch::PositionGM{});

auto tensor_tile = GetTile(tensor_8x16, tla::MakeCoord(2, 4), MakeShape(4, 8)); // (4,8):(_16,_1)