【Tilelang入门】Tilelang Puzzles 01-02

Tillang Puzzles

一个开源仓库https://github.com/tile-ai/tilelang-puzzles/tree/main

给出用tilelang实现经典算子的例子，附带讲解。分为10个puzzle，每个问题都有待补全文件，和参考实现，以及文字讲解。

采用循序渐进的思路，难度逐渐递增，01-05熟悉语法，06-09实现经典算子，10为挑战复杂实战算子

01 copy

实现一个一维张量的复制操作，主要是熟悉copy接口，多线程编程范式。

对应的cpu代码就是

py 复制代码

    for i in range(N):
        B[i] = A[i]

baseline选择pytorch，会和torch实现对比精度和效率

py 复制代码

def ref_copy_1d(A: torch.Tensor):
    assert len(A.shape) == 1
    assert A.dtype == torch.float16
    return A.clone()

单线程块单线程

主要是用来熟悉代码结构，只启动了一个线程

刚开始，到with T.Kernel中间这部分，是host侧，逻辑仍然是cpu编程，一般操作就是声明需要用的变量类型，shape，初始化一些标量值，或者声明一些变量为编译期常量
N = T.const("N")，类似于cpp里的声明一个变量为const，这样编译时可以基于这个变量的值进行优化。
A: T.Tensor((N,), T.float16)即为声明后面用到的张量，需要传入一个tuple声明shape，以及数据类型。这里(N,)的意思就是一维tuple，加上括号逗号是为了区别于标量表达式。
with T.Kernel(1, threads=1) as _:从这里开始就是device侧了，也就是CUDA编程中的核函数部分，T.Kernel(1, threads=1) 里的参数，最后一个表示每个线程块启动的线程数，前面每个值表示整体grid网格的一个维度的长度，前面有几个变量意味着grid有几个维度。这里的意思就是grid维度是1×11×11×1的，只有1个block，这个block里只有一个线程。返回的_，其实是Kernel这个函数会返回当前线程块的编号，这里没用上，随便拿个变量接收
T.copy(A, B)就是基础copy接口，如果只传名字，就会整体拷贝。也可以传切片，具体看后面的用法。

py 复制代码

@tilelang.jit
def tl_copy_1d_serial(A):
    # The host/declaration part of TileLang script.
    N = T.const("N")
    A: T.Tensor((N,), T.float16)
    B = T.empty((N,), T.float16)

    # The body of the kernel function is written in TileLang DSL.
    # We use T.Kernel to launch a kernel.
    with T.Kernel(1, threads=1) as _:
        # Here T.copy is a built-in TileOp in TileLang.
        # It will automatically utilize available threads in the block
        # to do efficient memory copy (including auto parallelism and vectorization)
        # As we only launch one thread here, it will be lowered into a serial loop copy
        # with certain bit width vectorization (like 128 bits per copy).
        T.copy(A, B)

    return B

单线程块多线程

线程块仍然是只有1×1=11×1=11×1=1个，每个线程块内改成256线程。核函数内部完全不用改，编译器会自动规划这256个线程的分工。这就是tilelang的强大之处。

py 复制代码

@tilelang.jit
def tl_copy_1d_multi_threads(A):
    # The host/declaration part of TileLang script.
    N = T.const("N")
    A: T.Tensor((N,), T.float16)
    B = T.empty((N,), T.float16)

    # TODO: Implement this function
    with T.Kernel(1, threads=256) as _:
        T.copy(A, B)

    return B

多线程块多线程

