【cutlass】cuTe 01 layout

简介

Layout将坐标映射到地址空间，其对阵列单元在内存中如何排布进行抽象，并提供了获取多维阵列的接口。用户可以以正常的方式（不需要做复杂的地址计算）写多维阵列存取代码，对于存取的不同方式可以改变Layout来实现，而不需要改代码。cuTe提供了对Layout的操作算法，这些方式可以组合/拆分以构建更复杂的layout。

Layout和Tensor的关系

Layout的逻辑地址空间表示了数据的逻辑形状，它和data共同组成了Tensor。Layout把逻辑地址映射到了一个index，index是一个偏移量用来索引地址数据。

Contains

一个Layout包含了Shape和Stride，他们都是IntTuple类型。IntTuple是通过一个单独的整数或者一个IntTuple的Tuple递归定义的。这意味着IntTuple可以任意组合到一起，包含了以下运算：

• get(IntTuple)：获取IntTuple的第I个单元，如果IntTuple只有一个整数单元那get<0>就是得到这个单元。
• rank(IntTuple)：IntTuple中elements的数量，单独的整数rank是1，一个tuple的rank是tuple_size。
• depth(IntTuple)：IntTuple的层数。一个单独的整数depth是0，一个tuple是1，一个tuple中含有另一个tuple是2，以此类推。
• size(IntTuple)：IntTuple中所有element的乘积
  我们用括号写“IntTuple”来表示层次结构。比如6, (2), (4,3), (3,(6,2),8) 都是IntTuple。

IntTuple和Layout是可以嵌套使用的，为了方便使用，cutlass定义了一些方便使用的函数，比如get<I0,I1,…,IN>(x) = get(…(get(get(x)))…)。这些函数包括：

• rank<I…>(x) := rank(get<I…>(x))：第I个单元的rank
• depth<I…>(x) := depth(get<I…>(x))：第I个单元的depth
• size<I…>(x) := size(get<I…>(x))：第I个单元的size

例子

• 向量
  • shape: (8); stride: (1) 代表了一个8个单元的向量(0,1,2,3,4,5,6,7)。如果stride为(2)，则向量变为(0,2,4,6,8,10,12,14)； 如果想要代表(0,1,2,3,4,5,6,7)八个数也可以按照 shape:((4,2)) stride:((1,4))放入内存中
• 矩阵
  • shape: (4,2)；stride: (1,4) 代表
    0 4
    1 5
    2 6
    3 7
  • 每个矩阵的模式都可以切分成多指数，比如shape: ((2,2),2)；stride: ((4,1),2)代表
    0 2
    4 6
    1 3
    5 7

c++实现方式

Layout可以通过compile-time int或者 run-time int构建，如下

auto layout_8s = make_layout(Int<8>{});
auto layout_8d = make_layout(8);

auto layout_2sx4s = make_layout(make_shape(Int<2>{},Int<4>{}));
auto layout_2sx4d = make_layout(make_shape(Int<2>{},4));

auto layout_2x4 = make_layout(make_shape (2, make_shape (2,2)),
                              make_stride(4, make_stride(2,1)));

其中make_layout函数返回一个Layout类型，类似的make_shape和make_stride函数返回一个shape和stride类型。

对Layout的使用是将坐标和索引相关联，比如打印一个Layout

template <class Shape, class Stride>
void print2D(Layout<Shape,Stride> const& layout)
{
  for (int m = 0; m < size<0>(layout); ++m) {
    for (int n = 0; n < size<1>(layout); ++n) {
      printf("%3d  ", layout(m,n));
    }
    printf("\n");
  }
}

总结

• layout的shape定义了坐标空间
  • • 每个layout都有一个一维的坐标空间。这个可以用来产生一个“通用colum-major”序列
    • 每个layout都有R维坐标空间，其中R是layout的rank。这个空间按照colexicographically（从右到左）排序，对序列的枚举与上面所说的一维类似。
    • 每个layout有h维坐标空间，其中h是层级。这个空间也是按照colexicographically排序，对序列的枚举与上面所说的一维类似。
  • h维坐标相对于shape都是连续的，这样坐标上的每个element都有相应的shape element
• layout的stride将坐标空间映射为地址
  • 通常情况，这种映射可以是一个从一维到地址的任意函数
  • 在cuTe中，我们使用一个内置推导，他是一个具有stride元素的h维坐标