cuda pipeline 源码解析 - 学习让我快乐

之前在学习 cutlass 实现时, 看到其使用了 cuda pipeline 设施, 然后看了下 pipeline 的文档似懂非懂 :-(

// 如下这段代码会被 thread block 中所有线程协作执行.
__shared__ cuda::pipeline_shared_state<thread_scope_block, stages_count> shared_state;
// 注意这里 pipeline 每个 thread 都有自己的副本: p1, p2, p3, p4
auto pipeline = cuda::make_pipeline(group, &shared_state);

pipeline.producer_acquire();  // #1
cuda::memcpy_async(group, pipeline);
pipeline.producer_commit();
for (batch_idx : batch_num) {
  pipeline.producer_acquire();
  cuda::memcpy_async(group, pipeline);
  pipeline.producer_commit();

  pipeline.consumer_wait();
  compute();
  pipeline.consumer_release();
}

在传统 cpu 编程中, 如上代码很清晰, 每个线程执行 producer_acquire 拿到一个类似 token 的东西, 该函数返回则意味着线程成功占用了 pipeline stage, 此时可以做一些事情. 但是在 cuda 执行模型中, 以一个 warp 为例, 一下子有 32 个线程执行 producer_acquire(), 那么这时的行为是啥? 只有一个线程拿到 token, 然后 32 个线程一个个去等? 这会产生严重的 warp diverge, 绝对不可能的! 文档看得稀里糊涂索性翻了下代码实现. 如下是 make_pipeline(group, shared_state) 之后 pipeline 状态:

# 位于 SM shared memory
pipeline_shared_state
  __stages[_Stages_count]: __pipeline_stage<_Scope>
    __produced barrier<_Scope>, init(__group_size)
    __consumed barrier<_Scope>, init(__group_size)

pipeline:
  __consumed_phase_parity: bool, 初始值 true
  __produced_phase_parity: bool, 初始值 false

P.S. 见鬼, 为啥老喜欢加双下划线啊!

producer_acquire() 的实现:

auto& b = __shared_state_get_stage(__head)->__consumed;
b.wait_parity(pipeline.__consumed_phase_parity); barrier 初始 parity 是 false, __consumed_phase_parity 初始是 true, 这里 wait 会立即就绪. 意味着所有的 consumer 都已经 release 这个 stage 了, producer 可以放心 acquire 了.
- Q: 设想情况 p1, p2, p3 wait_parity 都立刻结束了, p4 被调度卡了一会, 会不会出现一种情况, 在 p4 执行 wait_parity 时, b.phase 已经从 false 变为 true?
  
  A: 不会, b.phase 要从 false 变为 true, 即 consumer_release 中 __consumed.arrive(), 需要全部 4 个线程都参与了.

cuda::memcpy_async(group, pipeline);

__cp_async_shared_global<__copy_size>(__dest + __offset, __src + __offset), group 内所有线程协作执行完成 src 到 dst 拷贝. 该函数即 cp.async.ca.shared.global 指令封装.
__memcpy_completion_impl::__defer(__cm=__async_group, group, __size, pipeline); 就是简单的返回 async_contract_fulfillment::async.

pipeline<thread_scope_block>.producer_commit() 实现:

auto& __stage_barrier = __shared_state_get_stage(__head)->__produced;
cp.async.mbarrier.arrive.shared.b64 __stage_barrier. 注意这里没有使用 noinc, __stage_barrier pending arrive count 是 group size.
__stage_barrier.arrive();

这样看下来整体流程就比较清晰了. 回过头来再看 cuda pipeline, 其使用其实很简单, pipeline 连接着 producer, consumer, 其中 producer 是一堆线程, consumer 也是一堆线程. producer 执行 producer_acquire, 此时会等待 consumer 中一堆线程都执行了 consumer_release 之后返回, 此时这堆 producer 线程可以协作进行一些事情. 同理 consumer 中一坨线程执行 consumer_wait 此时会等待 producer 那堆线程都执行了 producer_commit 之后返回, 此时这组 consumer 线程可以协作做一些事情.

CATALOG