cuda编程笔记(39)--Asynchronous Barriers(异步屏障)

本文需要有cooperative_groups的API基础,可以见我以前的文章,虽然讲解的不太好,以后有时间重置一下。cuda编程笔记(35)-- Cooperative Groups_cuda grid sync-CSDN博客

可以将 cuda::barrier 理解为 __syncthreads() 的"进化增强版"。传统的 __syncthreads()阻塞式全员参与 的,而异步屏障则允许解耦(Split)"到达"和"等待"这两个动作。

在传统的 CUDA 同步中,调用 __syncthreads() 后,线程必须停下来等待所有线程到齐。 异步屏障的核心优势在于:

  • 非阻塞协调: 线程可以先声明"我到了"(Arrive),然后去执行不依赖同步结果的代码,最后再检查"大家都到了没"(Wait)。

  • 计算与通信重叠: 在等待数据传输(如异步内存拷贝)的过程中,GPU 可以继续进行计算,从而隐藏延迟。

  • 粒度更细: 不再局限于整个 Block,可以灵活指定参与同步的线程数量。

初始化

cpp 复制代码
#include <cuda/barrier>
#include <cooperative_groups.h>

__global__ void init_barrier()
{
  // 1. 定义一个共享内存中的屏障对象
  // 作用域设为 cuda::thread_scope_block,表示该屏障用于块内同步
  __shared__ cuda::barrier<cuda::thread_scope_block> bar;
  
  auto block = cooperative_groups::this_thread_block();

  // 2. 初始化屏障
  if (block.thread_rank() == 0) 
  {
    // 由 0 号线程负责调用 init 函数
    // 第二个参数 block.size() 是 "expected arrival count"(期望到达数)
    // 意味着屏障必须收到 block.size() 次 arrive 调用,才会解除阻塞
    init(&bar, block.size());
  }

  // 3. 关键:引导同步 (Bootstrapping)
  // 在屏障初始化完成之前,任何线程都不能调用 bar.arrive()
  // 所以必须先用传统的 __syncthreads() 或 block.sync() 挡一下
  block.sync(); 
}

初始化屏障也需要同步。这是一个"鸡生蛋"的问题:我们为了同步才创建屏障,但屏障本身的初始化也需要同步。官方建议在初始化 cuda::barrier 之后,必须使用传统的 __syncthreads() 确保所有线程都看到了初始化完成的状态,之后才能开始使用该屏障。

cuda::barrier

cpp 复制代码
template <thread_scope _Sco, class _CompletionF = _CUDA_VSTD::__empty_completion>
class barrier : public _CUDA_VSTD::__barrier_base<_CompletionF, _Sco>;

cuda::barrier 是 CUDA 在 C++ 标准屏障(std::barrier)基础上进行的扩展。它的作用是让一组线程在某个点汇合,并确保在跨过这个点时,之前的内存操作对所有参与线程都可见。

_Sco 定义了该屏障的内存一致性范围

当你写 cuda::barrier<cuda::thread_scope_block> bar; 时,你是在告诉 GPU:

"我要创建一个屏障,参与者都在同一个 Block 里,同步时只需要保证块内内存一致性即可,不要去管其他 Block 或 CPU 的缓存。

enum thread_scope

cpp 复制代码
enum thread_scope
{
  thread_scope_system = __ATOMIC_SYSTEM,
  thread_scope_device = __ATOMIC_DEVICE,
  thread_scope_block  = __ATOMIC_BLOCK,
  thread_scope_thread = __ATOMIC_THREAD
};
hread_scope_block (原子操作/屏障最常用)
  • 含义 :同步仅限在同一个 Thread Block(线程块) 内部。

  • 可见性:线程 A 写入共享内存(Shared Memory)或全局内存的数据,在屏障触发后,同一 Block 内的线程 B 保证能看到。

  • 性能:开销最小。因为它只涉及块内的 L1 缓存或 Shared Memory。

thread_scope_device
  • 含义 :同步范围扩展到整个 GPU 设备

  • 可见性:线程 A 在第一个 Grid 中写入的数据,对于同一个 GPU 上运行的其他线程(可能在不同的 Block 里)是可见的。

  • 场景:常用于多线程块之间的协作(Cooperative Groups),确保全局内存(Global Memory)的一致性。

thread_scope_system
  • 含义 :同步范围扩展到整个系统(包括 CPU 和其他 GPU)。

  • 可见性 :GPU 写入的内容,对于通过 PCIe/NVLink 连接的 CPU (Host) 也是可见的。

  • 场景 :用于使用了 cudaMallocManaged(统一内存)或系统分配的内存,且需要 GPU 与 CPU 进行极细粒度同步的场景。

thread_scope_thread
  • 含义:仅限当前线程。

  • 注意:这在屏障里几乎没有意义,因为屏障本质上是多线程协作。在原子操作中,它表示不进行跨线程的同步。

