tilelang-cuda-api

TileLang CUDA API 参考手册

本文档提供 TileLang 核心 API 的详细参考，包括函数签名、参数说明和使用示例。

1. 内核定义与编译

@tilelang.jit(out_idx)

@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"),
    ):
        # 内核实现
        pass
    return main

• 作用: 将 TileLang 函数编译为 GPU 内核
• 参数: out_idx - 指定输出张量的索引列表（如 [-1] 表示最后一个参数为输出）
• 调用: 设置 out_idx 后，运行内核时只需传入输入张量，输出由 TileLang 自动创建

tilelang.compile

kernel = tilelang.compile(my_func, out_idx=[-1])

• 作用: 编译 TileLang 函数（等价于 @tilelang.jit）

2. 内核上下文

T.Kernel(grid_x, grid_y, threads)

with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
    # bx, by 对应 blockIdx.x, blockIdx.y
    pass

• 参数:
  • grid_x: 网格 X 维度大小
  • grid_y: 网格 Y 维度大小（可选）
  • threads: 每个线程块的线程数
• 返回: 线程块索引 (bx, by)

T.ceildiv(a, b)

grid_size = T.ceildiv(N, block_N)

• 参数: a, b - 被除数和除数
• 返回: 向上取整的除法结果
• 用途: 计算网格大小

3. 内存分配 API

T.alloc_shared(shape, dtype)

A_shared = T.alloc_shared((block_M, block_K), "float16")

• 作用: 分配共享内存（对应 GPU 共享内存）
• 参数:
  • shape: 张量形状
  • dtype: 数据类型
• 用途: 缓存频繁访问的数据

T.alloc_fragment(shape, dtype)

C_local = T.alloc_fragment((block_M, block_N), "float")

• 作用: 分配寄存器片段（对应 GPU 寄存器文件）
• 参数:
  • shape: 张量形状
  • dtype: 数据类型
• 用途: 累加器和临时存储

T.alloc_local(shape, dtype)

temp = T.alloc_local((1,), "float32")

• 作用: 分配线程本地内存
• 参数:
  • shape: 张量形状
  • dtype: 数据类型
• 用途: 线程私有的临时变量

4. 数据操作 API

T.copy(src, dst)

# 全局内存到共享内存
T.copy(A[by * block_M, ko * block_K], A_shared)

# 寄存器到全局内存
T.copy(C_local, C[by * block_M, bx * block_N])

• 作用: 高效内存复制，自动合并访问
• 参数:
  • src: 源数据（可以是全局内存切片或寄存器片段）
  • dst: 目标数据

T.clear(tensor)

T.clear(C_local)

• 作用: 将张量清零
• 参数: tensor - 要清零的张量