集群编程

线程块集群（Thread Block Cluster） 是一组保证在同一 GPC（Graphics Processing Cluster）内的 SM 上同时运行的线程块。集群在 Hopper（SM 90）中引入，它提供了标准 CUDA 所没有的能力：跨 block 共享内存访问，无需经过全局内存。

在普通的 CUDA 中，线程块是相互独立的。Block 0 无法读取 block 1 的共享内存。有了集群后，硬件将所有集群成员的共享内存映射到一个统一的**分布式共享内存（DSMEM）**地址空间。集群中的任何线程都可以读取任何 block 的共享内存——直接访问，共享内存延迟，无需全局内存往返。

cuda-oxide 通过 cuda_device::cluster 和 #[cluster_launch] 属性暴露集群编程。本章介绍编程模型、DSMEM 访问模式，以及集群如何与 TMA 多播结合以支持高性能矩阵内核。

参见

CUDA Programming Guide — Thread Block Clusters 了解硬件规格、GPC 约束以及最大集群大小。

---

集群模型

一个 4-block 的集群。每个 block 有自己的共享内存，但硬件将所有四个映射到一个统一的 DSMEM 空间。Block 可以在 cluster_sync() 屏障之后读取彼此的共享内存。虚线箭头表示跨 block 的 DSMEM 读取。

集群在 launch 时定义。硬件将集群中的所有 block 调度到同一 GPC 内的 SM 上，确保低延迟的跨 block 通信。最大集群大小取决于架构——Hopper 支持每个集群最多 8 个 block。

声明集群 kernel

在 cuda-oxide 中，你用 #[cluster_launch] 注解 kernel：

use cuda_device::{kernel, cluster_launch, thread, cluster, SharedArray};

#[kernel]
#[cluster_launch(4, 1, 1)]
pub fn cluster_kernel(/* ... */) {
    // 该 kernel 以 X 维度 4 个 block 的集群运行
    let rank = cluster::block_rank();
    let size = cluster::cluster_size();
    // ...
}

#[cluster_launch(x, y, z)] 属性：

1. 注入一个 __cluster_config<X, Y, Z>() 函数，设置 PTX .reqnctapercluster 指令。

2. 在 host 端，launch 必须使用 launch_kernel_ex 并带有匹配的 cluster_dim 参数。

---

集群身份

每个线程都可以查询其在集群层次结构中的位置：

use cuda_device::cluster;

let rank = cluster::block_rank();       // 0..cluster_size()-1，在该集群内的编号
let size = cluster::cluster_size();     // 集群中的总 block 数
let cidx = cluster::cluster_idx();      // 在 grid 中的哪个集群
let ncls = cluster::num_clusters();     // grid 中的总集群数

block_rank() 是 DSMEM 访问的关键标识符——它告诉你应该读取哪个 block 的共享内存。可以把它看作集群内的"block 本地 device ID"。

对于完整的 3D 集群位置：

let cx = cluster::cluster_ctaidX();     // block 在集群中的 X 位置
let cy = cluster::cluster_ctaidY();     // block 在集群中的 Y 位置
let cz = cluster::cluster_ctaidZ();     // block 在集群中的 Z 位置

---

分布式共享内存（DSMEM）

集群的核心特性是跨 block 共享内存访问。有两种方式可以读取远程 block 的共享内存：

方法 1：map_shared_rank（指针重映射）

map_shared_rank 接受一个本地共享内存指针，并返回指向另一个 block 的共享内存中相同偏移量的指针：

use cuda_device::{cluster, SharedArray, thread};

static mut DATA: SharedArray<u32, 256> = SharedArray::UNINIT;

#[kernel]
#[cluster_launch(4, 1, 1)]
pub fn halo_exchange(/* ... */) {
    let tid = thread::threadIdx_x() as usize;
    let rank = cluster::block_rank();

    // 每个 block 写入自己的共享内存
    unsafe { DATA[tid] = compute_value(rank, tid); }

    thread::sync_threads();    // 首先本地屏障
    cluster::cluster_sync();   // 然后集群级屏障

    // 从下一个 block 的共享内存读取
    let neighbor = (rank + 1) % cluster::cluster_size();
    let remote_ptr = unsafe {
        cluster::map_shared_rank(DATA.as_ptr().add(tid), neighbor)
    };
    let neighbor_val = unsafe { *remote_ptr };
}

方法 2：dsmem_read_u32（推荐）

dsmem_read_u32 是读取远程共享内存的推荐方式。它编译为 ld.shared::cluster PTX 指令，硬件处理效率比重映射指针的通用加载更高：

let neighbor_val = unsafe {
    cluster::dsmem_read_u32(
        DATA.as_ptr() as *const u32,
        neighbor,  // 目标 rank
    )
};

差异虽然微妙，但对性能很重要：map_shared_rank 生成一个通过通用加载路径的指针，而 dsmem_read_u32 使用硬件可以优化的专用指令。对于 u32 大小的读取，使用 dsmem_read_u32；对于更大的类型，使用带有适当指针类型的 map_shared_rank 即可。

---

同步

集群在 block 同步和全局同步之间引入了一个新的同步级别：

原语	作用域	使用场景
thread::sync_threads()	Block	在单个 block 内同步
cluster::cluster_sync()	集群	在集群内的所有 block 之间同步
mbarrier_arrive_cluster(addr)	集群	向远程 block 中的屏障发出信号