init

cpp 复制代码
void init(barrier* __b, _CUDA_VSTD::ptrdiff_t __expected, 
_CUDA_VSTD::__empty_completion = _CUDA_VSTD::__empty_completion());

count 是核心参数。它告诉硬件:当有多少个"到达信号"产生时,这个屏障才算达成。在示例中设为 block.size(),意味着整个线程块的每个线程都要贡献一个"到达"。

特别提到:如果你只是想简单地对整个线程块或整个 Warp 进行同步,依然推荐使用传统的 __syncthreads()__syncwarp(),因为它们经过硬件高度优化,在简单同步场景下性能更好。

屏障的阶段:到达、倒计时、完成与重置

异步屏障通过参与线程调用 bar.arrive(),将"预期到达计数"向下减至零。当倒计时归零时,屏障在该当前阶段(Phase)即告完成。当最后一次 bar.arrive() 调用使计数归零时,计数器会自动且原子地重置------将计数重置回初始的预期值,并将屏障推进到下一个阶段。

调用 token = bar.arrive() 返回的 cuda::barrier::arrival_token 类对象,与屏障的当前阶段 相关联。调用 bar.wait(std::move(token))阻塞调用线程前提是屏障仍处于该 token 所对应的阶段 。如果在调用 wait 之前阶段已经推进(因为计数归零),则线程不会阻塞;

如果线程在 wait 中阻塞时阶段发生推进,线程将被唤醒。

了解重置发生的时机至关重要,尤其是在复杂的"到达/等待"同步模式中:

  • 线程调用 bar.arrive() 必须在屏障的当前阶段进行。

  • bar.wait(std::move(token)) 必须在相同阶段或紧接着的下一个阶段进行。

  • 线程调用 bar.arrive() 时,屏障的计数器必须为非零。初始化后,如果某个线程的 arrive() 导致计数归零,那么在屏障被用于后续的 arrive() 调用之前,必须先执行 wait()

  • bar.wait() 只能使用当前阶段或直接前一个阶段的 token。使用任何其他值的 token 会导致未定义行为。

Warp 纠缠(Warp Entanglement)

Warp 分歧(Warp-divergence)会影响屏障更新的次数。

  • 理想情况 :一个 Warp(32 线程)一起执行 bar.arrive()。硬件非常聪明,它会发现"哦,你们这 32 个人是一伙的",然后一次性把屏障计数减 32。这只产生 1 次 硬件原子操作开销。

  • 糟糕情况 :如果你的代码里有 if-else 导致 Warp 里的线程各跑各的,或者你没写 __syncwarp()。硬件就得老老实实处理 32 次 原子操作减 1。

  • 结论 :在调用 arrive 之前,随手写个 __syncwarp() 是极好的习惯。

示例代码

cpp 复制代码
// 1. 到达:拿票,然后该干嘛干嘛
// 这里屏障会自动减少计数
auto token = bar.arrive(); 

// 2. 重叠(Overlap):这里可以写几百行跟同步无关的代码
// 比如从 Global Memory 异步搬运下一轮的数据
do_independent_work(); 

// 3. 等待:凭票入场
bar.wait(std::move(token));

barrier::arrive

cpp 复制代码
using barrier<thread_scope_block>::arrival_token = uint64_t;

arrival_token arrive(_CUDA_VSTD::ptrdiff_t __update = 1);

__update 参数实际上是异步屏障灵活性的体现。简单来说,它允许一个线程代表多个"到场名额"进行签到。

如果一个 Warp 里的 32 个线程都调用 arrive(),默认会产生多次原子减法。 如果你能确定这 32 个线程都已经完成了任务,你可以只让其中一个线程 代表整个 Warp 调用一次 arrive(32)

barrier::wait

cpp 复制代码
void wait(arrival_token&& __phase) const;

只要屏障的计数(Countdown)还没有减到 0,屏障就还停留在该 token 所代表的"当前阶段(Current Phase)"。此时调用 wait(),线程就会进入阻塞状态

显式阶段跟踪 (Explicit Phase Tracking)

异步屏障可以根据线程同步和内存操作的次数拥有多个"阶段(Phases)"。除了使用 token(令牌)来跟踪屏障的阶段翻转外,我们还可以通过 cuda::ptx 提供的 mbarrier_try_wait_parity() 系列函数来直接跟踪阶段。

在最简单的形式中,cuda::ptx::mbarrier_try_wait_parity(uint64_t* bar, const uint32_t& phaseParity) 函数会等待具有特定"极性(Parity)"的阶段。phaseParity 操作数是屏障对象当前阶段或紧接前一个阶段的整数极性。偶数阶段的整数极性为 0,奇数阶段为 1。 当我们初始化屏障时,其初始阶段的极性为 0。因此,phaseParity 的有效值为 0 和 1。

