Triton/TileLang Tile-Level Programming

Why we need tile-level programming 为什么需要分块级编程

在最普通的SIMT中,我们会遇到线程协作、数据复用、shared memory这些内容都需要手动管理,增大了编程难度;且写起来繁琐而易错,优化比较困难.为了摆脱这种困境,Tile-Level是我们的首选.它把block内线程协作抽象为分块(tile)操作,相关工具负责把tile操作映射到具体的线程,使程序员不需要详细思考操作维护底层线程的逻辑,简化了实现难度.

在Tile-Level编程中,我们通常需要遵循下面的步骤或思路:

• 确定当前tile的位置
• 从global memory加载tile
• 对tile进行计算
• 将结果写回global memory
• 处理边界(不一定最后,根据实际处理)

而为了实现这一过程,我们采用Triton/TileLang工具,作为python的两种不同的扩展语言,它们是以上思想的两种DSL(Domain-Specific Language)表达.

---

Triton模型

Triton与CUDA主要有以下的区别:

1. CUDA

   • CUDA Kernal从thread级别出发,显式地分配计算每个线程负责的数据

2. Triton Kernal

   • Triton Kernal从program出发,每个program负责一个数据tile
   • 程序员描述tile的位置、形状以及load/compute/stroe操作
   • 编译器负责将tile内的计算映射到thread和warp

以下给出一个计算向量加法的kernal示例:

#引入相关模块
import torch
import triton
import triton.language as tl

@triton.jit    #声明kernal函数
def add_kernal(x, y, z, n: tl.constexpr, BLOCK_SIZE: tl.constexpr):
    '''
    传入x,y两个向量,并将结果写入z
    n是向量长度参数
    BLOCK_SIZE是block大小参数
    tl.constexpr表示是常量参数
    '''
    pid = tl.program_id(0)  #得到program id,0代表是一维
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   #全局坐标映射.tl.arange(a, b)获取块内从a到b所有位置id
    mask = offsets < n  #根据是否越界来得到mask的真假

    x_val = tl.load(x + offsets, mask=mask, other=0.0)  
    y_val = tl.load(y + offsets, mask=mask, other=0.0)
    #载入对应元素,other代表对超过mask的部分进行填0补充

    out = x_val + y_val #计算加和

    tl.store(z + offsets, out, mask=mask)   #将out结果写回z向量
    

#CPU host端入口,用于启动运行核函数
def add(x, y):
    z = torch.empty_like(x) #创建一个形状和数据类型与x完全相同,但内容未初始化的空张量
    n = x.numel()   #x向量长度
    block_size = 1024   #块长,即内含线程大小
    grid = (triton.cdiv(n, block_size),)    #计算分块数量,即描绘了gird的整体模样
    add_kernal[grid](x, y, z, n, BLOCK_SIZE=block_size) #启动核函数
    return z    #返回结果

类似地,得到矩阵分块乘法的代码:

import torch
import triton
import triton.language as tl

@triton.jit
def matmul_kernal(
    A, B, C,
    M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
    stride_am: tl.constexpr, stride_ak: tl.constexpr,
    stride_bk: tl.constexpr, stride__bn: tl.constexpr,
    stride_cm: tl.constexpr, stride__cn: tl.constexpr,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr
):
    '''
    A, B, C 计算的矩阵AB以及结果输出C矩阵
    M, N, K 矩阵大小参数
    stride_am等 行参数步长
    stride_kn等 列参数步长

    '''
    pid_m = tl.program(0)   #行分块id
    pid_n = tl.program(1)   #列分块id

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)    #m对应维度坐标映射
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)    #n对应维度坐标映射
    offs_k = tl.arange(0,BLOCK_K)

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)    #C对应区域的累加器,初始化为0

    for k0 in range(0, K, BLOCK_K):
        a_ptrs = A + offs_m[:, None] * stride_am + (k0 + offs_k[None, :]) * stride_ak
        b_ptrs = B + (k0 + offs_k[:, None]) * stride_bk + offs_n[None, :] * stride_bn   ##得到对应A中行和B中列的绝对地址,实际上是BLOCK_M*BLOCK_K维度矩阵和BLOCK_K*BLOCK_N维矩阵

        a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (k0 + offs_k[None, :] < K), other=0.0)
        b = tl.load(b_ptrs, mask=(k0 + offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0) #加载对应位置元素

        acc += tl.dot(a, b)

    c_ptrs = C + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn  #得到C中元素的绝对地址
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, acc, mask=c_mask)  #将累加器结果写入C对应位置