DSMEM 访问的正确同步顺序始终是：

1. 写入 本地共享内存

2. sync_threads() —— 确保所有本地线程已完成写入

3. cluster_sync() —— 确保所有 block 已到达此点

4. 读取 通过 DSMEM 从远程共享内存读取

缺少任何一个同步都是数据竞争。有 sync_threads() 而没有 cluster_sync() 意味着你的 block 已就绪，但邻居可能尚未就绪。有 cluster_sync() 而没有 sync_threads() 意味着集群已同步，但你自己的 block 的写入可能尚不可见。

---

TMA 多播与集群

集群解锁了 TMA 最强大的特性之一：多播拷贝。一次 TMA 加载可以将同一个分块同时写入多个 block 的共享内存：

use cuda_device::tma::cp_async_bulk_tensor_2d_g2s_multicast;

// CTA 掩码：bit 0..3 置位 → 集群中所有 4 个 block 都接收该分块
let cta_mask: u16 = 0b1111;

unsafe {
    cp_async_bulk_tensor_2d_g2s_multicast(
        smem_dst, desc, tile_x, tile_y, bar_ptr, cta_mask
    );
}

没有多播时，每个 block 会发出自己的 TMA 拷贝——四次独立的全局内存读取。有了多播，TMA 引擎读取一次数据，并将其分发给所有四个 block。这对于 GEMM 内核特别有价值，其中 A 或 B 的同一个分块被集群中的每个 block 所需要。

cta_mask 是一个位掩码，其中 bit i 被置位时表示 rank i 应当接收该拷贝。你可以选择性地多播到集群的一个子集。

---

实践示例：Halo 交换

集群的一个常见用例是模板计算中的 halo 交换。每个 block 处理网格的一个分块，需要从其邻居获取边界元素。没有集群时，这需要全局内存写入和读取（或小心的 stream 同步）。有了集群，这变成了本地操作：

use cuda_device::{kernel, cluster_launch, thread, cluster, SharedArray, DisjointSlice};

const TILE_W: usize = 256;
const HALO: usize = 1;
const SMEM_W: usize = TILE_W + 2 * HALO;

#[kernel]
#[cluster_launch(4, 1, 1)]
pub fn stencil_1d(input: &[f32], mut output: DisjointSlice<f32>, n: u32) {
    static mut SMEM: SharedArray<f32, { SMEM_W }> = SharedArray::UNINIT;

    let tid = thread::threadIdx_x() as usize;
    let rank = cluster::block_rank();
    let global_idx = rank as usize * TILE_W + tid;

    // 加载内部数据（偏移 HALO 以留出 halo 插槽）
    unsafe {
        if global_idx < n as usize {
            SMEM[tid + HALO] = input[global_idx];
        }
    }

    thread::sync_threads();
    cluster::cluster_sync();

    // 从前一个 block 加载左 halo
    if tid == 0 && rank > 0 {
        let prev_rank = rank - 1;
        unsafe {
            SMEM[0] = f32::from_bits(cluster::dsmem_read_u32(
                SMEM.as_ptr().add(TILE_W) as *const u32, prev_rank
            ));
        }
    }

    // 从后一个 block 加载右 halo
    if tid == 0 && rank < cluster::cluster_size() - 1 {
        let next_rank = rank + 1;
        unsafe {
            SMEM[TILE_W + HALO] = f32::from_bits(cluster::dsmem_read_u32(
                SMEM.as_ptr().add(HALO) as *const u32, next_rank
            ));
        }
    }

    thread::sync_threads();

    // 3 点模板：output[i] = 0.25 * left + 0.5 * center + 0.25 * right
    if global_idx < n as usize {
        let left = unsafe { SMEM[tid + HALO - 1] };
        let center = unsafe { SMEM[tid + HALO] };
        let right = unsafe { SMEM[tid + HALO + 1] };

        unsafe {
            *output.get_unchecked_mut(global_idx) = 0.25 * left + 0.5 * center + 0.25 * right;
        }
    }
}

如果没有集群，halo 交换将需要将边界元素写入全局内存、通过事件或 stream 同步，然后读回。有了集群，就是一次 cluster_sync() 加一次 dsmem_read_u32——共享内存延迟，无需全局内存。

---

约束与最佳实践

约束	详情
最大集群大小	8 个 block（取决于架构）
调度保证	所有集群 block 在同一 GPC 上同时运行
DSMEM 延迟	类似于本地共享内存（~5–10 cycles）
集群维度必须声明	#[cluster_launch(x, y, z)] + launch 配置中的 host cluster_dim
Block 数量必须可整除	每个维度的 grid block 必须是集群大小的倍数
混合集群/非集群	同一 kernel 中不支持

小技巧

集群施加了调度约束：硬件必须将所有 block 放置在同一 GPC 上。如果集群相对于 GPC 容量过大，占用率会下降。从较小的集群（2–4 个 block）开始，在扩展之前先进行测量。

参见

• Tensor Memory Accelerator — TMA 多播需要集群才能跨 CTA 分发

• 矩阵乘法加速器 — CG2 模式使用集群对来实现更宽的 MMA 分块

• 共享内存与同步 — DSMEM 所扩展的每 block 基础