当跟踪异步内存操作 时,显式阶段跟踪非常有用。它允许仅由单个线程到达(Arrive)屏障并设置事务计数(Transaction Count),而其他线程只需等待基于极性的阶段翻转。这比让所有线程都到达屏障并使用 token 更高效。此功能仅适用于线程块(Thread-block)和集群(Cluster)作用域的共享内存屏障。

在之前的章节中,我们使用 token = bar.arrive() 拿"票",然后用 bar.wait(token) 凭"票"入场。而显式阶段跟踪是一种更底层的控制方式:

  1. 极性循环 (Parity Cycle): 屏障的状态像是一个开关。第一轮(阶段 0)极性是 0,第二轮(阶段 1)极性是 1,第三轮又回到 0。你不再需要保存 token 对象,只需要记住现在是"开"还是"关"。

  2. 更高效的协作: 在复杂的异步拷贝(如 cp.async)中,可能只有部分线程负责发起搬运和"签到(Arrive)"。其他线程如果不参与签到,就没有 token。通过"极性",这些没参与签到的线程也可以通过观察屏障的"极性翻转"来判断数据是否搬运完成。

  3. 底层控制: 这里使用了 cuda::ptx 命名空间。这实际上是直接调用了硬件层面的 mbarrier 指令。这比高级 C++ API 更接近底层,开销更小。

代码示例

这段代码展示了如何在一个循环中使用"极性(Parity)"来手动控制同步。

cpp 复制代码
#include <cuda/ptx>
#include <cooperative_groups.h>

__device__ void compute(float *data, int iteration);

__global__ void split_arrive_wait(int iteration_count, float *data)
{
    using barrier_t = cuda::barrier<cuda::thread_scope_block>;
    
    // 在共享内存中定义屏障
    __shared__ barrier_t bar;
    
    // 初始化极性:初始阶段极性为 0
    int parity = 0; 
    
    auto block = cooperative_groups::this_thread_block();

    // 1. 初始化屏障
    if (block.thread_rank() == 0)
    {
        // 设置期望到达数为整个 Thread Block 的大小
        init(&bar, block.size());
    }
    // 确保屏障初始化完成且所有线程都已到齐
    block.sync();

    // 2. 循环同步
    for (int i = 0; i < iteration_count; ++i)
    {
        /* 到达前的操作 (例如:准备数据) */

        // 线程到达:这里使用了 PTX 底层接口
        // cuda::device::barrier_native_handle(bar) 获取屏障的底层原生句柄
        cuda::ptx::mbarrier_arrive(cuda::device::barrier_native_handle(bar));

        // 异步计算:在等待其他线程到齐的同时,可以做一些不依赖同步的计算
        compute(data, i);

        // 3. 等待阶段翻转
        // 使用 mbarrier_try_wait_parity 进行轮询
        // 如果屏障还没有翻转到下一个极性,此函数返回 false
        while (!cuda::ptx::mbarrier_try_wait_parity(
            cuda::device::barrier_native_handle(bar), 
            parity)
        ) {
            // 在循环中轮询,直到所有线程都 arrive,导致极性翻转
        }

        // 4. 翻转本地极性:为下一轮循环做准备
        // 0 变 1, 1 变 0
        parity ^= 1;

        /* 等待后的操作 (此时可以安全使用同步后的数据) */
    }
}
  • while(!...try_wait_parity) : 不同于 bar.wait() 会让线程挂起(yield),try_wait_parity 是一种非阻塞的检查。在代码中,我们手动写了一个 while 循环来实现阻塞等待。这种方式在微调性能时非常有用。

  • parity ^= 1 : 这一步至关重要。如果你的本地 parity 和屏障内部的 parity 没同步,下一轮循环你就会直接跳过等待,导致数据竞争。

提前退出 (Early Exit)

当一个参与了一系列同步序列的线程必须提前退出该序列时,该线程在退出前必须显式地取消参与(Drop out) 。这样,剩余的参与线程才能正常进行后续的 arrivewait 操作。

bar.arrive_and_drop() 操作会触发当前阶段的"到达"信号,以履行该线程在当前阶段 的到达义务;随后,它会减少下一阶段及后续所有阶段的"预期到达总数(Expected Arrival Count)"。这样,在后续的同步中,系统将不再期待该线程到达屏障。

在 CUDA 编程中,异步屏障是一个"严进严出"的机制。如果你在初始化时告诉屏障:"全块 256 个线程都会来报到",那么屏障的计数器就会死死地等够 256 次 arrive 才会翻转。

如果不调用 arrive_and_drop() 就直接 return

  1. 当前轮次死锁 :剩下的 255 个线程调用了 arrive,但计数器停在 1,所有人都在 wait 处永久等待。

  2. 计数器污染:即使当前轮次勉强通过了,下一轮屏障依然在等 256 个人。由于那个线程已经退出了(Kernel 结束或跳出循环),屏障将永远等不到最后一个人,导致整个 GPU 核心挂起。

