cuda编程笔记（41）--异步数据拷贝

基于前两节节的内容，本节将详细指导并演示 GPU 内存层级内的异步数据移动。内容涵盖：用于逐元素拷贝的 LDGSTS 、用于块状（一维和多维）传输的张量内存加速器 (TMA) ，以及用于寄存器到分布式共享内存拷贝的 STAS ；并展示了这些机制如何与异步屏障 (Asynchronous Barriers) 和流水线 (Pipelines) 集成。

使用 LDGSTS

许多 CUDA 应用程序需要在全局内存和共享内存之间进行频繁的数据移动。通常，这涉及拷贝较小的数据元素或执行不规则的内存访问模式。LDGSTS（计算能力 8.0+）的主要目标是：在进行较小的、逐元素的数据传输时，提供一种从全局内存到共享内存的高效异步传输机制，同时通过重叠执行来提高计算资源的利用率。

维度 (Dimensions)： LDGSTS 支持拷贝 4、8 或 16 字节。拷贝 4 或 8 字节时，始终处于所谓的 L1 ACCESS 模式，此时数据也会缓存在 L1 中；而拷贝 16 字节则可以启用 L1 BYPASS 模式，这种情况下不会污染 L1 缓存。
源与目标 (Source and Destination)： LDGSTS 异步拷贝操作仅支持一种方向：从全局内存到共享内存。根据拷贝数据的大小，指针需要进行 4、8 或 16 字节的对齐。当共享内存和全局内存的对齐均为 128 字节时，可以获得最佳性能。
异步性 (Asynchronicity)： 使用 LDGSTS 的数据传输是异步的，并被建模为"异步线程操作"（参见"异步线程与异步代理"章节）。这允许发起请求的线程在硬件异步拷贝数据时继续执行计算。实际上，数据传输是否能实现异步取决于硬件实现，并可能在未来发生变化。
完成信号： LDGSTS 必须提供一个操作完成的信号。它可以利用共享内存屏障 或流水线 作为提供完成信号的机制。默认情况下，每个线程仅等待其自身发起的 LDGSTS 拷贝。因此，如果您使用 LDGSTS 预取某些将与其他线程共享的数据，在与 LDGSTS 完成机制同步后，仍需要调用 __syncthreads()。

与cp.async的区别

简单来说，LDGSTS 和 cp.async 在功能上其实是指同一种技术，即从全局内存（Global Memory）到共享内存（Shared Memory）的异步拷贝。

它们的主要区别在于抽象层面 和文档称呼习惯的变化：

1. 术语层面的区别

cp.async ：这是 PTX（并行线程执行）指令的名字。当你阅读早期的 Ampere（架构 8.0）技术博客或编写底层 PTX 汇编时，你会看到这个词。它代表 "Copy Asynchronously"。
LDGSTS ：这是 SASS（机器汇编）指令 的名字，全称是 LoaD Global, STore Shared 。在 NVIDIA 较新的官方文档（特别是 Hopper 架构之后）中，为了更精确地描述硬件行为，开始统一使用 LDGSTS 这个术语。

2. 为什么现在强调 LDGSTS？

在 Ampere 架构（SM 8.0）时代，异步拷贝只有这一种主流方式。但到了 Hopper 架构（SM 9.0），NVIDIA 引入了 TMA（Tensor Memory Accelerator，张量内存加速器）。

为了区分不同的异步拷贝机制，文档进行了细化：

LDGSTS (Element-wise Copy) ：指传统的、由线程显式发起的、以 4/8/16 字节为单位的异步拷贝（也就是以前说的 cp.async）。
TMA (Bulk Copy)：指由专门硬件单元负责的、大块数据的、多维度的异步传输，不再依赖单个线程去循环读取。

简单理解就是，cp.async是PTX指令。最终会被编译成SASS指令，用的也是LDGSTS。

在条件代码中批量加载

在这个卷积（Stencil）示例中，线程块的第一个 Warp 负责集体加载所有必要的数据，包括中心数据以及左、右光晕（Halo）数据。

在使用同步拷贝时，由于代码具有条件分支特性，编译器可能会选择生成一系列"全局加载 (LDG) -> 存储到共享内存 (STS)"的指令序列，而不是先执行 3 个 LDG 再执行 3 个 STS。后者才是隐藏全局内存延迟的最优方式。

cpp 复制代码

