tilelang-cuda-optimization

TileLang CUDA 性能优化指南

1. 性能优化策略

1.1 分块大小选择

• 原则: 平衡并行度与资源占用
• 建议: 使用 2 的幂次
• 常用值: block_M/block_N = 64, 128, 256; block_K = 16, 32, 64

算子类型	推荐分块大小	线程数
Element-wise	block = 256-1024	128-256
GEMM	block_M=128, block_N=128, block_K=32	128
Reduce	block = 256-512	128-256

1.2 软件流水线优化

def pipelined_computation():
    # 选择合适的流水线深度
    num_stages = 3  # 通常 2-4 个阶段效果最好
    
    for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
        # 重叠内存操作和计算
        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)

流水线深度选择：

• num_stages=2: 最少的共享内存使用
• num_stages=3: 通常最优（推荐默认值）
• num_stages=4: 更多重叠但占用更多共享内存
• num_stages=5+: 可能超出共享内存限制

1.3 并行化策略

# 1. 细粒度并行
for i, j in T.Parallel(block_M, block_N):
    # 自动映射到线程
    pass

# 2. 向量化优化
for k in T.vectorized(TILE_K):
    A_local[k] = A[bk * BLOCK_K + tk * TILE_K + k]

# 3. 串行循环（必要时使用）
for k in T.serial(block_K):
    # 顺序执行
    pass

1.4 数据类型优化

# 1. 使用混合精度
input_dtype = "float16"    # 输入数据
accum_dtype = "float"      # 累加器使用更高精度

# 2. 类型转换优化
result = A[i].astype(accum_dtype) * B[i].astype(accum_dtype)

# 3. 避免不必要的类型转换（在计算前统一转换）

2. 内存优化策略

2.1 内存层次结构优化

def memory_optimized_matmul(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")):
        
        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
            # 1. 共享内存分配 - 缓存频繁访问的数据
            A_shared = T.alloc_shared((block_M, block_K), "float16")
            B_shared = T.alloc_shared((block_K, block_N), "float16")
            
            # 2. 寄存器片段分配 - 累加和临时存储
            C_local = T.alloc_fragment((block_M, block_N), "float")
            
            # 3. 启用 swizzle 以提高 L2 缓存局部性
            T.use_swizzle(panel_size=10, enable=True)
            
            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)
            
            T.copy(C_local, C[by * block_M, bx * block_N])
    
    return main