arrive_and_drop() 的两个动作:

  • 动作 A(对当下负责):给当前计数器减 1。相当于说:"这轮我赶上了,大家别等我,你们先过。"

  • 动作 B(对未来负责) :将屏障内部永久维护的 expected count 减 1。相当于说:"后面几轮我也没假了,以后别算我的人头。"

这种提前退出的情况也很常见,典型例子就是Block的线程数大于数组大小,有的线程就可以提前return。

示例代码

cpp 复制代码
#include <cuda/barrier>
#include <cooperative_groups.h>

__device__ bool condition_check(); // 某种判断是否需要提前退出的逻辑

__global__ void early_exit_kernel(int N)
{
  // 1. 定义共享内存屏障
  __shared__ cuda::barrier<cuda::thread_scope_block> bar;
  auto block = cooperative_groups::this_thread_block();

  // 2. 初始化:初始期望到达数为 block.size() (例如 128)
  if (block.thread_rank() == 0)
  {
    init(&bar, block.size());
  }
  // 引导同步:确保所有线程都看到了初始化的屏障
  block.sync();

  for (int i = 0; i < N; ++i)
  {
    // 3. 检查退出条件(例如:计算已收敛、发生错误、或该线程分配的任务已完成)
    if (condition_check())
    {
      // 【核心】:我要走了,帮我把当前计数减 1,并且把以后每轮的期望数也减 1
      // 执行完这一行后,下一次循环屏障的 expected count 就会自动变为 127
      bar.arrive_and_drop();
      return; // 安全退出线程
    }

    // 4. 正常参与的线程执行常规流程
    // 剩下的线程继续正常工作
    auto token = bar.arrive();

    /* 在 arrive 和 wait 之间可以执行一些计算 */

    // 等待本轮所有还在参与的线程到齐
    bar.wait(std::move(token));

    /* wait 之后的处理逻辑 */
  }
}

完成函数 (Completion Function)

cuda::barrier API 支持一个可选的完成函数 。对于 cuda::barrier<Scope, CompletionFunction> 类型的屏障,其 CompletionFunction 在每一轮阶段中仅执行一次。执行时机是在最后一个线程"到达"(Arrive)之后,且在任何线程从"等待"(Wait)状态被唤醒之前。

在当前阶段中,所有已到达屏障的线程所执行的内存操作,对于执行 CompletionFunction 的线程是可见的;而 CompletionFunction 内部执行的所有内存操作,在所有等待线程被唤醒后,对它们也是可见的。

串行语义 :虽然有 128 个线程在跑,但 CompletionFunction 只由其中一个线程 执行一次。你不需要在逻辑里写 if(tid == 0),硬件和驱动会自动帮你处理。

代码示例

这个例子展示了如何利用完成函数在每一轮迭代中自动进行一次"归约(Reduction)"操作。

cpp 复制代码
#include <cuda/barrier>
#include <cooperative_groups.h>
#include <functional>

namespace cg = cooperative_groups;

__global__ void psum(int *data, int n, int *acc)
{
  auto block = cg::this_thread_block();
  constexpr int BlockSize = 128;
  __shared__ int smem[BlockSize]; // 共享内存缓存

  // --- 1. 定义完成函数 ---
  // 这个 Lambda 函数定义了当所有人到达屏障后,该做什么汇总工作
  auto completion_fn = [&] 
  {
    int sum = 0;
    for (int i = 0; i < BlockSize; ++i)
    {
      sum += smem[i]; // 将本轮所有线程写入 smem 的值累加
    }
    *acc += sum; // 更新全局累加器
  };

  // --- 2. 屏障类型定义与内存准备 ---
  using completion_fn_t = decltype(completion_fn);
  using barrier_t = cuda::barrier<cuda::thread_scope_block, completion_fn_t>;

  // 由于 completion_fn 捕获了外部变量,barrier 无法自动构造
  // 我们在共享内存中开辟一块对齐的地址空间
  __shared__ std::aligned_storage_t<sizeof(barrier_t), alignof(barrier_t)> bar_storage;
  
  // 将地址转换为屏障指针
  barrier_t *bar = reinterpret_cast<barrier_t *>(&bar_storage);

  // --- 3. 手动初始化 ---
  if (block.thread_rank() == 0)
  {
    // 使用 placement new 在指定共享内存位置构造屏障对象
    // 并传入 block 大小和完成函数对象
    new (bar) barrier_t{block.size(), completion_fn};
    /* 相当于: init(bar, block.size(), completion_fn); */
  }
  
  // 必须进行一次显式同步,确保屏障对象在所有线程开始使用前已构造完毕
  block.sync();

  // --- 4. 主循环 ---
  for (int i = 0; i < n; i += block.size())
  {
    // 每个线程负责将数据从全局内存搬运到共享内存
    // 这里使用了 *acc,它是上一轮 completion_fn 计算出的结果
    smem[block.thread_rank()] = data[i] + *acc;

    // 线程宣告到达
    auto token = bar->arrive();

    // 【此处可以执行与 smem 无关的独立计算任务】

    // 线程等待
    // 在 wait 返回之前,硬件会自动挑选一个线程执行 completion_fn
    bar->wait(std::move(token));

    // 当线程运行到这里时,smem 里的数据已经可以安全地被下一轮循环覆盖了
    // 因为 completion_fn 已经读取完毕,且所有线程都通过了 wait
  }
}

