Warp 级编程#

Warp 是 CUDA 的基本调度单元：32 个线程在同一个 SM 上以 lockstep 方式执行。由于所有 32 个线程共享同一个指令指针，它们可以通过 warp shuffle 指令直接交换数据——这是寄存器到寄存器的传输，成本大约为一个周期，不需要共享内存、不需要屏障、不需要同步。

cuda-oxide 通过 cuda_device::warp 暴露了完整的 warp 内联函数集。本章介绍 shuffle、vote 以及它们解锁的模式：warp 规约、广播、扫描以及基于 ballot 的过滤。

参见

CUDA Programming Guide — Warp Shuffle Functions 了解 PTX 编码细节和完整的宽度变体集合。

Lane 和 Warp 身份#

Block 中的每个线程在其 warp 中都有一个 lane ID（0–31）：

use cuda_device::warp;

let lane = warp::lane_id();     // 0..31，硬件寄存器 %laneid
let warp = warp::warp_id();     // threadIdx.x / 32

warp_id() 由 threadIdx.x / 32 推导得出。对于多维 block，这仅考虑 x 维度——这通常没问题，因为大多数关心 lane 身份的 kernel 使用 1D block。

Shuffle：寄存器到寄存器的数据交换#

四种 shuffle 变体允许一个线程在不经过内存的情况下读取另一个线程的寄存器：

函数

作用

源 lane

shuffle(val, src)

从指定 lane 读取

src

shuffle_xor(val, mask)

从 lane_id ^ mask 读取

lane_id ^ mask

shuffle_down(val, delta)

从 lane_id + delta 读取

lane_id + delta

shuffle_up(val, delta)

从 lane_id - delta 读取

lane_id - delta

每种变体都有 u32 和 f32 两个版本：

let partner_val = warp::shuffle_xor_f32(my_val, 1);
let broadcast   = warp::shuffle_f32(my_val, 0);     // lane 0 的值广播到所有 lane
let neighbor    = warp::shuffle_down_f32(my_val, 1); // 下一个 lane 的值

所有 shuffle 操作都是warp 同步的——它们隐式地同步了 warp。不需要 sync_threads()，实际上在 shuffle 模式中调用 sync_threads() 既无必要又浪费。

Warp 规约#

最常见的 shuffle 模式是蝶形规约（butterfly reduction）：在 ⌈log₂(32)⌉ = 5 个步骤中，每个 lane 累加所有 32 个值的总和（或最小值、最大值等）。不需要共享内存，不需要屏障，只需五条指令。

advanced/images/warp-shuffle-reduction.svg

使用 shuffle_xor 的蝶形规约。每一步，各 lane 与其 XOR 伙伴交换值并相加。经过 5 步（掩码 16, 8, 4, 2, 1）后，lane 0 持有全部 32 个值的总和。#

use cuda_device::warp;

fn warp_reduce_sum(mut val: f32) -> f32 {
    val += warp::shuffle_xor_f32(val, 16);
    val += warp::shuffle_xor_f32(val, 8);
    val += warp::shuffle_xor_f32(val, 4);
    val += warp::shuffle_xor_f32(val, 2);
    val += warp::shuffle_xor_f32(val, 1);
    val
}

规约完成后，全部 32 个 lane 都持有总和（因为 XOR 是对称的——双方都累加）。如果你只需要在 lane 0 中获得结果，可以使用 shuffle_down 替代：

fn warp_reduce_sum_lane0(mut val: f32) -> f32 {
    val += warp::shuffle_down_f32(val, 16);
    val += warp::shuffle_down_f32(val, 8);
    val += warp::shuffle_down_f32(val, 4);
    val += warp::shuffle_down_f32(val, 2);
    val += warp::shuffle_down_f32(val, 1);
    val
}

使用 shuffle_down 时，只有 lane 0 持有正确的结果——其他 lane 持有部分和。当只有 lane 0 写输出时，这是可以的。

小技巧

需要 block 级规约？在每个 warp 内使用 shuffle 进行规约，将 32 个每 warp 的结果写入共享内存，sync_threads()，然后用最后一次 warp 规约来规约 warp 级结果。这种混合方法比纯共享内存树更快，因为它消除了 5 层屏障。

广播#

将 lane 0 的值广播到所有 lane 只需要一次 shuffle：

let leader_val = warp::shuffle_f32(my_val, 0);

任何 lane 都可以作为源。这取代了共享内存模式中的"lane 0 写入共享内存，同步，所有 lane 读取"——一条指令替代了三个操作。