def matmul(A, B):
    M, K = A.shape
    K2, N = B.shape
    assert K == K2

    C = torch.empty((M, N), device=A.device, dtype=torch.float32)
    block_m, block_n, block_k = 16, 16, 32
    grid = (triton.cdiv(M, block_m), triton.cdiv(N, block_n))

    matmul_kernal[grid](
        A, B, C,
        M, N, K,
        A.stride(0), A.stride(1),
        B.stride(0), B.stride(1),
        C.stride(0), C.stride(1),   #获取对应矩阵的行元素数量与列元素数量以得到对应步长
        BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k 
    )
    return C

所以,对于上述代码,只要你读了,就可以看到Triton的局限:

• 编译器是黑盒,内部全靠编译器
• 生成的PTX/SASS不透明,性能问题难以定位
• 复杂的pipeline或warp specialization难表达
• 复杂kernal写起来别扭
• 报错和性能回退难分析

因此,我们需要TileLang

---

TileLang Programming Model TileLang编程模型

TileLang更强调显示写出层次结构,拥有GPU kernal的同时,可以输出显存、共享内存、CTA内共享缓存块等,也可以沿着某维度分块循环,更接近于手写的CUDA kernal结构

以下是一个示例图:

右侧补全的完整矩阵乘法代码贴在这里:

import tilelang
import tilelang.language as T

@tilelang.jit
def matmul(
    M: int,
    N: int,
    K: int,
    BLOCK_M: int,
    BLOCK_N: int,
    BLOCK_K: int,
    dtype: str = "float16",
    accum_dtype: str = "float32",
):
    @T.prim_func
    def main(
        A: T.Buffer((M, K), dtype),
        B: T.Buffer((K, N), dtype),
        C: T.Buffer((M, N), accum_dtype),
    ):
        with T.Kernel(
            T.ceildiv(N, BLOCK_N),
            T.ceildiv(M, BLOCK_M),
            threads=128,
        ) as (bx, by):
            A_shared = T.alloc_shared((BLOCK_M, BLOCK_K), dtype)
            B_shared = T.alloc_shared((BLOCK_K, BLOCK_N), dtype)
            C_local = T.alloc_fragment((BLOCK_M, BLOCK_N), accum_dtype)

            T.clear(C_local)

            for k in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3):
                T.copy(A[by * BLOCK_M, k * BLOCK_K], A_shared)
                T.copy(B[k * BLOCK_K, bx * BLOCK_N], B_shared)
                T.gemm(A_shared, B_shared, C_local)

            T.copy(C_local, C[by * BLOCK_M, bx * BLOCK_N])

    return main

---

Conclusion 小结

对于Triton,TileLang,以及传统CUDA,三者有着明显不同的应用场景以及优势: 1. Triton适合快速简洁的效率场景,尤其是不关注底层以及加速的算法. 2. TileLang则适合既需要一定简洁性,又能够表达出一定底层内存层次和数据情况的场景,易于讲清楚高性能kernal结构. 3. CUDA的控制粒度最细,深入底层结构,对线程级别规划以及内存控制最为直观,适合针对底层的编写或者优化场景,对于基于硬件结构的极端性能调控或者特殊硬件机制研究有无可比拟的绝对优势,但代价是代码比较繁复,维护理解较困难.

擅长python语言或有python对GPU编程基础的读者,对前两种可能更容易掌握;对于CUDA,则更适合于C/C++的编程者.