为什么一定要通过原始内存构造barrier?

1. 构造函数的"调用权"问题

在 CUDA 中,如果你声明一个普通的共享内存变量:

cpp 复制代码
__shared__ cuda::barrier<...> bar;

编译器会自动尝试在每个线程进入内核时都去调用它的默认构造函数。

如果这个屏障带了 Completion Function(完成函数),情况会变成这样:

  • 这个函数通常是一个 Lambda 表达式,它捕获了当前的上下文(比如引用或指针)。

  • 这类对象没有默认构造函数(Non-default-constructible)。

  • 编译器不知道该给它传什么参数,因此无法直接通过 __shared__ 这种声明式语法完成初始化。

2. 内存对齐与大小控制

cuda::barrier 内部可能包含硬件特定的对齐要求。

  • 使用 std::aligned_storage 可以确保我们在共享内存中开辟的空间,不仅大小足够(sizeof),而且内存起始地址是对齐的alignof)。

  • 如果对齐不正确,访问屏障时可能会触发非法内存访问或性能大幅下降。

3. 灵活控制初始化时机

通过"原始内存 + 手动构造",我们实现了对屏障生命周期的手术刀式精准控制

  1. 预留空间__shared__ bar_storage 仅仅是挖了一个坑。

  2. 单线程执政 :通过 if (block.thread_rank() == 0),我们确保只有0号线程调用构造函数。

  3. 屏障可见性 :初始化后,通过 block.sync() 确保所有线程都看到了这个变量,之后大家才能进去参加同步。

跟踪异步内存操作 (Tracking Asynchronous Memory Operations)

异步屏障可用于跟踪异步内存拷贝 。当异步拷贝操作绑定到屏障时,该拷贝操作在启动时会自动增加屏障当前阶段的"预期计数",并在完成后自动减少该计数。这种机制确保了屏障的 wait() 操作会一直阻塞,直到所有关联的异步内存拷贝全部完成,为同步多个并发内存操作提供了便利。

计算能力 9.0 (Hopper 架构) 开始,位于共享内存中且具有线程块或集群作用域的异步屏障可以显式 跟踪异步内存操作。我们称这些屏障为异步事务屏障 (Asynchronous Transaction Barriers) 。除了"预期到达人数"外,屏障对象还可以接收一个"事务计数 (Transaction Count)"。事务计数跟踪尚未完成的异步事务数量,其单位由具体的异步内存操作指定(通常是字节数)。

当前阶段要跟踪的事务计数可以通过 cuda::device::barrier_arrive_tx() 在到达时设置,或通过 cuda::device::barrier_expect_tx() 直接设置。当屏障使用事务计数时,它会在 wait 操作处阻塞线程,直到:

  1. 所有生产者线程都执行了 arrive

  2. 所有事务计数的总和达到了预期值(即归零)。

这是 CUDA 同步机制的一次重大进化。之前的屏障只关心"人到齐了没",现在的屏障不仅关心人,还关心"货(数据)到了没"。

  1. 双重计数机制

    • Arrival Count (到达计数):传统的线程同步,解决"谁参与"的问题。

    • Transaction Count (事务计数) :新型的数据同步,解决"数据搬完没"的问题。通常以搬运的字节数为单位。

  2. Hopper 架构的杀手锏 : 这是为了配合 TMA (Tensor Memory Accelerator) 硬件设计的。以前搬运数据,CPU/GPU 必须时刻盯着进度;现在你可以告诉屏障:"我要搬 1024 字节",然后屏障硬件会自动监听总线,直到这 1024 字节全部落入共享内存,屏障才放行。

  3. 为什么更高效?

    • 减少指令数 :不需要每个线程搬完后手动 arrive

    • 硬件级监听:由硬件直接检测内存事务的完成情况,完全不占用计算单元(SM)的指令发射周期。

代码示例

cpp 复制代码
#include <cuda/barrier>
#include <cooperative_groups.h>

