TileLang学习 · 0130W Notes

因为后面可能在笔记中会用到T之类的，在这里先导入一下文件(雾)

import tilelang
import tilelang.language as T
from tilelang import jit

Defining a Kernel with `@T.prim_func`

在Tilelang中，kernel是由@T.prim_func装饰器生成的TVM IR(TIR)，即Apache TVM的中间表示，装饰器会将普通的Python函数变成一段IR描述，将函数体翻译成一棵TIR的表达树，然后Tilelang就可以复用TVM的优化和后端。

参数一般通过T.Tensor或者T.Buffer进行表示，例如

@T.prim_func
def add(A: T.Tensor((256,), T.float32)):
    pass

TODO: T.Tensor和T.Buffer的区别

T.Tensor需要传入Tensor的尺寸和dtype，这里的dtype可以是:

string类型，如float32
TileLang dtype，如T.float32
framework dtype，如pytorch中的torch.float32

对于要传入的尺寸，最简单的当然是直接写死尺寸，如上面的示例代码，也可以是动态确定的参数, 这里有两种方式：

@jit装饰器，例如说：

@jit
def add(N: int, dtype):
    @T.prim_func
    def add_kernel(
            A: T.Tensor((N, ), dtype),
            B: T.Tensor((N, ), dtype),
            C: T.Tensor((N, ), dtype),
    ):
        pass
    return add_kernel


def main():
    add_kernel = add(256, T.float32)

这个add可以类比c++中的模板类，在编译时传入N的具体数值，Tilelang就按这个N生成一个特化的kernel，尺寸在构建kernel时确定，不同的配置的Kernel需要重新编译，相同的则可以复用。

T.dyn和T.dynamic注解符号 (似乎我找不到T.dyn?)

N = T.dyn['N']  # or N = T.dynamic('N', 'int32') --> by default, int32

@T.prim_func
def add_kernel(
        A: T.Tensor((N, ), T.float32),  # --> 这里的这种N不需要在手动指定 N: ...， TileLang会去自动将传入的Tensor的形状与N做绑定(数值，类型)
        B: T.Tensor((N, ), T.float32),
        C: T.Tensor((N, ), T.float32),
):
    # N = A.shape[0]
    pass

这个是真正动态的kernel，具体的尺寸绑定会在传入Tensor后完成。T.dyn一般建议用于仅表示形状关系的情况，本身并不产生一个变量，单纯负责描述签名，实际上可以通过A.shape[0]之类的获取具体的形状值。T.dynamic适合于那些既需要描述形状，又需要在函数体中参与运算，如做循环边界，tile切分等，T.dynamic真实地产生了一个TIR变量。

Lanuching Work with `T.Kernel`

定义完一个kernel后，我们需要启动该kernel，TileLang提供了如下的语法对kernel配置参数以及context进行声明：

@T.prim_func
def add_kernel(
        A: T.Tensor((K, ), T.float32),
        B: T.Tensor((K, ), T.float32),
        C: T.Tensor((K, ), T.float32),
):
    N = A.shape[0]
    thread_num = 128
    grid_x = T.ceildiv(N, thread_num)   # 可支持dim3
    with T.Kernel(grid_x, thread_num) as bx:    # bx --> blockIdx.x (bx, by, bz)
        pass

在Tile模型中，比较少需要用到thread的下标，大多数情况下在T.Kernel中都是使用T.serial, T.unroll, T.parallel, T.Pipelined。当然也可以拿，用T.get_thread_binding(dim)拿线程下标。

Loops and Control Flow

TileLang中的循环主要有以下四种：

T.serial(start, stop[, step]): 最普通的循环
T.unroll(start, stop[, step]): 展开循环，适合于次数小且确定的内层循环，类似CUDA中#pragma unroll

比如说

for k in range(4):
    s += a[k] + b[k]

这个如果unroll之后，会变成

s = a[0] + b[0]
s = a[1] + b[1]
s = a[2] + b[2]
s = a[3] + b[3]

省去了循环所带来的开销，不过编译器对大循环也不一定会全部展开，可能是部分展开。这样还能够让指令流水线更加直接，更有利于编译器做数据搬运、数据并行、向量化等优化。

T.Parallel(ext0, ext1): 适合element-wise的循环，声明循环之间无依赖，这样编译器可以将这ext0 * ext1个迭代分配给所在block的线程执行，并对访问模式做向量化和仿存合并的优化。
T.Pipelined(iters, num_stages=N): 用于生产者和消费者的软件流水线，生产者生产数据，消费者等待数据进行消费，Pipelined会在第K步执行计算的时候，先将K+1, K+2, … , K+N步需要的数据放进缓冲区中。

控制流就是普通的Python代码的if-else逻辑，if的条件必须是TIR表达式，纯Python Bool值会在编译期折叠。

TileLang有个LegalizeSafeMemoryAccess机制，能够在可能越界的访问处自动插入guard用于防止越界访问(CUDA中则需要手动管理idx下标)。但多维的边界判断还是建议通过T.all_of和T.any_of来写清楚

Memory Scopes and Allocation

TileLang提供了以下的几种内存层级(scope):

global: 对应device memory, 声明T.Tensor时默认对应的张量放在这块
Shared: 共享内存，block内私有，通过T.alloc_shared(shape, dtype)申请
Fragment和scalars: 线程寄存器上，通过T.alloc_fragment(shape, dtype)和T.alloc_var申请。但是在编程视角来看，并不用关心某个fragment应该持有哪几个元素，直接按照block内共享的视角来写就行，TileLang会做布局推断(layout inference)在编译时决定哪个元素应该落在对应的哪个线程的哪个寄存器，并将并行循环切成线程私有的访问。

Moving Data

通过T.copy(src, dst)能够将tiles在不同的内存层级之间移动，不同层级之间，TileLang编译器会生成不同的设备代码，例如

global -> shared: 一般是block内线程协作搬运，线程算好自己要读的地址，编译器安排让相邻线程访问相邻地址，凑成合并访存(coalesced), 从而更好利用显卡带宽。在支持异步拷贝的卡上还可能会被降级为cp.async或者Tensor Memory Accelerator, TMA(sm_90后加入的一块专门负责搬数据的硬件单元)。

T.async_copy则是发起异步搬运，不等待，后面读之前需要通过T.ptx_wait_group(n)(等待最近的n组异步搬运，传0全部等待，类似于CUDA中cp.async.wait_group n)，大多数TileLang程序不用手写类似于__syncthreads()这样的东西，编译器里有ThreadSync("shared")这个pass(中间流程)，会将T.tvm_storage_sync("shared")这段代码在第一次读dst之前插入。

Debugging and Printing

T.print(A, msg="...")

由于GPU线程太多，TileLang对shared/fragment层级的内容打印的时候只会输出一次(tid.x/y/z = 0)，避免刷屏

整理自TileLang Language Basic

Defining a Kernel with @T.prim_func