tilelang-cuda-patterns

TileLang CUDA 编程模式

1. 逐元素操作模式

适用于元素级运算：加法、乘法、激活函数等。

标准代码结构

import tilelang
import tilelang.language as T

@tilelang.jit(out_idx=[-1])
def elementwise_op(M, N, block_M, block_N, threads):
    @T.prim_func
    def main(A: T.Tensor((M, N), "float32"),
             B: T.Tensor((M, N), "float32"),
             C: T.Tensor((M, N), "float32")):
        
        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
            for (local_y, local_x) in T.Parallel(block_M, block_N):
                y = by * block_M + local_y
                x = bx * block_N + local_x
                C[y, x] = A[y, x] + B[y, x]
    
    return main

适用算子

• 算术运算: add, mul, sub, div
• 激活函数: relu, sigmoid, tanh, gelu
• 数学函数: exp, log, sqrt, pow
• 类型转换: cast
• 广播操作: broadcast

关键要点

• 使用 T.Parallel 映射到线程
• 从块索引计算全局索引
• 直接访问全局内存，无需共享内存
• 适当处理边界条件

条件逐元素操作

@tilelang.jit(out_idx=[-1])
def conditional_elementwise(N, threads):
    @T.prim_func
    def main(A: T.Tensor((N,), "float32"),
             B: T.Tensor((N,), "float32"),
             C: T.Tensor((N,), "float32")):
        
        with T.Kernel(T.ceildiv(N, threads), threads=threads) as bx:
            for i in T.Parallel(threads):
                idx = bx * threads + i
                if idx < N:
                    C[idx] = T.if_then_else(
                        A[idx] > 0,
                        A[idx] + B[idx],
                        A[idx] - B[idx]
                    )
    
    return main

2. 归约模式

适用于求和、最大值、归一化等聚合操作。

标准代码结构

@tilelang.jit(out_idx=[-1])
def reduction_op(M, N, block_size):
    @T.prim_func
    def main(A: T.Tensor((M, N), "float32"),
             C: T.Tensor((M,), "float32")):
        
        with T.Kernel(M, threads=block_size) as bx:
            # 分配寄存器片段
            A_local = T.alloc_fragment((N,), "float32")
            C_local = T.alloc_fragment((1,), "float32")
            
            # 加载数据
            T.copy(A[bx, 0:N], A_local)
            
            # ✅ 使用内置归约函数
            T.reduce_sum(A_local, C_local, dim=0)
            
            # 写回结果
            C[bx] = C_local[0]
    
    return main