__global__ void track_kernel()
{
  // 定义共享内存屏障
  __shared__ cuda::barrier<cuda::thread_scope_block> bar;
  auto block = cooperative_groups::this_thread_block();

  // 1. 初始化:设置期望到达的线程数
  if (block.thread_rank() == 0)
  {
    init(&bar, block.size());
  }
  // 引导同步,确保屏障初始化完成
  block.sync();

  // 2. 核心操作:执行带有事务计数的到达 (Arrive with Transaction)
  // 参数说明:
  // bar: 屏障对象
  // 1:  Arrival count 减去的数值(代表当前这 1 个线程到了)
  // 0:  Expected transaction count 增加的数值
  //     这里设为 0,意味着虽然使用了 tx 接口,但实际上并没有要跟踪的异步搬运事务
  auto token = cuda::device::barrier_arrive_tx(bar, 1, 0);

  // 3. 等待
  // 线程会在这里阻塞,直到:
  // a) 所有 block.size() 个线程都调用了 arrive (使 arrival count 归零)
  // b) 所有的事务字节数也都完成了 (由于这里 tx 是 0,所以这一项默认满足)
  bar.wait(cuda::std::move(token));
}

在实际的 TMA异步拷贝 场景中,用法通常是这样的:

  • 生产者线程(发起拷贝的线程):
cpp 复制代码
// 告诉屏障:我到了,而且我还发起了 1024 字节的异步搬运任务
cuda::device::barrier_arrive_tx(bar, 1, 1024);
// 发起真正的异步拷贝 (例如 TMA 指令)

消费者线程

cpp 复制代码
// 消费者只需要拿 token 然后 wait
bar.wait(std::move(token)); 
// 当 wait 返回时,保证 1024 字节的数据已经完整地躺在共享内存里了

cuda::device::barrier_arrive_tx

cpp 复制代码
inline barrier<thread_scope_block>::arrival_token barrier_arrive_tx(
  barrier<thread_scope_block>& __b,
  _CUDA_VSTD::ptrdiff_t __arrive_count_update,
  _CUDA_VSTD::ptrdiff_t __transaction_count_update);

两个数值代表arrive数量和搬运的字节数

该函数不支持SM90之前的架构,所以30系列,4060这种都用不了,我也没办法做测试。

使用屏障的生产者-消费者模式

线程块可以在空间上进行划分,以允许不同的线程执行独立的操作。这通常通过将线程块内不同 Warp(线程束)的线程分配给特定的任务来实现。这种技术被称为Warp 特化 (Warp Specialization)

本节展示了一个在生产者-消费者模式中进行空间划分的示例,其中一个线程子集(生产者)产生数据,而另一个不相交的线程子集(消费者)并发地消耗数据。生产者-消费者空间划分模式需要两次**"单向同步 (One-sided Synchronizations)"**来管理两者之间的数据缓冲区。

生产者 (Producer) 消费者 (Consumer)
等待缓冲区处于"准备好被填充"状态 发出"缓冲区已准备好被填充"的信号
生产数据并填充缓冲区
发出"缓冲区已填满"的信号 等待缓冲区被填满
消耗已填满缓冲区中的数据

生产者线程等待消费者发出"缓冲区准备好被填充"的信号;然而,消费者线程并不等待这个信号 。消费者线程等待生产者发出"缓冲区已填满"的信号;然而,生产者线程并不等待这个信号。为了实现完全的生产者/消费者并发,这种模式使用了(至少)双重缓冲(Double Buffering),其中每个缓冲区需要两个屏障。

在此示例中,第一个 Warp 被特化为生产者,其余的 Warp 被特化为消费者。所有的生产者和消费者线程都参与了全部四个屏障的同步(调用 bar.arrive()bar.arrive_and_wait()),因此预期的到达总数等于 block.size()

生产者线程等待消费者线程发出"共享内存缓冲区可以被填充"的信号。为了等待一个屏障,生产者线程必须先调用 ready[i%2].arrive() 获取令牌,然后使用该令牌调用 ready[i%2].wait(token)。为了简单起见,ready[i%2].arrive_and_wait() 将这两个操作结合在了一起(相当于 bar.wait(bar.arrive()))。

生产者线程计算并填充就绪的缓冲区,然后通过到达"已填充"屏障 filled[i%2].arrive() 来发出缓冲区已满的信号。生产者此时不等待,而是直接去等待下一次迭代的缓冲区(双重缓冲)准备好被填充。

消费者线程首先发出"两个缓冲区都已准备好被填充"的信号。消费者此时不等待 ,而是等待本次迭代的缓冲区被填满,即调用 filled[i%2].arrive_and_wait()。在消费者消耗完缓冲区的数据后,它们发出信号表示缓冲区又可以再次被填充了 ready[i%2].arrive(),然后等待下一次迭代的缓冲区被填满。

为什么需要 4 个屏障?

因为使用了双缓冲 (Ping-Pong Buffer)

