wenyi-li TileLang CUDA 同步和线程安全
本文档详细说明 TileLang 中 T.sync_threads() 的使用规范和线程安全最佳实践。同步问题是 TileLang 编程中最常见也是最危险的错误来源。
1. T.sync_threads() 使用规范
⚠️ 严格禁止的使用场景
1.1 条件分支中的同步
# ❌ 错误示例 - 会导致死锁
if condition:
T.sync_threads() # 只有部分线程执行，其他线程永远等待
# ✅ 正确做法 - 所有线程都执行同步
T.sync_threads()
if condition:
# 同步后的操作
1.2 循环中的条件同步
# ❌ 错误示例 - 死锁风险
for i in range(n):
if tid < threshold:
T.sync_threads() # 死锁风险
# ✅ 正确做法
for i in range(n):
T.sync_threads() # 所有线程都同步
if tid < threshold:
# 同步后的操作
1.3 共享内存分配后的条件同步
# ❌ 错误示例
if tid < N:
shared_mem = T.alloc_shared((N,), dtype)
T.sync_threads() # 只有部分线程分配了共享内存
# ✅ 正确做法
shared_mem = T.alloc_shared((N,), dtype)
T.sync_threads()
if tid < N:
# 使用共享内存
2. ❌ 绝对禁止：手动归约
手动归约是导致线程卡死的最常见原因。必须使用内置归约函数。
手动归约的危险
# ❌ 绝对禁止：手动归约会导致线程卡死
while stride > 0:
if tid < stride:
shared[tid] += shared[tid + stride]
T.sync_threads() # 死锁风险，线程卡死
stride //= 2
# ❌ 绝对禁止：条件分支中的同步
if condition:
T.sync_threads() # 死锁风险
# ❌ 绝对禁止：循环中的条件同步
for i in range(n):
if tid < threshold:
T.sync_threads() # 死锁风险
症状识别
- UTL（GPU 利用率）打满但 MEM 低: 通常是线程卡死在同步点
- 内核永远不返回: 死锁导致的无限等待
- 性能极差: 不当的同步模式导致串行化
3. ✅ 正确的同步模式
3.1 在共享内存操作前后同步
# ✅ 写入共享内存后同步
shared_mem[tid] = value
T.sync_threads() # 确保所有写入完成
# ✅ 读取共享内存前同步
T.sync_threads() # 确保所有写入完成
result = shared_mem[tid]
3.2 确保所有线程参与同步
# ✅ 正确的同步模式
T.sync_threads() # 所有线程都必须执行
# 后续操作
3.3 避免不必要的同步
# ❌ 过度同步
for i in range(n):
T.sync_threads() # 每次迭代都同步，开销大
# ✅ 只在必要时同步
# 只在数据依赖需要时添加同步
4. ✅ 推荐的内置函数（替代手动同步）
4.1 内置归约函数
# ✅ 推荐：内置归约函数，无需手动同步
T.reduce_sum(input_tensor, output_tensor, dim=axis) # 求和归约
T.reduce_max(input_tensor, output_tensor, dim=axis) # 最大值归约
T.reduce_min(input_tensor, output_tensor, dim=axis) # 最小值归约
T.reduce_mean(input_tensor, output_tensor, dim=axis) # 平均值归约
4.2 内置内存操作
# ✅ 推荐：内置内存操作，自动处理同步
T.copy(src, dst) # 高效的内存复制
T.clear(tensor) # 清零操作
T.fill(tensor, value) # 填充操作
4.3 内置数学函数
# ✅ 推荐：内置数学函数
T.rsqrt(x) # 平方根倒数
T.sqrt(x) # 平方根
T.exp(x) # 指数函数
T.log(x) # 对数函数
5. 内置归约函数使用示例
5.1 LayerNorm 中的归约
A_pow_local = T.alloc_fragment((M, N), "float32")
A_powsum = T.alloc_fragment((M,), "float32")
# ✅ 使用内置归约替代手动同步
T.reduce_sum(A_pow_local, A_powsum, dim=1) # 按行求和
5.2 Softmax 中的归约
A_exp = T.alloc_fragment((M, N), "float32")
A_sum = T.alloc_fragment((M,), "float32")
A_max = T.alloc_fragment((M,), "float32")
# ✅ 使用内置归约
T.reduce_max(A_exp, A_max, dim=1) # 按行求最大值
T.reduce_sum(A_exp, A_sum, dim=1) # 按行求和
5.3 完整的安全归约示例
@tilelang.jit(out_idx=[-1])
def safe_reduction(M, N, block_size):
@T.prim_func
def main(x: T.Tensor((M, N), "float32"),
y: T.Tensor((M,), "float32")):
with T.Kernel(M, threads=block_size) as bx:
# 分配寄存器片段
x_local = T.alloc_fragment((N,), "float32")
sum_local = T.alloc_fragment((1,), "float32")
# 加载数据
for tid in T.Parallel(N):
x_local[tid] = x[bx, tid]
# ✅ 安全的归约方式
T.reduce_sum(x_local, sum_local, dim=0)
# 写回结果
y[bx] = sum_local[0]
return main
6. 推荐的并行计算模式
6.1 使用 T.Parallel
# ✅ 推荐：使用 T.Parallel 进行并行计算
for i, j in T.Parallel(M, N):
result[i, j] = input[i, j] * scale[i]
6.2 使用 T.vectorized
# ✅ 推荐：向量化操作
for i in T.vectorized(N):
result[i] = input[i] * scale
6.3 使用 T.Pipelined
# ✅ 推荐：软件流水线
for k in T.Pipelined(K, num_stages=3):
T.copy(A[k], shared_A)
T.gemm(shared_A, shared_B, result)
7. 线程索引获取规范
✅ 推荐：T.Parallel
# ✅ 推荐方式：使用 T.Parallel 获取线程索引
for tn in T.Parallel(BLOCK_N):
# tn 对应 threadIdx.x，自动处理线程映射
pass
⚠️ 不推荐：T.get_thread_binding
# ⚠️ 不推荐：可能导致问题
tn = T.get_thread_binding(0) # threadIdx.x
tk = T.get_thread_binding(1) # threadIdx.y
为什么不推荐：
- 可能导致线程映射错误
- 手动管理线程索引容易出错
T.Parallel提供更安全、更简洁的替代方案
8. 重要警告总结
❌ 绝对禁止
- 手动归约: 会导致线程卡死，UTL 打满但 MEM 低
- 条件分支中的同步: 会导致死锁
- 循环中的条件同步: 会导致死锁
- 条件分支中的共享内存分配: 会导致未定义行为
✅ 必须遵循
- 使用内置归约函数:
T.reduce_sum(),T.reduce_max()等 - 确保所有线程参与同步:
T.sync_threads()必须在所有线程执行路径上 - 使用 T.Parallel 获取线程索引: 替代
T.get_thread_binding() - 使用内置内存操作:
T.copy(),T.clear()等
调试建议
- UTL 高但 MEM 低: 检查是否有手动归约导致的线程卡死
- 内核不返回: 检查是否有条件分支中的同步导致死锁
- 结果错误: 检查同步点是否正确，数据是否在同步后才使用