__global__ void stencil_kernel(const float *left, const float *center, const float *right) {
    // 缓冲区结构：左光晕(8元素) - 中心(32元素) - 右光晕(8元素)
    __shared__ float buffer[8 + 32 + 8];
    const int tid = threadIdx.x;

    // 同步拷贝写法：编译器可能无法很好地对这些访存进行流水线编排
    if (tid < 8) {
        buffer[tid] = left[tid]; // 加载左光晕
    } else if (tid >= 32 - 8) {
        buffer[tid + 16] = right[tid]; // 加载右光晕
    }

    if (tid < 32) {
      buffer[tid + 8] = center[tid]; // 加载中心数据
    }

    __syncthreads();
    // 执行卷积计算
}

为了确保以最优方式加载数据，我们可以将同步内存拷贝替换为异步拷贝。这不仅能通过将数据直接从全局内存拷贝到共享内存来减少寄存器占用，还能确保所有全局内存加载指令同时处于"在途（In-flight）"状态。

cpp 复制代码

#include <cooperative_groups.h>
#include <cuda/barrier>

__global__ void stencil_kernel(const float *left, const float *center, const float *right) {
    auto block = cooperative_groups::this_thread_block();
    auto thread = cooperative_groups::this_thread();
    const int tid = threadIdx.x;

    using barrier_t = cuda::barrier<cuda::thread_scope_block>;
    __shared__ barrier_t barrier;
    __shared__ float buffer[8 + 32 + 8];

    // 初始化异步屏障对象
    if (block.thread_rank() == 0) {
        init(&barrier, block.size());
    }
    __syncthreads();

    // --- 版本 1：在各个线程中单独发起拷贝 ---
    if (tid < 8) {
        // 加载左光晕，通过对齐参数告知编译器使用 LDGSTS
        cuda::memcpy_async(buffer + tid, left + tid, cuda::aligned_size_t<4>(sizeof(float)), barrier);
    } else if (tid >= 32 - 8) {
        // 加载右光晕
        cuda::memcpy_async(buffer + tid + 16, right + tid, cuda::aligned_size_t<4>(sizeof(float)), barrier);
    }

    if (tid < 32) {
        // 加载中心数据
        cuda::memcpy_async(buffer + 8 + tid, center + tid, cuda::aligned_size_t<4>(sizeof(float)), barrier);
    }

    // --- 版本 2：跨所有线程协作式地发起批量拷贝 ---
    // 这种方式更简洁，API 内部会处理负载均衡
    // cuda::memcpy_async(block, buffer, left, cuda::aligned_size_t<4>(8 * sizeof(float)), barrier); 
    // cuda::memcpy_async(block, buffer + 8, center, cuda::aligned_size_t<4>(32 * sizeof(float)), barrier);
    // cuda::memcpy_async(block, buffer + 40, right, cuda::aligned_size_t<4>(8 * sizeof(float)), barrier);

    // 等待所有异步拷贝完成
    barrier.arrive_and_wait();
    __syncthreads();

    // 执行卷积计算
}

异步屏障与 memcpy_async 的协作 ： cuda::memcpy_async 针对 cuda::barrier 的重载版本非常强大。它在创建拷贝任务时会自动增加屏障当前阶段的"预期计数（Expected count）"，并在拷贝完成时自动递减该计数。只有当所有参与线程都到达屏障，且所有绑定的 memcpy_async 操作都完成后，屏障阶段才会推进。
线程级 vs 集体级拷贝 ：你可以选择由各个线程根据 if 条件发起拷贝（版本 1），也可以直接调用接受 block 参数的集体接口（版本 2）。版本 2 中，API 会自动在底层处理如何分配拷贝任务。
性能优化的关键：对齐 ：代码中使用了 cuda::aligned_size_t<4>()。这是在告诉编译器：数据是 4字节 对齐的，且拷贝大小也是 4 的倍数。这对于触发底层的 LDGSTS 指令至关重要。

cuda::memcpy_async

cpp 复制代码

template <class _Tp, typename _Size, thread_scope _Sco, typename _CompF>
_LIBCUDACXX_INLINE_VISIBILITY async_contract_fulfillment
memcpy_async(_Tp* __destination, _Tp const* __source, _Size __size, barrier<_Sco, _CompF>& __barrier);

这个函数是 libcudacxx（CUDA 的 C++ 标准库实现）中定义的 memcpy_async 的一个重载版本。它最核心的特点是深度集成了 cuda::barrier，从而实现了对异步生命周期的自动化管理。