有两个缓冲区(Buffer 0 和 Buffer 1),每个缓冲区都需要跟踪两种状态(空、满),所以需要 4 个屏障:

  • bar[0]:Buffer 0 空了(Ready)

  • bar[1]:Buffer 1 空了(Ready)

  • bar[2]:Buffer 0 满了(Filled)

  • bar[3]:Buffer 1 满了(Filled)

示例代码

cpp 复制代码
#include <cuda/barrier>
#include <cooperative_groups.h>
#include <cuda_runtime.h>
#include <iostream>
#include <vector>

using barrier_t = cuda::barrier<cuda::thread_scope_block>;

// ==========================================
// 生产者逻辑 (由 Warp 0 执行)
// ==========================================
__device__ void produce(barrier_t ready[], barrier_t filled[], float *buffer, int buffer_len, float *in, int N) {
    int warp_tid = threadIdx.x % 32; // 在当前 Warp 内的编号

    for (int i = 0; i < N / buffer_len; ++i) {
        // 1. 等待消费者信号:Buffer i%2 是否为空
        ready[i % 2].arrive_and_wait();

        // 2. 生产:将数据从全局内存搬到共享内存
        // 只有 Warp 0 参与,步长为 32
        float* current_smem = buffer + (i % 2) * buffer_len;
        float* current_in   = in + i * buffer_len;

        for (int j = warp_tid; j < buffer_len; j += 32) {
            current_smem[j] = current_in[j];
        }

        // 3. 通知消费者:Buffer i%2 已填满
        // 注意:即使这里只有 Warp 0 干活,但屏障初始化时是整个 Block 的 size,
        // 所以整个 Block 的线程都必须在各自的逻辑分支里调用 arrive/wait。
        filled[i % 2].arrive();
    }
}

// ==========================================
// 消费者逻辑 (由其他 Warps 执行)
// ==========================================
__device__ void consume(barrier_t ready[], barrier_t filled[], float *buffer, int buffer_len, float *out, int N) {
    int consumer_tid = threadIdx.x - 32; // 消费者线程的相对编号
    int num_consumers = blockDim.x - 32; // 消费者总数

    // 1. 初始通知:告诉生产者所有 Buffer 都是空的
    ready[0].arrive();
    ready[1].arrive();

    for (int i = 0; i < N / buffer_len; ++i) {
        // 2. 等待生产者信号:Buffer i%2 是否已填满
        filled[i % 2].arrive_and_wait();

        // 3. 消费:处理数据并写回全局内存
        float* current_smem = buffer + (i % 2) * buffer_len;
        float* current_out  = out + i * buffer_len;

        for (int j = consumer_tid; j < buffer_len; j += num_consumers) {
            current_out[j] = current_smem[j] * 2.0f; // 业务逻辑:乘以 2
        }

        // 4. 通知生产者:Buffer i%2 我用完了,你可以继续填了
        ready[i % 2].arrive();
    }
}

// ==========================================
// 主 Kernel 函数
// ==========================================
__global__ void producer_consumer_pattern(int N, float *in, float *out, int buffer_len) {
    constexpr int warpSize = 32;

    // 动态共享内存:大小需在启动时指定为 2 * buffer_len * sizeof(float)
    extern __shared__ float buffer[];

    // 屏障存放在共享内存中
    #pragma nv_diag_suppress static_var_with_dynamic_init
    __shared__ barrier_t bar[4];

    // 初始化 4 个屏障
    if (threadIdx.x < 4) {
        init(bar + threadIdx.x, blockDim.x);
    }
    __syncthreads(); // 必须同步,确保初始化完成

    if (threadIdx.x < warpSize) { 
        // 前 32 个线程作为生产者
        produce(bar, bar + 2, buffer, buffer_len, in, N); 
    } else { 
        // 剩下的 96 个线程(假设 block 为 128)作为消费者
        consume(bar, bar + 2, buffer, buffer_len, out, N); 
    }
}

