Intro-ops/docs/tilelang-vs-cuda.md

TileLang vs CUDA 对比教程

intro-ops 的特色之一是同时支持 CUDA 和 TileLang 两种后端。本文通过 copy 算子的两种实现，帮助理解"同一个算法，两种 DSL 的表达差异"。

概述对比

维度	CUDA	TileLang
抽象层级	线程级（thread/warp/block）	tile 级（fragment/tile）
编写难度	需关注同步、shared memory 细节	声明式，编译器处理调度
编译方式	nvcc 编译 .cu 文件	JIT 编译，无需构建
调优手段	手动调整 grid/block/shared memory	T.Parallel / T.Serial 语义
适用场景	极致性能调优	快速原型 + 跨硬件
学习曲线	陡峭	平缓

---

案例一：copy 算子

CUDA 实现

template <typename T>
__global__ void copy_contiguous_kernel(T *dst, const T *src, int64_t n) {
    // 1. 计算全局线程索引
    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    // 2. 计算网格总步长
    int64_t stride = gridDim.x * blockDim.x;
    // 3. grid-stride loop
    for (int64_t i = idx; i < n; i += stride) {
        dst[i] = src[i];
    }
}

关键细节：

• 你需要手动管理线程索引（blockIdx.x、blockDim.x、threadIdx.x）
• 你需要手动设计 grid-stride loop 来处理任意大小的 tensor
• launch 时你需要自己指定 grid size 和 block size
• 你需要理解"一个线程负责哪些元素"的映射关系

TileLang 实现

@tilelang.jit
def copy_kernel(src, BLOCK_N: int, dtype):
    N = T.const("N")
    src: T.Tensor((N,), dtype)
    out = T.empty((N,), dtype)

    for i in T.Parallel(N // BLOCK_N):
        tile = T.copy(src[i * BLOCK_N : (i + 1) * BLOCK_N])
        out[i * BLOCK_N : (i + 1) * BLOCK_N] = tile

    return out

关键细节：

• T.Parallel 声明一个 tile 级别的并行循环——编译器负责把 tile 分配到线程
• T.copy 加载一个 tile 到局部 fragment——不需要手写线程索引
• 没有显式的 grid/block/thread 概念
• 编译器自动决定最优的 launch 参数

核心差异：谁负责什么

职责	CUDA	TileLang
线程→元素映射	你手动写 idx = blockIdx.x * blockDim.x + threadIdx.x	编译器生成
循环边界检查	你手动写 i < n	编译器生成
shared memory 管理	你手动声明、加载、同步	编译器处理
launch 参数	你手动指定 grid/block size	编译器自动选择
内存合并访问	你需要确保访问模式正确	编译器优化

---

案例二：reduce_sum 算子

CUDA 实现

树形归约需要关注 shared memory 管理和线程同步：

template <typename T>
__global__ void reduce_sum_kernel(T *out, const T *in, int64_t N, int64_t C) {
    extern __shared__ float smem[];
    int64_t row_id = blockIdx.x;
    int tid = threadIdx.x;

    // 1. 每个线程累加自己负责的列
    float sum = 0.0f;
    for (int64_t c = tid; c < C; c += blockDim.x) {
        sum += (float)in[row_id * C + c];
    }
    smem[tid] = sum;
    __syncthreads();

    // 2. 树形归约
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            smem[tid] += smem[tid + s];
        }
        __syncthreads();
    }

    // 3. 输出结果
    if (tid == 0) {
        out[row_id] = (T)smem[0];
    }
}

陷阱点：

• __syncthreads() 必须放在条件分支外
• shared memory 大小需要作为 launch 参数传入
• 树形归约的 stride 和边界条件容易写错

TileLang 实现

@tilelang.jit
def reduce_sum_kernel(inp, BLOCK_C: int, dtype):
    N = T.const("N")
    C = T.const("C")
    inp: T.Tensor((N, C), dtype)
    out = T.empty((N,), dtype)

    for n in T.Parallel(N):
        acc = T.alloc_fragment((1,), dtype)
        acc[0] = 0.0
        for c in T.Serial(C // BLOCK_C):
            tile = T.copy(inp[n, c * BLOCK_C : (c + 1) * BLOCK_C])
            acc[0] += T.reduce_sum(tile)
        out[n] = acc[0]

    return out

关键差异：

• shared memory 归约被 T.reduce_sum 替代——编译器负责生成高效的归约代码
• T.Serial 保证分块累加的顺序（因为累加有依赖）
• 不需要手动管理 shared memory 大小和同步

---

选择建议

什么时候用 CUDA？

• 需要极致性能，想手动控制每一个优化细节
• 算子涉及复杂的 shared memory / warp-level 操作
• 需要嵌入已有 CUDA 生态（CUTLASS、cuBLAS）

什么时候用 TileLang？

• 快速原型：写一个可工作的 kernel 远比性能重要
• 跨硬件：同一份代码可编译到 NVIDIA / MetaX / 甚至未来其他硬件
• 学习阶段：先理解算法逻辑，不必被线程细节分心

intro-ops 的建议路径

1. 入门：先写 TileLang 版本，快速跑通算法逻辑
2. 进阶：再写 CUDA 版本，理解线程级别的执行细节
3. 专精：对比两个版本，分析编译器生成的差异，手动调优 CUDA