template <class _Tp, typename _Size, thread_scope _Sco, typename _CompF>:
- _Tp: 待拷贝的数据类型（如 float, int4 等）。
- _Size: 拷贝大小的类型，可以是普通的 size_t，但在高性能场景通常是 cuda::aligned_size_t<N>。
- _Sco: 屏障的作用域（Scope），例如 thread_scope_block 表示块级同步。
- _CompF: 屏障到达后的回调函数类型（Completion Function）。
async_contract_fulfillment : 这是一个标记返回类型。它告诉编译器和开发者：该函数并不保证此时数据已经拷贝完成，它只是完成了"合约的履行"------即成功将拷贝请求提交到了硬件队列中。

参数名称	作用
*`_Tp __destination`**	目标地址：通常是指向 Shared Memory 的指针。
*`_Tp const __source`**	源地址：通常是指向 Global Memory 的指针。
`_Size __size`	拷贝大小：单位是字节。如果传入 `cuda::aligned_size_t<16>(32)`，不仅告诉了大小，还告知了硬件地址是对齐的，从而触发 LDGSTS。
`barrier<...> & __barrier`	异步屏障：这是该重载的核心。它负责跟踪这个异步任务的进度。

这个重载版本之所以比传统的异步拷贝更好用，是因为它在底层自动完成了以下几件事：

自动增加预期计数 (Arrive on Creation) ：当你调用这个函数时，它会检测传入的 __barrier。它会自动将屏障当前阶段（Phase）的"预期完成计数 (Expected Count)"加 1。你不需要手动写代码去计算有多少个拷贝在运行。
绑定异步代理 (Binding to Async Proxy)：它启动硬件的异步拷贝引擎（在 Ampere+ 架构上通常是 LDGSTS 指令）。这个过程不占用当前线程的计算资源，线程可以继续往下执行其他指令。
自动释放计数 (Signal on Completion) ：当硬件层面的数据传输真正完成后，硬件会自动向 __barrier 发送一个信号，将计数减 1。
与计算重叠 ：因为有了屏障，你可以先发起一堆 memcpy_async，然后去做一些不依赖这些数据的计算，最后调用 barrier.wait()。

假设你的线程块中有 N 个线程。

初始化阶段：

你调用 init(&barrier, N)。此时，屏障的"预期计数"（Expected Count）等于 N。这代表它在等待 N 个线程发出到达信号。
发起异步拷贝：

当你执行 cuda::memcpy_async(..., barrier) 时，每发起一次拷贝，屏障内部的预期计数就会动态自增 1。
- 如果你发起了 M 个异步拷贝任务，当前的预期计数就变成了 N + M。
调用 arrive_and_wait()：

这个操作分为两步：
- Arrive (到达) ：当前线程宣告"我完成了我的任务"，计数器 -1 。当所有 N 个线程都调用了 arrive，计数器减去了 N。
- Wait (等待) ：此时计数器还剩下 M （即那 M 个还没完成的异步拷贝）。线程会在这里阻塞，直到硬件搬运完最后一字节数据并发出信号，让计数器减到 0。

cuda::aligned_size_t

cpp 复制代码

template <_CUDA_VSTD::size_t _Alignment>
struct aligned_size_t
{
  static constexpr _CUDA_VSTD::size_t align = _Alignment;
  _CUDA_VSTD::size_t value;
  _LIBCUDACXX_INLINE_VISIBILITY explicit constexpr aligned_size_t(size_t __s)
      : value(__s)
  {}
  _LIBCUDACXX_INLINE_VISIBILITY constexpr operator size_t() const
  {
    return value;
  }
};

这个结构体非常简单，主要包含两部分：

静态部分 (static constexpr _Alignment)：在编译时就确定的对齐数值（如 4, 8, 16）。
动态部分 (value)：在运行时实际要拷贝的字节数。

它的设计目的不是为了存储数据，而是为了作为一个 "带有属性的尺寸标签"。

在 C++ 中，如果你只传递一个普通的 size_t，编译器只知道要拷贝多少字节，但它不敢保证这些字节的起始地址和结束地址是否是对齐的。

当你使用 cuda::aligned_size_t<16>(32) 时：

编译时提示：你告诉编译器："我保证源地址、目标地址和拷贝长度都至少是 16 字节对齐的。"
触发优化路径 ：cuda::memcpy_async 内部使用了模板重载或 if constexpr。当它检测到参数类型是 aligned_size_t<16> 时，它会直接生成底层硬件支持的最快指令（例如 LDGSTS.128），而不是生成一堆通用的、慢速的逐字节拷贝指令。

