cute 教程 03 Tensor

Cute教程 03 tensor

这篇博客将描述Tensor, 它是应用layout概念的核心载体。

本质上，Tensor表示一个多维数组，Tensor排除了数组元素如何进行组织的细节以及数组元素如何存储。

一个Tensor可以通过两个模板参数描述, Engine和Layout. 数据的layout通过Layout来提供，而真实数据通过iterator提供，这种数据可以存储在任何memory中，gmem smem rmem ... 甚至是on the fly的转换或生成。

基础操作

提供类容器的操作来访问元素：

• data(): Tensorholds的iterator
• size(): Tensor的logical size
• operator[](Coord): 访问逻辑坐标Coord的对应元素
• operator()(Coord): 同上
• operator()(Coords...):

cute的tensor提供了与layout相似的层次化操作:

• rank<I...>(Tensor): tensor 的 I..th mode的rank
• depth<I...>(Tensor):
• shape
• size
• layout
• tensor

Tensor engine

Engine的概念是对数据的iterator或数组的一种封装; 通常来说用户不需要显式构造Engine. 当一个Tensor构造时，合适的engine--通常是ArrayEngine<T,N>, ViewEngine<Iter> ConstViewEngine<Iter>--将会构造。

tagged iterators

任何随机访问的iterator都可以被用来构造一个Tensor，但是用户也需要能够tag特定的内存区域--比如说明这个iterator是访问gmem还是smem? 有如下两种方法:

• make_gmem_ptr(g) or make_gmem_ptr<T>(g): 标注g为gmem的iterator
• make_smem_ptr(g) or make_smem_ptr<T>(g)

标注内存可以使得cute的tensor算法可以使用更适合内存scope的实现。例如，一些优化的copy操作要求copy的src在gmem，dst在smem上。

创建Tensor

Tensor有两种方式构造: owning & non-owning

• Owning: 就像是std::array，当拷贝Tensor时，会深拷贝其中的元素, Tensor的析构函数也会释放这个数组
• Non-owning: 像是raw pointer, 拷贝Tensor并不会拷贝元素, 析构也不会释放这个数组

注意到Tensor作为输入时，我们尽量call-by-ref，否则可能发生深拷贝

Nonowning Tensor

Tensor通常是先有内存的non-owning的view；可以通过调用make_tensor()来创建，输入为一个随机访问的iterator 和 一个layout或是创建一个layout的参数。下面是构造tensor的例子:

float* A = ...;

// non-tagged
Tensor tensor_8 = make_tensor(A, make_layout(Int<8>{})); //
Tensor tensor_8s = make_tensor(A, Int<8{}>);
Tensor tensor_8d2 = make_tensor(A, 8, 2);

// gmem tag
Tensor gmem_8s = make_tensor(make_gmem_ptr(A), Int<8>{});
Tensor gmem_8dx16s = make_tensor(make_gmem_ptr(A), make_shape(8, Int<16>{}), make_stride(Int<16>{}, Int<1>{}));

// shared
Layout smem_layout = make_layout(make_shape(Int<4>{}, Int<8>{}));
__shared__ float smem[decltype(cosize(smem_layout))::value];
Tensor smem_4x8_row = make_tensor(make_smem_ptr(smem), shape(smem_layout), LayoutRight{});

Owning Tensors

一个Tensor也可以是内存的owning数组; Owning tensor可以通过make_tensor<T>来进行创建，这里的T是数组中每个元素的类型, 输入是Layout或是构造Layout的参数。数组类似于std::array<T,N>的分配，必须要静态分配，因此owning tensor的layout必须是静态的shape和stride。

有一些创建owning tensor的例子:

Tensor rmem_4x8_col = make_tensor<float>(Shape<_4, _8>{});
Tensor rmem_4x8_row = make_tensor<float>(Shape<_4, _8>{}, LayoutRight{});
Tensor rmem_4x8_like = make_tensor(rmem_4x8_row);

访问tensor

用户可以通过[]或()来访问对应坐标的元素；当用户给一个坐标时，首先使用Layout映射到对应的物理坐标offset上，根据这个offset可以访问实际物理存储上的元素。

TODO: verify

Tensor A = make_tensor<float>(Shape <Shape <_4, _5>, Int<13>>{}, 
    Stride<Stride<_12,_1>,_64>{});

float* b_ptr = ...;
Tensor B = make_tensor(b_ptr, make_shape(13,20));
// natural coord
for(int m0 = 0; m0 < size<0,0>(A); ++m0){ // mode 0 中的 mode 0
    for(int m1=0; m1 < size<0,1>(A); ++m1){ // mode 0 中的 mode 1
        for(int n = 0; n < size<1>(A); ++n)
            A[make_coord(make_coord(m0,m1), n)] = n + 2 * m0;
    }
}