with T.Kernel(N // BLOCK_N, threads=256) as pid_n:，类似于数据结构中的分块思想，把一个长N的一维数组，划分成N // BLOCK_N个块，每个分给一个256线程的线程块。这里线程块编号有用了，用个变量pid_n来接收
T.copy( A[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N], B[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N], )，copy接口的另一个用法，可以接受张量切片，切片语法和py原生语法完全相同，左闭右开。这里的思想就是每个线程块，负责自己这一块的拷贝，下标范围根据块编号pid_n可计算

py 复制代码

@tilelang.jit
def tl_copy_1d_parallel(A, BLOCK_N: int):
    # The host/declaration part of TileLang script.
    N = T.const("N")
    A: T.Tensor((N,), T.float16)
    B = T.empty((N,), T.float16)

    # TODO: Implement this function
    with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
        T.copy(
            A[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N],
            B[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N],
        )

    return B

02 vector-add

两个一维张量的相加，也就是向量相加

定义

py 复制代码

    for i in range(N):
        C[i] = A[i] + B[i]

baseline

py 复制代码

def ref_add_1d(A: torch.Tensor, B: torch.Tensor):
    assert len(A.shape) == 1
    assert len(B.shape) == 1
    assert A.shape[0] == B.shape[0]
    assert A.dtype == B.dtype == torch.float16
    return A + B

向量加实现

with T.Kernel(N // BLOCK_N, threads=256) as pid_n:和前面的copy框架类似，都是一维算子，对第一维度分块。
for i in T.Parallel(BLOCK_N):第一个坑点，在核函数内部，不能写传统的for i in range(n)了，想要循环，要么for i in T.Parallel(n):，要么for i in T.Serial(n):含义是一个循环，并且是否进行并行化。这里的并行化和copy类似，都是自动规划的，交给编译器了。
C[base_idx + i] = A[base_idx + i] + B[base_idx + i]块内计算下标，进行相加。

py 复制代码

@tilelang.jit
def tl_add_1d(A, B, BLOCK_N: int):
    N = T.const("N")
    A: T.Tensor((N,), T.float16)
    B: T.Tensor((N,), T.float16)
    C = T.empty((N,), T.float16)

    with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
        base_idx = pid_n * BLOCK_N
        for i in T.Parallel(BLOCK_N):
            C[base_idx + i] = A[base_idx + i] + B[base_idx + i]

    return C

乘法Relu

并不是向量加了，但也是element-wise逐元素操作，放在一起讲了。

定义

py 复制代码

    for i in range(N):
        C[i] = max(0, A[i] * B[i])

baseline

py 复制代码

def ref_mul_relu_1d(A: torch.Tensor, B: torch.Tensor):
    assert len(A.shape) == 1
    assert len(B.shape) == 1
    assert A.shape[0] == B.shape[0]
    assert A.dtype == B.dtype == torch.float16
    return (A * B).relu_()

看起来只是元素操作改了一下？用个if-else或者max函数就行？第二个坑点来了，不能在核函数直接写if-else语句，需要用一个函数C[base_idx + i] = T.if_then_else(cond,v1,v2)实现，类似cpp的三目运算符，第一个表达式为真则用v1，否则v2。想用max?也行，但是也要套上T.max(x,y)

py 复制代码

@tilelang.jit
def tl_mul_relu_1d(A, B, BLOCK_N: int):
    N = T.const("N")
    A: T.Tensor((N,), T.float16)
    B: T.Tensor((N,), T.float16)
    C = T.empty((N,), T.float16)

    # TODO: Implement this function
    with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
        base_idx = pid_n * BLOCK_N
        for i in T.Parallel(BLOCK_N):
            C[base_idx + i] = T.if_then_else(
                A[base_idx + i] * B[base_idx + i] > 0,
                A[base_idx + i] * B[base_idx + i],
                0,
            )

    return C

共享内存优化Relu

前面的Relu只是结果对了，性能还和torch有较大的差距。这里需要引入GPU编程中的内存层次

GPU内存层次一般分为：全局内存，共享内存，寄存器。全局内存每个线程块都能访问，共享内存只属于一个线程块，寄存器也是。三者的访问速度逐级变快，但容量逐级变小。这是计算机组成原理里提到过的内存层次，越靠近计算单元的存储单元，速度越快，但对应的容量也越小。或者说，为了让一个存储单元接近计算单元，必须缩小他的容量才能实现
A_local = T.alloc_fragment((BLOCK_N), dtype)，意思是申请一个寄存器上的张量，括号的意思是这是一个元组。用这个元组声明张量shape。这里用寄存器是因为，只有256个float数据，寄存器还存的下。所以用寄存器速度最快。
T.copy(A[base_idx], A_local)声明后，把数据从全局内存搬运到寄存器。后续计算时，直接从寄存器读取数据，速度更快

细心的读者可能会注意到，这里搬运到寄存器再计算，每个位置的计算，读取次数是一次全局内存读，一次寄存器读，写次数是一次寄存器写，一次全局内存写（假设编译器够聪明，对这里进行了变量替换C_local[i] = A_local[i] * B_local[i] C_local[i] = T.if_then_else(C_local[i] > 0, C_local[i], 0)），但前一版，只有一次全局内存读，一次全局内存写，不仅没减少次数，还增加了一次寄存器读，一次寄存器写？这算什么优化？

实际上这个判断是对的，从运行时间上来看，延迟几乎不变。没有变的更好，但也没变得更差。没有更差是因为，共享内存的延迟远小于全局内存，基本可以忽略不计。

但是问题还没有解决，既然没有优化效果，为什么还要做这个优化？答案是只是在这个算子没有优化，因为这里搬运过来后的数据，每个只会用一次，所以优化前后读取次数分别是：1全局 vs 1全局+1寄存器。但如果每个数据在计算是会被用到多次呢？比如矩阵乘法中，每个元素都会被用到O(n)O(n)O(n)次。

那么此时优化前后的读取次数分别为：n全局 vs 1全局+(n-1)寄存器，那么优化后的优势就是显然的了。

py 复制代码

def tl_mul_relu_1d_mem(A, B, BLOCK_N: int):
    N = T.const("N")
    dtype = T.float16
    A: T.Tensor((N,), dtype)
    B: T.Tensor((N,), dtype)
    C = T.empty((N,), dtype)

    # TODO: Implement this function
    with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
        base_idx = pid_n * BLOCK_N
        A_local = T.alloc_fragment((BLOCK_N), dtype)
        B_local = T.alloc_fragment((BLOCK_N), dtype)
        C_local = T.alloc_fragment((BLOCK_N), dtype)

        T.copy(A[base_idx], A_local)
        T.copy(B[base_idx], B_local)

        for i in T.Parallel(BLOCK_N):
            C_local[i] = A_local[i] * B_local[i]
            C_local[i] = T.if_then_else(C_local[i] > 0, C_local[i], 0)

        T.copy(C_local, C[base_idx])

    return C