tilelang-cuda-synchronization

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)    # 平均值归约