// Transpose A into B via variadic op()
for(int m =0; m < size<0>(A); ++m) {
    for(int n=0; n < size<1>(A); ++n)
        B(n, m) = A(m, n);
}

// copy B to A
for(int i = 0; i < A.size(); ++i){
    A[i] = B[i];
}

tiling tensor

我们一般使用composition和divide，而不使用product，因为它会增加layout的size。

composition(Tensor, Tiler);
logical_divide
zipped_divide
tiled_divide
flat_divide

slicing tensor

使用坐标访问只会返回一个元素，slicing一个tensor将会返回一个subtensor。

slice操作也可以通过()使用, 通过传递_可以获得跟torch中:的效果。

// ((_3,2),(2,_5,_2)):((4,1),(_2,13,100))
Tensor A = make_tensor(ptr, make_shape (make_shape (Int<3>{},2), make_shape (       2,Int<5>{},Int<2>{})),
                            make_stride(make_stride(       4,1), make_stride(Int<2>{},      13,     100)));

// ((2,_5,_2)):((_2,13,100))
Tensor B = A(2,_);

// ((_3,_2)):((4,1))
Tensor C = A(_,5);

// (_3,2):(4,1)
Tensor D = A(make_coord(_,_),5);

// (_3,_5):(4,13)
Tensor E = A(make_coord(_,1),make_coord(0,_,1));

// (2,2,_2):(1,_2,100)
Tensor F = A(make_coord(2,_),make_coord(_,3,_));

partition tensor 划分

有以下三种比较常见的模式

Inner和outer划分

Tensor A = make_tensor(ptr, make_shape(8,24));
auto tiler = Shape<_4, _8>{};

Tensor tiled_a = zipped_divide(A, tiler); // ((_4, _8), (2,3))

假设我们想给每个threadgroup这些4x8 tiles中的一个，那么我们需要使用我们的threadgroup的coord来在第二个mode中索引.

Tensor cta_a = tiled_a(make_coord(_, _), make_coord(blockIdx.x, blockIdx.y)); // (_4, _8)

我们称这种为inner partition，因为它维持了内在的tile mode, 这种模式进一步封装成了inner_partition(Tensor, Tiler, Coord), 你也经常能看到local_tile(Tensor, Tiler, Coord)这其实是它另一个名字. local_tile经常应用在将tensor划分成threadgroup的tiles上。

作为替代，假设我们有32个线程，并且想给每个线程每个4x8tiles中的一个元素，那么我们可以使用我们线程来在第一个mode进行索引:

Tensor thr_a = tiled_a(threadIdx.x, make_coord(_,_)); // (2,3)

我们称这种parition为outer; 这种操作被封装为outer_partition(Tensor, Tiler, Coord), 也称为local_partition(Tensor, Layout, Idx); 更近一步的示例参见gemm例子。

thread-value划分

在这种模式下，我们构造一个layout来表示所有线程和所有值的映射，每个线程都将受到目标数据的坐标

// 构建一个tv-layout
// 将8个线程编号和4个索引编号映射到4x8 tensor的1D坐标
// (T8, V4) -> (M4, N8)
auto tv_layout = Layout<Shape <Shape <_2, _4>, Shape<_2, _2>>,
    Stride<Stride<_8, _1>, Stride<_4, _16>>>{}; // mode 1 理解为 线程号; mode 2 理解数据编号

//
Tensor A = make_tensor<float>(Shape<_4, _8>{}, LayoutRight{}); // (4,8) tensor

// 映射
Tensor tv = compostion(A, tv_layout); // 跟tv_layout的shape一致 (8,4)

// 每个线程有4个tensor
Tensor v = tv(threadIdx.x, _);

例子

从gmem拷贝一个subtile到rmem

Tensor gmem = make_tensor(ptr, make_shape(Int<8>{}, 16));
Tensor rmem = make_tensor_like(gmem(_,0)); // _8

for (int j = 0; j < size<1>(gmem); ++j) {
    copy(gmem(_,j), rmem);
    do_something(rmem);
}

// 使用tiler模式
Tensor gmem = make_tensor(ptr, make_shape(24, 16));

auto tiler = Shape<_8, _4>{};

Tensor gmem_tiled = zipped_divide(gmem, tiler);
Tensor rmem = make_tensor_like(gmem_tiled(_,0));
for (int j = 0; j < size<1>(gmem_tiled); ++j) {
    copy(gmem_tiled(_,j),rmem);
    do_something(rmem);
}