tilelang-cuda-basics

TileLang CUDA 编程基础

1. 核心概念

TileLang 简介

• 定义: TileLang 是专为高性能 GPU/CPU 内核开发设计的领域特定语言（DSL），采用类似 Python 的语法，底层基于 TVM 编译器
• 特点: 专注于生产力而不牺牲底层优化能力，提供三层抽象级别

编程接口层次

• Level 1 (硬件无关): 编译器自动处理内存层次和硬件特定优化，适合快速原型开发
• Level 2 (硬件感知 + Tile库): 提供预定义 Tile 库操作和模式，适合大多数高性能计算应用
• Level 3 (硬件感知 + 线程原语): 提供线程原语和低级构造的直接访问，适合极致性能优化

内核 (Kernel)

• 定义: 使用 @tilelang.jit 装饰的函数，编译后在 GPU 上并行执行
• 结构: 内部包含 @T.prim_func 装饰的主函数，通过 T.Kernel 上下文管理器定义并行执行逻辑

网格 (Grid) 与线程块

• 网格: 内核启动时的并行维度配置，使用 T.ceildiv 计算块数
• 线程块: 每个块包含指定数量的线程，通过 threads 参数设置
• 块索引: T.Kernel 上下文返回 (bx, by) 对应 blockIdx.x, blockIdx.y

内存层次

• 全局内存 (Global Memory): GPU 主内存（HBM），所有线程可访问
• 共享内存 (Shared Memory): SM 内共享，通过 T.alloc_shared 分配
• 寄存器片段 (Fragment): 对应 GPU 寄存器文件，通过 T.alloc_fragment 分配
• 本地内存 (Local): 线程本地存储，通过 T.alloc_local 分配

2. 标准内核结构

TileLang 内核的标准结构模式：

import tilelang
import tilelang.language as T

@tilelang.jit(out_idx=[-1])
def my_kernel(M, N, K, block_M, block_N, block_K):
    @T.prim_func
    def main(A: T.Tensor((M, K), "float16"),
             B: T.Tensor((K, N), "float16"),
             C: T.Tensor((M, N), "float16")):
        
        # 1. 定义内核上下文（网格和线程配置）
        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
            # 2. 内存分配
            A_shared = T.alloc_shared((block_M, block_K), "float16")
            B_shared = T.alloc_shared((block_K, block_N), "float16")
            C_local = T.alloc_fragment((block_M, block_N), "float")
            
            # 3. 初始化
            T.clear(C_local)
            
            # 4. 计算逻辑（含数据加载和计算）
            for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
                T.copy(A[by * block_M, ko * block_K], A_shared)
                T.copy(B[ko * block_K, bx * block_N], B_shared)
                T.gemm(A_shared, B_shared, C_local)
            
            # 5. 结果写回
            T.copy(C_local, C[by * block_M, bx * block_N])
    
    return main