基于前两节 节的内容本节将详细指导并演示 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.0NVIDIA 引入了TMATensor 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。后者才是隐藏全局内存延迟的最优方式。__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”状态。#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::barriercuda::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_t4(sizeof(float)), barrier); } else if (tid 32 - 8) { // 加载右光晕 cuda::memcpy_async(buffer tid 16, right tid, cuda::aligned_size_t4(sizeof(float)), barrier); } if (tid 32) { // 加载中心数据 cuda::memcpy_async(buffer 8 tid, center tid, cuda::aligned_size_t4(sizeof(float)), barrier); } // --- 版本 2跨所有线程协作式地发起批量拷贝 --- // 这种方式更简洁API 内部会处理负载均衡 // cuda::memcpy_async(block, buffer, left, cuda::aligned_size_t4(8 * sizeof(float)), barrier); // cuda::memcpy_async(block, buffer 8, center, cuda::aligned_size_t4(32 * sizeof(float)), barrier); // cuda::memcpy_async(block, buffer 40, right, cuda::aligned_size_t4(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_t4()。这是在告诉编译器数据是4字节对齐的且拷贝大小也是4的倍数。这对于触发底层的LDGSTS指令至关重要。cuda::memcpy_asynctemplate 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);这个函数是libcudacxxCUDA 的 C 标准库实现中定义的memcpy_async的一个重载版本。它最核心的特点是深度集成了cuda::barrier从而实现了对异步生命周期的自动化管理。template class _Tp, typename _Size, thread_scope _Sco, typename _CompF:_Tp: 待拷贝的数据类型如float,int4等。_Size: 拷贝大小的类型可以是普通的size_t但在高性能场景通常是cuda::aligned_size_tN。_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_t16(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_ttemplate _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_t16(32)时编译时提示你告诉编译器“我保证源地址、目标地址和拷贝长度都至少是 16 字节对齐的。”触发优化路径cuda::memcpy_async内部使用了模板重载或if constexpr。当它检测到参数类型是aligned_size_t16时它会直接生成底层硬件支持的最快指令例如LDGSTS.128而不是生成一堆通用的、慢速的逐字节拷贝指令。回到我们之前的讨论LDGSTS支持 4、8 或 16 字节的拷贝如果对齐是4 或 8硬件走 L1 缓存。如果对齐是16硬件可以走L1 Bypass模式减少缓存污染。如果你不使用aligned_size_t编译器为了保证程序的正确性万一地址没对齐呢通常会退化成最保守的、效率最低的拷贝方式。cuda::memcpy_async的集体接口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”。#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::pipelinecuda::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_t4(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_priorpending_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_t4(sizeof(int)), pipeline); } // 即使没有数据可取也要 commit 以保持流水线计数平衡 pipeline.producer_commit(); // 轮转阶段索引 stage (stage 1) % num_stages; } }constexpr size_t pending_batchesnum_stages-1; cuda::pipeline_consumer_wait_priorpending_batches(pipeline);这行代码是cuda::pipeline机制中最核心的“同步刹车”。它的作用是确保流水线中“最老”的那一个数据批次已经搬运完毕可以安全地开始计算。让我们跟踪一个num_stages 3三缓冲的例子初始状态你连续提交了 Batch 0, 1, 2。此时流水线里有3个任务。执行wait_prior2程序会检查流水线里现在有几个任务发现有 3 个。由于我们要等待直到只剩 2 个所以它会阻塞直到 Batch 0 完成。Batch 0 一完流水线里就只剩下 Batch 1 和 2共2个满足了条件程序“放行”。结果你现在可以安全地去处理 Batch 0 的数据了而此时 Batch 1 和 2 依然在硬件后台异步搬运完全没有浪费时间。cuda::pipeline_consumer_wait_priortemplate uint8_t _Prior _LIBCUDACXX_INLINE_VISIBILITY void pipeline_consumer_wait_prior(pipelinethread_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字节对齐修改代码比较麻烦这里只给出对应表// 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_asyncCP_ASYNC_CG将拷贝请求扔进异步代理队列。.cg表示缓存策略。pipeline.producer_commit()CP_ASYNC_COMMIT_GROUP()在队列中划下一道“分割线”标记这一批请求为一个 Group。wait_priorN(pipeline)CP_ASYNC_WAIT_GROUP(N)强制阻塞直到队列中剩下的 Group 数量 $\le N$。pipeline.producer_acquire()(无对应指令)纯软件层面的资源申请。pipeline.consumer_release()(无对应指令)纯软件层面的状态释放。TMA (Tensor Memory Accelerator)很遗憾我的GPU只到了sm89的地步。该技术只支持SM90以上的版本。现在这里占个位。等我有钱了我再考虑学习。