回到我们之前的讨论，LDGSTS 支持 4、8 或 16 字节的拷贝：

如果对齐是 4 或 8，硬件走 L1 缓存。
如果对齐是 16 ，硬件可以走 L1 Bypass 模式，减少缓存污染。

如果你不使用 aligned_size_t，编译器为了保证程序的正确性（万一地址没对齐呢？），通常会退化成最保守的、效率最低的拷贝方式。

cuda::memcpy_async的集体接口

cpp 复制代码

template <typename _Group, class _Tp, _CUDA_VSTD::size_t _Alignment, thread_scope _Sco, typename _CompF>
_LIBCUDACXX_INLINE_VISIBILITY async_contract_fulfillment memcpy_async(
  _Group const& __group,
  _Tp* __destination,
  _Tp const* __source,
  aligned_size_t<_Alignment> __size,
  barrier<_Sco, _CompF>& __barrier);

在第一种方式中，每个线程必须明确计算自己要搬运的地址。而在第二种方式（集体接口）中：

输入一致性 ：线程组中的所有线程都必须调用这个函数，并且传入相同的参数（相同的源地址、目标地址和大小）。
内部负载均衡 ：你不再需要写 if (tid < 8) 这样的逻辑。API 会根据线程组的大小和总数据量，自动在底层分配每个线程应该负责搬运哪一部分数据。

虽然这是一个集体操作，但它对 __barrier 的操作逻辑与单线程版本一致，只是规模不同：

Arrive (创建时)：当这个集体函数被调用时，它依然会自动增加屏障的"预期完成计数"。
动态调整：无论底层 API 决定开启多少个并发的硬件拷贝流，它都会确保在所有数据搬运完成后，屏障的计数会归位。

数据预取 (Prefetching Data)

在本例中，我们将演示如何使用异步数据拷贝将数据从全局内存预取到共享内存。在"拷贝与计算"循环往复的模式中，这种方法可以用当前迭代的计算来掩盖未来迭代的数据传输延迟，从而增加"在途字节数（Bytes-in-flight）"。

cpp 复制代码

#include <cooperative_groups.h>
#include <cuda/pipeline>

template <size_t num_stages = 2 /* 默认 2 级流水线 */>
__global__ void prefetch_kernel(int* global_out, int const* global_in, size_t size, size_t batch_size) {
    auto grid = cooperative_groups::this_grid();
    auto block = cooperative_groups::this_thread_block();
    const int tid = threadIdx.x;

    // 外部共享内存：大小为 (阶段数 * 线程块大小 * sizeof(int)) 字节
    extern __shared__ int shared[]; 
    size_t shared_offset[num_stages];
    for (int s = 0; s < num_stages; ++s) shared_offset[s] = s * block.size();

    // 创建线程作用域的流水线对象
    cuda::pipeline<cuda::thread_scope_thread> pipeline = cuda::make_pipeline();

    // 辅助 lambda：计算当前 batch 在全局内存中的起始偏移
    auto block_batch = [&](size_t batch) -> int {
        //同个block负责的不同batch是grid-stride的
        return block.group_index().x * block.size() + grid.size() * batch;
    };

    // --- 启动阶段：填充流水线前 num_stages 个批次 ---
    for (int s = 0; s < num_stages; ++s) {
        pipeline.producer_acquire(); // 申请一个生产者名额
        // 异步加载数据到对应的共享内存阶段
        cuda::memcpy_async(shared + shared_offset[s] + tid, 
                           global_in + block_batch(s) + tid, 
                           cuda::aligned_size_t<4>(sizeof(int)), 
                           pipeline);
        pipeline.producer_commit();  // 提交该阶段的任务
    }

    int stage = 0;
    // compute_batch: 下一个要处理的批次
    // fetch_batch:   下一个要从全局内存预取的批次
    for (size_t compute_batch = 0, fetch_batch = num_stages; 
         compute_batch < batch_size; 
         ++compute_batch, ++fetch_batch) {

        // 1. 等待流水线中最旧的一个请求完成
        // 这里保留 num_stages - 1 个批次在后台运行
        constexpr size_t pending_batches = num_stages - 1;
        cuda::pipeline_consumer_wait_prior<pending_batches>(pipeline);
        
        // 如果数据要在线程间共享，则需要屏障
        __syncthreads(); 

        // 2. 在当前批次上执行计算
        compute(global_out + block_batch(compute_batch) + tid, 
                shared + shared_offset[stage] + tid);
                
        // 3. 释放当前阶段，告诉生产者该 Buffer 现在可以重新使用了
        pipeline.consumer_release();
        __syncthreads(); 

        // 4. 生产者：预取未来的第 fetch_batch 个批次
        pipeline.producer_acquire();
        if (fetch_batch < batch_size) {
            cuda::memcpy_async(shared + shared_offset[stage] + tid, 
                               global_in + block_batch(fetch_batch) + tid, 
                               cuda::aligned_size_t<4>(sizeof(int)), 
                               pipeline);
        }
        // 即使没有数据可取，也要 commit 以保持流水线计数平衡
        pipeline.producer_commit();
        
        // 轮转阶段索引
        stage = (stage + 1) % num_stages;
    }
}