// ==========================================
// Main 函数
// ==========================================
int main() {
    const int N = 1024 * 1024;    // 数据总量
    const int buffer_len = 1024;  // 缓冲区大小(单个)
    const size_t size = N * sizeof(float);

    // 1. 准备数据
    std::vector<float> h_in(N), h_out(N);
    for(int i=0; i<N; ++i) h_in[i] = static_cast<float>(i);

    float *d_in, *d_out;
    cudaMalloc(&d_in, size);
    cudaMalloc(&d_out, size);
    cudaMemcpy(d_in, h_in.data(), size, cudaMemcpyHostToDevice);

    // 2. 配置启动参数
    // 我们启动一个 Block,包含 128 个线程(1 个生产 Warp,3 个消费 Warp)
    int threadsPerBlock = 128;
    int blocksPerGrid = 1; // 示例中只使用一个 Block 演示逻辑
    
    // 动态共享内存大小:2 个缓冲区容量 + 屏障本身占用的空间(这里简化处理,只算数据部分)
    // 注意:实际严谨写法应考虑对齐和 barrier 对象大小,这里直接在 kernel 里静态声明了 bar[4]
    size_t sharedMemSize = 2 * buffer_len * sizeof(float);

    producer_consumer_pattern<<<blocksPerGrid, threadsPerBlock, sharedMemSize>>>(N, d_in, d_out, buffer_len);

    // 3. 检查结果
    cudaMemcpy(h_out.data(), d_out, size, cudaMemcpyDeviceToHost);
    
    bool success = true;
    for(int i=0; i<10; ++i) { // 检查前 10 个数据
        if(h_out[i] != h_in[i] * 2.0f) success = false;
        std::cout << "In: " << h_in[i] << " -> Out: " << h_out[i] << std::endl;
    }

    std::cout << (success ? "Success!" : "Failed!") << std::endl;

    cudaFree(d_in);
    cudaFree(d_out);
    return 0;
}
cpp 复制代码
In: 0 -> Out: 0
In: 1 -> Out: 2
In: 2 -> Out: 4
In: 3 -> Out: 6
In: 4 -> Out: 8
In: 5 -> Out: 10
In: 6 -> Out: 12
In: 7 -> Out: 14
In: 8 -> Out: 16
In: 9 -> Out: 18

详细拆解

虽然所有人都在"签到",但只有部分人在"等"。我们要看谁调用了 wait(或者 arrive_and_wait)。

为了方便理解,我们假设只有两个 Buffer(0 和 1)和两个状态(Ready 和 Filled)。

第一步:开局(消费者先发话)
  • 消费者(Warps 1-3) :一上来就执行 ready[0].arrive()ready[1].arrive()

    • 同步点 :它们不等待,只是在那刷了两下卡,然后立刻冲向 filled[0].arrive_and_wait()由于此时生产者还没搬货,消费者在此处卡住(第一个真正的阻塞点)。
  • 生产者(Warp 0) :一上来执行 ready[0].arrive_and_wait()

    • 同步点 :因为它自己签到了(32个信号),消费者也签到了(96个信号),总数够 128 了!生产者不阻塞,直接通过,开始往 Buffer 0 搬货。
第二步:生产中(各司其职)
  • 生产者 :搬完 Buffer 0 的货,执行 filled[0].arrive()

    • 同步点 :它不等待,刷完卡立刻回头去等 ready[1](准备搬下一份货)。
  • 消费者 :本来在 filled[0].arrive_and_wait() 等着。

    • 同步点 :因为它自己早就签过到了(在 consume 函数的循环里),现在生产者也签到了,总数够 128 了!消费者被唤醒,开始处理 Buffer 0。
第三步:并行流水线(核心高能区)

当时间线推进到这一刻,GPU 内部发生了"分身":

1. 生产者(Warp 0):抢跑 Buffer 1

  • 生产者刚刚在 filled[0] 刷完卡,它根本不看 消费者有没有开始吃 Buffer 0,而是直接进入了 i=1 的循环。

  • 它遇到了 ready[1].arrive_and_wait()

  • 奇迹发生了 :因为消费者在**第一步(初始化)**的时候,已经预先执行了 ready[1].arrive()。现在生产者一打卡,ready[1] 计数立刻归零翻转。

  • 结果 :生产者完全不阻塞,直接冲进 Buffer 1 开始搬运数据。

2. 消费者(Warps 1-3):消化 Buffer 0

  • 与此同时,消费者刚刚从 filled[0].arrive_and_wait() 被唤醒。

  • 它现在正满负荷地在处理 Buffer 0 里的数据(比如做复杂的浮点运算)。

  • 完成后,ready[0].arrive()打卡。生成者搬完Buffer 1后,又可以继续搬Buffer 0.

相关推荐
U盘失踪了2 小时前
Go 结构体
笔记·golang
xuhaoyu_cpp_java5 小时前
连接池学习
数据库·经验分享·笔记·学习
东京老树根6 小时前
SAP学习笔记 - BTP SAP Build05 - SAP BTP BPA简介,Email Destination Settings(TODO)
笔记·学习
talen_hx2966 小时前
《零基础入门Spark》学习笔记 Day 17
大数据·笔记·学习·spark
北山有鸟6 小时前
Linux第一宏:container_of
笔记·嵌入式硬件·学习
做cv的小昊7 小时前
【TJU】研究生应用统计学课程笔记(2)——第一章 数理统计的基本知识(1.3 统计中常用的分布族)
笔记·线性代数·数学建模·矩阵·概率论·学习方法·抽象代数
ouliten7 小时前
cuda编程笔记(38)--CUDA 异步回调
笔记·cuda
Better Bench8 小时前
《八十天环游地球》阅读笔记
笔记·读书笔记·八十天环游地球
sheeta19988 小时前
LeetCode 每日一题笔记 日期:2026.04.21 题目:1722. 执行交换操作后的最小汉明距离
笔记·算法·leetcode