包含前缀和（扫描）#

**包含式扫描（inclusive scan）**计算运行总计：lane i 持有从 lane 0 到 lane i 的值之和。该模式使用 shuffle_up：

fn warp_inclusive_scan(mut val: f32) -> f32 {
    let mut offset = 1u32;
    while offset < 32 {
        let n = warp::shuffle_up_f32(val, offset);
        if warp::lane_id() >= offset {
            val += n;
        }
        offset *= 2;
    }
    val
}

经过 5 步之后，每个 lane 持有直到并包括其自身值在内的前缀和。这是流压缩（stream compaction）、直方图构建以及并行扫描算法的基础。

Vote：Warp 级谓词#

Vote 操作允许 warp 集体评估布尔条件：

函数

返回值

warp::all(pred)

如果每个活跃 lane 的 pred 都为 true 则返回 true

warp::any(pred)

如果任意活跃 lane 的 pred 为 true 则返回 true

warp::ballot(pred)

一个 u32 位掩码——如果 lane i 的 pred 为 true，则 bit i 被置位

warp::popc(pred)

人口计数：有多少个活跃 lane 的 pred == true

使用 ballot 进行过滤#

一种常见模式是压缩数组，只保留通过谓词的元素。ballot + popc 提供计数和每 lane 的写入偏移量：

use cuda_device::{kernel, thread, warp, DisjointSlice};

#[kernel]
pub fn compact_positive(
    input: &[f32],
    mut output: DisjointSlice<f32>,
    mut count: DisjointSlice<u32>,
) {
    let idx = thread::index_1d();
    let val = input[idx.get()];
    let is_positive = val > 0.0;

    let mask = warp::ballot(is_positive);
    let lane = warp::lane_id();

    // 计算该 lane 以下的置位 bit 数，得到写入位置
    let offset = (mask & ((1u32 << lane) - 1)).count_ones();

    if is_positive {
        unsafe {
            *output.get_unchecked_mut(offset as usize) = val;
        }
    }

    // Lane 0 记录该 warp 的总计数
    if lane == 0 {
        unsafe {
            *count.get_unchecked_mut(warp::warp_id() as usize) = mask.count_ones();
        }
    }
}

ballot 掩码在一个寄存器中编码了整个 warp 的谓词结果。不需要通信，不需要共享内存——硬件在单个周期内计算完成。

何时使用 warp 原语 vs. 共享内存#

任务

Warp shuffle

共享内存

规约 32 个值

5 次 shuffle，~5 cycles

5 次加载 + 5 次同步，~50 cycles

规约 256 个值

Shuffle + 1 次同步 + shuffle

树形规约，~10 次同步

模板（邻域访问）

shuffle_up/shuffle_down

适用于 2D 模板

数据对其他 warp 可见

不可能

必需

随机访问模式

不支持

自由索引

跨 warp 持久化

不适用

在 block 生命周期内持久化

经验法则：如果数据适合单个 warp（32 个元素）且访问模式是规则的，shuffle 更快。如果需要跨 warp 通信、更大的数据集或随机访问，共享内存是合适的工具。

完整示例：Warp 级点积#

将 shuffle 和 vote 结合起来，下面是一个使用 warp 规约计算两个向量点积的 kernel：

use cuda_device::{kernel, thread, warp, DisjointSlice};

#[kernel]
pub fn warp_dot_product(
    a: &[f32],
    b: &[f32],
    n: u32,
    mut result: DisjointSlice<f32>,
) {
    let idx = thread::index_1d();

    // 每个线程计算逐元素乘积的一个元素
    let product = if idx.get() < n as usize {
        a[idx.get()] * b[idx.get()]
    } else {
        0.0f32
    };

    // Warp 级规约
    let mut sum = product;
    sum += warp::shuffle_xor_f32(sum, 16);
    sum += warp::shuffle_xor_f32(sum, 8);
    sum += warp::shuffle_xor_f32(sum, 4);
    sum += warp::shuffle_xor_f32(sum, 2);
    sum += warp::shuffle_xor_f32(sum, 1);

    // 每个 warp 的 lane 0 写入其部分和
    if warp::lane_id() == 0 {
        unsafe {
            *result.get_unchecked_mut(warp::warp_id() as usize) = sum;
        }
    }
}

对于完整的点积，启动第二趟（pass）来规约每个 warp 的结果——可以再用一个 warp kernel 或使用原子操作。第一趟仅使用寄存器 shuffle 就消除了绝大部分工作。

参见