cpp 复制代码

constexpr size_t pending_batches=num_stages-1;
cuda::pipeline_consumer_wait_prior<pending_batches>(pipeline);

这行代码是 cuda::pipeline 机制中最核心的"同步刹车"。它的作用是：确保流水线中"最老"的那一个数据批次已经搬运完毕，可以安全地开始计算。

让我们跟踪一个 num_stages = 3（三缓冲）的例子：

初始状态 ：你连续提交了 Batch 0, 1, 2。此时流水线里有 3 个任务。
执行 wait_prior<2>：
- 程序会检查：流水线里现在有几个任务？
- 发现有 3 个。
- 由于我们要等待直到只剩 2 个，所以它会阻塞，直到 Batch 0 完成。
- Batch 0 一完，流水线里就只剩下 Batch 1 和 2（共 2 个），满足了条件，程序"放行"。
结果：你现在可以安全地去处理 Batch 0 的数据了，而此时 Batch 1 和 2 依然在硬件后台异步搬运，完全没有浪费时间。

cuda::pipeline_consumer_wait_prior

cpp 复制代码

template <uint8_t _Prior>
_LIBCUDACXX_INLINE_VISIBILITY void
pipeline_consumer_wait_prior(pipeline<thread_scope_thread>& __pipeline);

pipeline_consumer_wait_prior<_Prior> 是 cuda::pipeline 的"消费者"接口。

它的语义 ：强制当前线程阻塞，直到流水线中处于"在途（In-flight）"状态的任务批次数量小于或等于 _Prior。
硬编码优化 ：由于 _Prior 是作为一个 Template Parameter（模板参数） 传入的，编译器在编译阶段就知道了这个数值。这允许它生成极其精简的汇编指令（如 SASS 层的 DEPBAR 指令），而不是在运行时再去解析变量，这对于性能极其敏感的内层循环（Inner Loop）至关重要。

特别注意，这个接口只能给 thread_scope_thread 用。

和cp.async的对应

上述代码使用cp.async指令也可以实现。不过由于cp.async.cg.shared.global只支持16字节对齐，修改代码比较麻烦，这里只给出对应表

cpp 复制代码

// PTX 宏定义
#define CP_ASYNC_CG(dst, src, Bytes) \
    asm volatile("cp.async.cg.shared.global [%0], [%1], %2;\n" :: "r"(dst), "l"(src), "n"(Bytes))

#define CP_ASYNC_COMMIT_GROUP() \
    asm volatile("cp.async.commit_group;\n" ::)

#define CP_ASYNC_WAIT_GROUP(N) \
    asm volatile("cp.async.wait_group %0;\n" :: "n"(N))

高级 API (cuda::pipeline)	PTX 指令 / 宏	硬件行为
`cuda::memcpy_async`	`CP_ASYNC_CG`	将拷贝请求扔进异步代理队列。`.cg` 表示缓存策略。
`pipeline.producer_commit()`	`CP_ASYNC_COMMIT_GROUP()`	在队列中划下一道"分割线"，标记这一批请求为一个 Group。
`wait_prior<N>(pipeline)`	`CP_ASYNC_WAIT_GROUP(N)`	强制阻塞，直到队列中剩下的 Group 数量 $\\le N$ 。
`pipeline.producer_acquire()`	(无对应指令)	纯软件层面的资源申请。
`pipeline.consumer_release()`	(无对应指令)	纯软件层面的状态释放。

TMA (Tensor Memory Accelerator)

很遗憾，我的GPU只到了sm89的地步。该技术只支持SM90以上的版本。现在这里占个位。等我有钱了，我再考虑学习。