CUDA 异步数据搬运学习笔记:从 cp.async 到 TMA

目录：Overlap 挑战 - 同步基线 - cp.async - PTX 语法 - Async-group - TMA 与 Tensor Map - 传输模式 - 完成机制 - 高级特性

在 GPU kernel 里，两个地方存在很大的挑战：数据离计算单元太远，以及搬运与计算的时序被绑死。

global memory（DRAM）的带宽很高，但访问延迟也高。如果每次计算都临时去 global 取数，计算单元会花大量时间等数据。传统写法把「数据拷进 shared」与「使用 shared 做计算」写成顺序代码，再用 __syncthreads() 形成严格的阶段边界。这样做是安全的，但会限制重叠空间。

异步数据搬运解决的核心问题是 Overlap。CUDA 允许多个任务重叠执行：host 计算、device 计算、host 与 device 之间的传输、device 内部的搬运，这些都可以并发。但可以并发不等于自动并发，能否重叠取决于代码写法、同步机制以及硬件资源。

异步搬运把发起拷贝与等待完成分离：先发起搬运，让数据在后台移动；在搬运尚未完成时做其他工作；在真正使用数据前，再用显式的 completion 机制等待搬运完成。这是单 kernel 内部的异步，与 cudaMemcpyAsync 那种应用层流式异步属于不同层次。

现代 CUDA/PTX 提供多种机制支持 global 到 shared 的异步搬运：

• 同步时代：host 侧 cudaMemcpy/cudaMemcpyAsync 与 kernel 内 ld + st + block 同步
• Ampere (sm_80+)：cp.async (PTX) / memcpy_async (CUDA C++) 分离发起与等待
• Hopper (sm_90+)：TMA (Tensor Memory Accelerator) 支持更大块、更高维度的 tile 搬运

术语表：

blocking vs non-blocking

异步搬运的关键在于区分 blocking（阻塞）与 non-blocking（非阻塞）。

两者的核心区别在于线程是否在搬运期间被阻塞：同步模式下，线程发起拷贝后必须原地等待完成，期间无法做其他事；异步模式下，线程发起后立即返回，可以继续执行其他工作（如计算上一轮数据），只在真正需要数据时才显式等待。这种"搬运与计算重叠"正是异步机制的价值所在。

• 同步模式：线程在发起拷贝后必须等待数据搬运完成，期间无法做其他事情
• 异步模式：线程发起拷贝后立即返回，在数据搬运的同时可以执行其他工作，只在真正需要数据时才显式等待

Completion（完成）机制是异步操作的核心安全机制，确保数据已准备就绪。

异步拷贝的生命周期分为三个阶段：发起（cp.async 或 memcpy_async）→ 传输中（数据在后台移动，destination 内容未定义）→ 确认完成（wait_group 或 barrier.arrive_and_wait()）。只有在确认完成后，destination 数据才保证可用；在此之前读取属于未定义行为（UB）。

关键规则：

• 从发起拷贝到确认完成之间，destination 的内容是未定义的，读取属于数据竞争/UB
• 在确认完成之前，source 也不能被修改（否则可能拷贝到错误的数据）
• 只有通过 completion 机制（如 wait_group、barrier.arrive_and_wait()）确认完成后，数据才保证可用
• 跨线程消费：即使本线程等到了也不代表其他线程等到了，跨线程消费还需要 __syncthreads() 等同步（后文 Async-group 章节详述）

把异步（async）理解为无需等待是一个常见的误读。正确的理解是：异步操作把等待从发起点推迟到了使用点，但等待本身是不可省略的。

同步时代的基线

Host 侧：cudaMemcpy 与 cudaMemcpyAsync

从 API 行为看，cudaMemcpy 是同步拷贝：调用不会返回，直到拷贝完成。cudaMemcpyAsync 是异步拷贝：调用会尽快返回，但我们必须用 stream 同步、event、或显式同步 API 在真正使用拷贝结果之前建立同步点。

cudaStream_t stream;cudaStreamCreate(&stream);cudaMemcpyAsync(dst_device, src_host, bytes, cudaMemcpyHostToDevice, stream);// ... enqueue other work in the same stream ...cudaStreamSynchronize(stream);  // synchronization point

这类异步的本质是：我们把等待从拷贝调用处移动到了后面的某个同步点。这为 overlap（例如与其他 kernel、其他拷贝、或 host 侧工作）提供了空间，但不自动保证总是重叠或总是更快。

一个工程上高频遇到的条件：当拷贝涉及 host memory 并期望异步行为时，host buffer 需要是 pinned（page-locked）。如果 host buffer 不是 pinned，会有两个影响：①host 侧 API 调用可能会阻塞（因为需要先拷贝到临时 pinned buffer）；②与其他 host 工作的 overlap 会受限。具体的 DMA 传输与 device 侧工作能否并行，还取决于 stream 配置与硬件 copy engine。建议用 CUDA event 和时间线工具验证实际行为。

Pinned memory（page-locked memory）是一种不会被操作系统换出（page out）到磁盘的内存。普通的 malloc 分配的内存可能被操作系统换页，这会影响 DMA 传输的性能：GPU 的 DMA 引擎需要直接访问物理内存地址，如果使用 pageable host memory，CUDA runtime 会先分配临时 pinned buffer 并做一次额外拷贝（staging），这会导致 host 侧阻塞，减少可实现的 overlap。

使用 cudaMallocHost 或 cudaHostAlloc 分配 pinned memory：

float* h_pinned;cudaMallocHost(&h_pinned, size);  // 分配 pinned memory// ... 使用 h_pinned ...cudaFreeHost(h_pinned);           // 释放 pinned memory

为什么 pinned memory 能实现真正的异步？ 使用普通内存时，cudaMemcpyAsync 需要先分配临时 pinned buffer 并拷贝数据过去，这两步会阻塞 CPU；使用 pinned memory 时，DMA 可以直接发起，CPU 立即返回继续其他工作。

Kernel 内：经典 global→shared 的寄存器中转

在 kernel 内部，最常见的模式之一是先搬运到 shared，再在 shared 上计算。典型语句：

shared[local_idx] = global[global_idx];

这行代码看起来像直接赋值，但实际会被展开为两步：

• 从 global load 到寄存器
• 再从寄存器 store 到 shared

PTX 是一个 load-store machine：ALU 指令的操作数必须在寄存器里，ld 和 st 指令负责在寄存器与可寻址 state space 之间搬运数据。这意味着传统的 global→shared 搬运必须经过寄存器中转。

在迭代算法里，这种写法通常会伴随两次 block 级同步：

// iteration patternshared[t] = global[g];       // load global -> reg -> shared__syncthreads();             // ensure all threads finished loadingcompute_using(shared);       // use shared data__syncthreads();             // ensure no thread overwrites shared prematurely

第一次同步确保 shared 写入完成后再开始 compute；第二次同步避免某些线程提前覆写 shared，而其他线程仍在使用该数据。

这段代码在正确性上很直观，但它把执行划分为严格阶段：只有当搬运阶段全部结束时计算阶段才能开始。

LDG 与 STS 指令

这两个是 PTX/SASS 层面的指令助记符：

• LDG = LoaDGlobal：从 global memory 加载数据到寄存器
• STS = STore Shared：从寄存器存储数据到 shared memory

传统的 shared[i] = global[i] 会被编译为两条独立指令：

LDG.E R0, [global_addr]    ; 从 global 加载到寄存器 R0（高延迟，数百 cycles 量级）STS   [shared_addr], R0    ; 从 R0 存储到 shared（低延迟）

LDG/STS 交错 vs 批量发射

编译器有两种调度策略：

策略 1: 交错执行（latency 未完全隐藏）LDG R0, [g0] → 等待 → STS [s0], R0LDG R1, [g1] → 等待 → STS [s1], R1LDG R2, [g2] → 等待 → STS [s2], R2              ↑ 每次都要等待 global latency策略 2: 批量发射（更好的 latency hiding）LDG R0, [g0]  ─┐LDG R1, [g1]   ├── 一次发射多个 load（它们可以并行执行）LDG R2, [g2]  ─┘    ... 等待所有 load 完成 ...STS [s0], R0  ─┐STS [s1], R1   ├── 批量 storeSTS [s2], R2  ─┘

策略 2 更好，因为多个 global load 可以并行 in-flight，共同隐藏延迟。

CUDA 文档指出，在条件代码下使用同步拷贝时，编译器可能生成 LDG/STS 交错序列；而先发多次 LDG 再发多次 STS被描述为更优的隐藏 global latency 的方式。但无论编译器如何优化，传统路径的根本限制在于：我们没有显式的发起与等待分离，重叠空间受限于编译器的调度能力。

同步搬运的局限

同步搬运不是错的，但在特定场景下会成为瓶颈：

• 寄存器压力：global→shared 需要寄存器中转，如果同时发起大量搬运，寄存器占用会限制 occupancy
• 阶段边界：__syncthreads() 形成硬性边界，所有线程必须等到同一点才能继续
• 重叠受限：如果下一轮数据具有可预取性，我们无法在当前计算的同时让下一轮数据在后台移动

cp.async：Ampere 的突破

Ampere 架构（sm_80, compute capability 8.0）引入了 kernel 内的异步数据拷贝能力。CUDA 文档将其称为 LDGSTS，用于 global→shared 的 element-wise 异步传输。PTX 对应指令为 cp.async。

non-blocking 的核心语义

PTX ISA 对 cp.async 的定义：

• non-blocking：指令发起后，控制会在拷贝完成之前返回到执行线程
• 地址空间：src 必须位于 global state space，dst 必须位于 shared state space
• 必须显式等待：执行线程随后必须使用 async-group completion 或 mbarrier-based completion 来等待完成；PTX 明确指出没有其他同步机制可以保证异步拷贝完成

第三点最重要。non-blocking 不等于结果立刻可用，我们仍然必须显式等待完成。

LDGSTS：CUDA 文档中的机制描述

LDGSTS 机制

LDGSTS (LoaD Global STore Shared) 是一条复合指令，而非两条独立指令。

回顾前文，传统的 global→shared 搬运需要两步：

• LDG：global → 寄存器
• STS：寄存器 → shared

传统路径需要两步：LDG（Global→寄存器）+ STS（寄存器→Shared）。LDGSTS 把这两步合并成一条指令，跳过寄存器中转，直接从 Global 写入 Shared。

LDGSTS 的优势：

• 减少寄存器压力：不需要临时寄存器来存放中间数据
• 原生异步：硬件直接支持 non-blocking 语义，数据可以在后台移动
• 更高带宽利用率：减少指令数量，让硬件有更多资源处理其他工作

CUDA Programming Guide 将 LDGSTS 描述为 compute capability 8.0+ 的异步 global→shared 数据传输机制。其设计目标是高效的 element-wise 拷贝并支持 overlapped execution。

LDGSTS 的关键特性：

L1 ACCESS vs L1 BYPASS 模式

这两种模式决定了数据是否经过 L1 cache：

选择建议：

• 如果相同的 global 数据后续还会被其他线程/block 读取 → ACCESS（利用 L1 cache）
• 如果是 streaming 访问，数据只用一次 → BYPASS（保护 L1 给其他数据）

LDGSTS 的 16B 变体可以选择 BYPASS 模式（对应 PTX 的 .cg cache operator），而 4/8B 只支持 ACCESS 模式。

LDGSTS 需要完成信号：可用 shared memory barriers 或 pipelines。默认情况下每线程只等待自身拷贝，如果预取的数据要被其他线程共享，则在 completion 机制之后还需要 __syncthreads()。

CUDA 侧的三种入口

CUDA 提供了多种方式来使用这个能力，从高层到低层分别是：

入口 1：cuda::memcpy_async + cuda::barrier

#include<cooperative_groups.h>#include<cuda/barrier>__global__ voidkernel(constfloat* __restrict__ src){auto block = cooperative_groups::this_thread_block();usingbarrier_t = cuda::barrier<cuda::thread_scope_block>;    __shared__ barrier_t barrier;__shared__ alignas(16)float buffer[256];  // 16B 对齐以匹配 aligned_size_t<16>if (block.thread_rank() == 0) {        init(&barrier, block.size());    }    __syncthreads();// Initiate async copy, bind completion to barrier    cuda::memcpy_async(block, buffer, src,                        cuda::aligned_size_t<16>(256 * sizeof(float)), barrier);    barrier.arrive_and_wait();  // 等待异步拷贝完成    __syncthreads();             // 跨线程同步：确保 block 内所有线程都能安全访问 buffer（见后文 Async-group 章节）// Now buffer is ready}

init(&barrier, block.size()) 中的第二个参数 block.size() 是 expected count（期望到达数），即 barrier 期望有多少个线程调用 arrive。当这么多线程都 arrive 后，barrier 的当前 phase 才算完成。

假设 block 有 256 个线程：init(&barrier, 256);  // 期望 256 个线程到达// 每个线程调用 arrive_and_wait：// - arrive: 告诉 barrier 我到了// - wait: 等待所有 256 个线程都到达barrier.arrive_and_wait();

Thread Scope（线程作用域）

cuda::thread_scope_block 指定 barrier 的 可见性范围。CUDA 定义了三种 scope：

对于 shared memory 中的 barrier，通常使用 thread_scope_block，因为 shared memory 本身就是 block 内可见的。

cuda::memcpy_async(block, buffer, src, size, barrier) 的最后一个参数将异步拷贝的 completion 绑定到 barrier。

Barrier Phase 完成条件

当使用 cuda::memcpy_async(..., barrier) 时，barrier 的 phase 完成需要两个条件同时满足：

• 线程 arrive 达到 expected count
• 绑定到该 phase 的所有异步拷贝完成（通过 transaction bytes 机制追踪，详见后文完成机制小节）

注意：expected count 应该设置为参与同步的线程数，不需要为拷贝操作加额外的 count。

cuda::aligned_size_t<16>(...) 不是装饰，而是向编译器传达对齐与 size 是 16B 的倍数。这个信息帮助生成高效代码，使底层更可能使用异步机制。

入口 2：cooperative_groups::memcpy_async + cg::wait

#include<cooperative_groups.h>#include<cooperative_groups/memcpy_async.h>namespace cg = cooperative_groups;__global__ voidkernel(constfloat* __restrict__ src){    cg::thread_block block = cg::this_thread_block();    __shared__ float buffer[256];    cg::memcpy_async(block, buffer, src, 256 * sizeof(float));    cg::wait(block);    __syncthreads();// Now buffer is ready}

cg::wait(block) 等待此前的异步拷贝完成。在 copy 发起到 completion 之间，如果我们修改 global source、或读取/写入 shared destination，都可能引入 data race。

入口 3：Primitives API

#include<cuda_pipeline.h>__global__ voidkernel(constfloat* __restrict__ src){    __shared__ float buffer[256];int tid = threadIdx.x;// Each thread copies its portion    __pipeline_memcpy_async(&buffer[tid], &src[tid], sizeof(float));    __pipeline_commit();    __pipeline_wait_prior(0);    __syncthreads();// Now buffer is ready}

三个 primitive 的角色划分：

• __pipeline_memcpy_async：提交异步拷贝请求
• __pipeline_commit：把此前提交的拷贝归入一个提交点
• __pipeline_wait_prior(N)：等待完成，语义是等待直到 pipeline 中除最近 N 次提交外的更早提交都已完成

满足 only-asynchronous 的条件

术语定义：only-asynchronous

"only-asynchronous" 指拷贝操作完全异步执行，即硬件真正使用异步数据路径（如 LDGSTS），发起线程不被阻塞。 不满足条件时，API 可能 fallback 到同步行为（实际变成 blocking copy，线程等待拷贝完成后才返回）。

cooperative_groups::memcpy_async 的文档明确指出：只有当 src 是 global 且 dst 是 shared 且满足 16/8/4B 对齐时，拷贝才是 only-asynchronous。不满足这些条件时，可能会 fallback 到同步行为。

这意味着我们不能假设调用了 memcpy_async 就一定是异步的。对齐和方向都是硬约束。最佳性能建议是 16B 对齐。

选择哪种入口

无论选择哪一个入口，底线一样：异步搬运的正确性依赖 completion。不等待就访问 destination，或者在完成前修改 source，都可能触发 data race 或未定义行为。

完整示例：memcpy_async 的实际使用流程

基础版：加载→等待→计算

#include<cooperative_groups.h>#include<cooperative_groups/memcpy_async.h>namespace cg = cooperative_groups;// 基础示例：展示正确的 cp.async 使用流程// 注意：对于 partial tile（尾块），未拷贝的元素是 shared 中的旧值__global__ voidbasic_async_copy(constfloat* __restrict__ src,float* __restrict__ dst, int n){    cg::thread_block block = cg::this_thread_block();constexprint TILE_SIZE = 256;    __shared__ float smem[TILE_SIZE];int tile_offset = blockIdx.x * TILE_SIZE;int tid = threadIdx.x;// 先初始化 shared memory（处理 partial tile 的边界情况）    smem[tid] = 0.0f;    __syncthreads();// 计算实际要拷贝的字节数（可能小于 TILE_SIZE * sizeof(float)）int elements_to_copy = min(TILE_SIZE, n - tile_offset);if (elements_to_copy > 0) {// 阶段 1：发起异步拷贝（立即返回）size_t bytes_to_copy = elements_to_copy * sizeof(float);        cg::memcpy_async(block, smem, src + tile_offset, bytes_to_copy);    }// 阶段 2：等待完成    cg::wait(block);    __syncthreads();// 阶段 3：使用数据（smem 中未拷贝的部分是初始化的 0）if (tile_offset + tid < n) {        dst[tile_offset + tid] = smem[tid] * 2.0f;    }}

上面的基础版虽然语法正确，但没有实际的 overlap：发起拷贝后立刻等待，与同步写法效果相似。

进阶版：Double Buffering：展示异步的真正价值

前置条件：本示例假设 n_elements 是 TILE_SIZE 的整数倍，且 grid 配置为 (n_elements / TILE_SIZE / TILES_PER_BLOCK) 个 block。如需处理非对齐情况，需额外添加尾块逻辑。

#include<cuda_pipeline.h>// 进阶示例：使用 pipeline primitives 实现真正的 double buffering// 关键：使用 __pipeline_wait_prior(1) 允许 1 个 stage 仍在传输中__global__ voiddouble_buffer_async(constfloat* __restrict__ input,float* __restrict__ output,int n_elements){constexprint TILE_SIZE = 256;constexprint TILES_PER_BLOCK = 4;  // 每个 block 处理 4 个 tile// 双缓冲：两个 buffer 交替使用    __shared__ float buffer[2][TILE_SIZE];int tid = threadIdx.x;int block_start = blockIdx.x * TILES_PER_BLOCK * TILE_SIZE;// ========== Prologue: 预取第一个 tile 到 buffer[0] ==========int offset0 = block_start;if (offset0 + tid < n_elements) {        __pipeline_memcpy_async(&buffer[0][tid], &input[offset0 + tid], sizeof(float));    }    __pipeline_commit();  // commit group 0// ========== Main Loop: 计算与预取重叠 ==========for (int t = 0; t < TILES_PER_BLOCK; t++) {int current_buf = t % 2;int next_buf = (t + 1) % 2;int current_offset = block_start + t * TILE_SIZE;int next_offset = current_offset + TILE_SIZE;// Step 1: 发起下一个 tile 的拷贝（如果有）if (t + 1 < TILES_PER_BLOCK && next_offset + tid < n_elements) {            __pipeline_memcpy_async(&buffer[next_buf][tid], &input[next_offset + tid], sizeof(float));        }        __pipeline_commit();  // commit group (t+1)// Step 2: 等待当前 tile 就绪// __pipeline_wait_prior(1) 表示允许最近 1 个 group 仍 pending// 所以 group t 会完成，而 group t+1 可以继续传输        __pipeline_wait_prior(1);        __syncthreads();// Step 3: 处理当前 tile（同时 next tile 在后台传输！）if (current_offset + tid < n_elements) {float val = buffer[current_buf][tid];            val = val * 2.0f + 1.0f;            output[current_offset + tid] = val;        }        __syncthreads();  // 确保当前 buffer 用完再被覆写    }}

时间线对比：同步 vs 异步双缓冲

同步写法下，拷贝和计算串行执行：拷贝 tile0 → 计算 tile0 → 拷贝 tile1 → 计算 tile1 → ...，总时间是所有操作时间之和。异步双缓冲下，拷贝 tile1 与计算 tile0 并行，拷贝 tile2 与计算 tile1 并行，以此类推。总时间近似为 max(总拷贝, 总计算)，而非两者之和。

关键洞察：

• 异步的价值在于拷贝与计算的时间重叠
• 如果拷贝后立即等待，与同步写法无本质区别
• Double buffering 是实现 overlap 的经典模式

常见错误写法：

//  错误：没有等待就读取cg::memcpy_async(block, s_a, a + offset, size);float val = s_a[tid];  // UB! 数据可能还没到//  错误：在完成前修改 sourcecg::memcpy_async(block, s_a, a + offset, size);a[offset + tid] = 999.0f;  // UB! source 正在被读取cg::wait(block);//  正确：先等待，再读取cg::memcpy_async(block, s_a, a + offset, size);cg::wait(block);__syncthreads();float val = s_a[tid];  // OK!

cp.async PTX 语法详解

精确控制异步拷贝的行为，或者理解 CUDA 编译器生成的代码，需要了解 PTX 层面的 cp.async 指令。

PTX 背景知识

State Space（状态空间）

PTX 把 GPU 的内存划分为多个逻辑上独立的 state space，每个 space 有不同的特性：

cp.async 的方向限制就是基于 state space：src 必须是 .global，dst 必须是 .shared。

Generic Addressing（通用地址）

有时候指针的目标 state space 在编译时不确定。PTX 提供了 generic address space，它是一个统一的地址空间，可以在运行时解析到具体的 state space。

编译时 state space 已知的指针直接映射到 .global/.shared/.local；而 generic address 是统一地址空间，运行时才解析到具体 state space。

当我们使用 generic address 时，需要注意：

• 某些操作（如 prefetch）只对特定 state space 有效
• 如果 generic address 不落在期望的 state space 窗口内，行为可能是 undefined

在 CUDA C++ 中，普通的指针（如 float*）默认就是 generic 的。

Weak Memory Operation（弱内存操作）

PTX 文档把 cp.async 归类为 weak memory operation。这意味着：

• 不提供强制的 ordering guarantee：即使我们在代码中写了 A 在 B 之前，硬件实际执行时 B 可能先于 A 完成
• 类似 relaxed memory order：如果我们熟悉 C++ memory model，weak operation 类似于 memory_order_relaxed
• 必须显式同步：只有通过 completion 机制（如 wait_group、mbarrier）才能建立可靠的 happens-before 关系

这与普通的 load/store 形成对比：

这就是后文 ordering 小节所述「cp.async 之间没有默认顺序」是关键约束的原因。

语法骨架

PTX ISA cp.async 语法骨架（保留关键字段）：

cp.async.ca.shared{::cta}.global{.level::cache_hint}{.level::prefetch_size}    [dst], [src], cp-size{, src-size}{, cache-policy};cp.async.cg.shared{::cta}.global{.level::cache_hint}{.level::prefetch_size}    [dst], [src], 16{, src-size}{, cache-policy};

还有一种带 ignore-src predicate 的变体：

cp.async.ca.shared{::cta}.global [dst], [src], cp-size{, ignore-src};cp.async.cg.shared{::cta}.global [dst], [src], 16{, ignore-src};

需区分语义硬约束（写错可能 UB）与性能 hint（不改变正确性）。

cp-size 与 src-size：边界规则

cp-size 是拷贝粒度，只能是 {4, 8, 16} 的整数常量，没有其他选项。

src-size 是可选的 32-bit integer operand，用于处理实际需要的数据不足 cp-size 的情况。规则如下：

需注意：src-size 大于 cp-size 导致 UB，不是截断。

一个典型场景是处理数组尾部不足一个完整 cp-size 的情况：

// assume we need to copy only 12 bytes at the tail// use cp-size=16, src-size=12cp.async.ca.shared.global [dst], [src], 16, 12;// dst[0..11] = src[0..11], dst[12..15] = 0

ignore-src：整段置零

ignore-src 是一个 predicate operand（不是 immediate），当它为 true 时：

• 忽略 src 数据
• 向 dst 拷贝 0（整个 cp-size 字节都是 0）
• 默认值是 false

处理越界时可以让越界的 lane 依然发起 cp.async，但用 predicate 把它变成对 destination 的置零，从而避免分支发散与复杂控制流，同时保持 shared tile 内容在语义上可定义：

// tid >= valid_count means out of boundssetp.ge.s32 %p, %tid, %valid_count;cp.async.ca.shared.global [dst], [src], 16, %p;  // if %p=true, zero-fill dst

注意 ignore-src 和 src-size 是不同的边界处理机制：

• src-size：部分拷贝 + 尾部 zero-fill
• ignore-src：整段置零

在 CUDA C++ 中使用 ignore-src 功能

CUDA C++ 高层 API（如 cg::memcpy_async）不直接暴露 ignore-src predicate。如果我们需要这个功能，有两种方式：

• 条件拷贝（高层方式）：

if (tid < valid_count) {    cg::memcpy_async(block, &smem[tid], &global[tid], sizeof(float));} else {    smem[tid] = 0.0f;  // 手动置零越界部分}

• 使用 ignore-src predicate（需要 inline PTX）：

ignore-src 需要通过 inline PTX 使用，但正确编写 inline PTX 涉及寄存器约束、地址空间转换、predicate 生成等复杂细节。除非我们熟悉 PTX 且有性能需求，否则建议使用方式 1。

// 概念性语法（非直接可编译）：// 需要：setp 生成 predicate、正确的地址约束、clobber 声明cp.async.ca.shared.global [smem], [global], 16, ignore_predicate;

方式 1 更简单且可维护，方式 2 在避免分支时可能有性能收益，但实现复杂度高。

Cache operator：.ca vs .cg

.ca 和 .cg 是 cache operator，用于影响缓存层级行为：

这两个是 performance hints only，不会改变 memory consistency 行为。我们不能依赖 .cg 来实现某种同步语义，它只是告诉硬件这个数据可能不需要在 L1。

使用 .cg 时，cp-size 必须是 16。

.ca vs .cg 选择建议

例如，在 tiled matrix multiplication 中：

• 输入矩阵的 tile 会被 block 内的线程多次访问 → .ca
• 如果是 reduction 操作，每个元素只读一次 → .cg

Prefetch hint 与 cache policy

.level::prefetch_size 允许指定预取大小：.L2::64B、.L2::128B、.L2::256B。这也是 performance hint，用于提示硬件可能会继续访问后续数据。

约束：

• 只能用于 .global（或 generic addressing 指向 global）
• 如果 generic address 不落在 global 的地址窗口内，prefetch 行为是 undefined

.level::cache_hint + cache-policy 用于指定 cache eviction policy，同样是 hint。

对齐要求

PTX 文档在 Addresses as Operands 章节给出：地址必须自然对齐到访问大小的倍数，不对齐会导致 undefined behavior。

对于 cp.async：

• cp-size=4 需要 4B 对齐
• cp-size=8 需要 8B 对齐
• cp-size=16 需要 16B 对齐

CUDA 文档进一步建议：最佳性能时 shared 与 global 均 128B 对齐。

ordering：cp.async 之间没有默认顺序

如果两条 cp.async 没有用 cp.async.wait_all / cp.async.wait_group 或 mbarrier 指令显式同步，那么它们之间没有 ordering guarantee。

这意味着：

• 不能把我在同一线程里按顺序写了两条 cp.async当作完成顺序与可见性顺序的证明
• 异步操作的完成与可见性不自动跟随源码顺序（program order）
• completion 机制是我们唯一可靠的把异步搬运变成可消费数据的手段

cp.async 被 PTX 视为 weak memory operation，它提供的 ordering guarantee 比普通 load/store 更弱。

最小 PTX 示例

// basic: copy 16 bytes from global to sharedcp.async.ca.shared.global [shrd], [gbl], 16;// with L2 prefetch hintcp.async.ca.shared.global.L2::128B [shrd], [gbl+64], 16;// with src-size for tail handling (copy only %remaining_bytes, zero-fill rest)cp.async.ca.shared.global [shrd], [gbl], 16, %remaining_bytes;// with ignore-src predicate for out-of-bounds handlingcp.async.ca.shared.global [shrd], [gbl], 16, %out_of_bounds;

这些示例只展示指令形态。完整正确用法必须配套 completion 机制。

Async-group 机制

PTX 异步拷贝 completion 机制分为两类：async-group mechanism 和 mbarrier-based mechanism。async-group mechanism 是 cp.async 最直接的等待方式。

核心概念：per-thread 的组

async-group 机制的关键特征是 per-thread：它以执行线程为单位建模，不是 warp 级也不是 block 级。

每个线程维护自己的 async-group 序列。当我们调用 commit_group 时，它会把该线程此前发起、且尚未被提交的异步操作归入一个新的 group。

per-thread 含义

这是一个关键的概念区分：

线程 0 的视角:                    线程 1 的视角:┌─────────────────┐              ┌─────────────────┐│ cp.async A       │              │ cp.async X       ││ cp.async B       │              │ cp.async Y       ││ commit_group ────│── group_0    │ commit_group ────│── group_0│ wait_group 0     │              │ wait_group 0     │└─────────────────┘              └─────────────────┘       │                                   │       │ 线程 0 的 wait_group              │ 线程 1 的 wait_group       │ 只等待 A 和 B                    │ 只等待 X 和 Y       │ 不会等待 X 和 Y！                │ 不会等待 A 和 B！       ▼                                   ▼

关键结论：如果线程 0 需要读取线程 1 加载的数据，wait_group 是不够的！ 必须在 wait_group 之后加 __syncthreads() 来建立跨线程的同步。

Thread 0 的视角:  cp.async [dst0], [src0], 16    ──┐  cp.async [dst1], [src1], 16    ──┼─── commit_group ───> group_0  cp.async [dst2], [src2], 16    ──┘  cp.async [dst3], [src3], 16    ──┐  cp.async [dst4], [src4], 16    ──┼─── commit_group ───> group_1

cp.async.commit_group

cp.async.commit_group 的行为：

• 为执行线程创建一个新的 cp.async-group
• 把该线程此前未提交的 cp.async 批量归入这个 group
• 如果没有未提交的操作，会产生一个 empty group（这个 empty group 被视为 trivially complete）

语法就是一行：

cp.async.commit_group;

cp.async.wait_group N

这是最容易误读的指令。

cp.async.wait_group N 的语义是：执行线程等待，直到最近的 N 个 cp.async-groups 仍可能 pending，且更早提交的 groups 都已完成。

换句话说，它允许最近 N 个 group 还没完成，但保证比最近 N 个更早的 group 都完成了。

举个例子，假设线程提交了 group_0、group_1、group_2 三个 group：

// group_0 已提交// group_1 已提交  // group_2 已提交cp.async.wait_group 1;// 此时：group_0 和 group_1 保证完成//       group_2 可能仍在进行（最近 1 个允许 pending）

特殊情况：

• wait_group 0 等待所有 prior groups 完成
• N 是整数常量，不是 runtime 值

cp.async.wait_all

cp.async.wait_all 等价于：

cp.async.commit_group;cp.async.wait_group 0;

先把当前未提交的操作归入一个 group，然后等待所有 group 完成。这是最简单的等待所有东西完成的方式。

Group 的完成顺序

PTX 给出两条规则：

• Group 之间有序：同一线程提交的 async-groups 按提交顺序完成。如果 group_0 先于 group_1 提交，那么 group_0 一定先于或同时与 group_1 完成。

• Group 内无序：同一 group 内的异步操作之间没有 ordering guarantee。我们不能假设 group 内的第一条 cp.async 一定先于第二条完成。

可见性语义

Memory Visibility (内存可见性)

在并发编程中，一个线程写入内存后，另一个线程（或同一线程的后续操作）不一定能立即看到这个写入。这涉及两个概念：

• 完成（Completion）：写操作在硬件层面执行完毕
• 可见（Visibility）：写入的数据对特定观察者可读

线程 A:                    线程 B:  写入 x = 1                 读取 x       │        ←─ 可见性问题 ─→         │       │    即使 A 写完了，B 可能          │       │    看到的还是旧值                │

对于 cp.async，问题更复杂：异步操作的完成与可见性不自动跟随线程的 program order，所以即使代码上后面的指令也不能自动看到异步写入的结果。只有通过 completion 机制（wait_group、mbarrier 等），才能建立异步写入→后续读取的可见性保证。

cp.async 写入对执行线程可见的时机，PTX 给出三种情况：

• wait_all 完成后
• 所属 group 的 wait_group 完成后
• 跟踪该 cp.async 的 mbarrier.test_wait 返回 True（后续章节会讲）

重要限制：cp.async.wait_group / wait_all 仅对 cp.async 提供完成与可见性语义，不对其他内存操作提供 ordering/visibility 保证。

常见陷阱

陷阱 1：把 commit/wait 当成 warp/block 级

PTX 在描述中反复使用 executing thread，并定义 per-thread group。每个线程独立维护自己的 group 序列，一个线程的 wait_group 只等待它自己的 group。

如果我们希望 block 内所有线程的数据都就绪，需要在 wait_group 之后再加 __syncthreads()。

陷阱 2：把 wait_group N 误读为等待第 N 组

wait_group 2 不是等待 group_2，而是允许最近 2 个 group 仍 pending。这个语义对于实现 multi-stage pipeline 很重要：我们可以让最近几个 stage 的数据仍在搬运，同时使用更早 stage 的数据做计算。

陷阱 3：同一 group 内多次写同一位置

如果同一 group 内有多条 cp.async 写入同一 location，可能导致 UB。因为 group 内无 ordering guarantee，最终结果是不确定的。

陷阱代码对比：

//  错误：wait_group 后没有 __syncthreads()// 线程 0 可能读取线程 1 还未加载完成的数据__pipeline_memcpy_async(&smem[tid], &global[tid], sizeof(float));__pipeline_commit();__pipeline_wait_prior(0);float val = smem[other_tid];  // UB! other_tid 的数据可能还没到//  正确：wait_group + __syncthreads()__pipeline_memcpy_async(&smem[tid], &global[tid], sizeof(float));__pipeline_commit();__pipeline_wait_prior(0);__syncthreads();  // 确保所有线程的拷贝都完成float val = smem[other_tid];  // OK!//  错误：误读 wait_group N 的语义// 以为 wait_group 2 是 等待 group_2__pipeline_commit();  // group_0__pipeline_commit();  // group_1__pipeline_commit();  // group_2__pipeline_wait_prior(2);  // 这不是等待 group_2！// 而是允许最近 2 个仍 pending// 即 group_0 完成，group_1/2 可能未完成//  正确理解：wait_group N = 允许最近 N 个 pending__pipeline_commit();  // group_0__pipeline_commit();  // group_1__pipeline_commit();  // group_2__pipeline_wait_prior(0);  // 等待所有 group 完成

与 CUDA primitive 的对应关系

CUDA 的 __pipeline_wait_prior(N) 与 PTX 的 wait_group N 语义类似：对该线程的 commit 序列（索引 0..L），等待至少直到 batch L-N 完成。

// CUDA side__pipeline_memcpy_async(...);__pipeline_memcpy_async(...);__pipeline_commit();  // batch 0__pipeline_memcpy_async(...);__pipeline_commit();  // batch 1__pipeline_wait_prior(1);  // wait until batch 0 is done// batch 1 may still be in progress

完整示例：两阶段预取

// Stage 0: initiate first batchcp.async.ca.shared.global [shrd0], [gbl0], 16;cp.async.ca.shared.global [shrd0+16], [gbl0+16], 16;cp.async.commit_group;  // group_0// Stage 1: initiate second batchcp.async.ca.shared.global [shrd1], [gbl1], 16;cp.async.ca.shared.global [shrd1+16], [gbl1+16], 16;cp.async.commit_group;  // group_1// Wait for group_0 to complete (allow group_1 still pending)cp.async.wait_group 1;// Now safe to use shrd0, while shrd1 data is still being loaded// ... compute using shrd0 ...// Wait for allcp.async.wait_group 0;// Now safe to use shrd1

这种模式是 software pipelining 的基础：计算当前 stage 的数据，同时预取下一个 stage 的数据。

CUDA C++ 版本：两阶段预取完整示例

下面是上述 PTX 示例的 CUDA C++ 等价实现。与前文「Double Buffering」示例的区别在于：这里假设单 block 顺序处理所有 tiles，适合教学演示；前文版本则是多 block 并行，适合实际部署。

前置条件：本示例假设使用单个 block 处理所有 tiles。

关键点：软件流水线的核心不在于 API 是 C++ 还是 PTX，而在于 prologue（预取） -> loop（计算+预取next） -> epilogue（收尾） 的结构。详细代码结构请参考前文「Double Buffering」示例，只需注意本例是单 block 串行处理即可。

TMA 与 Tensor Map

到目前为止，我们讨论的 cp.async 是 element-wise 的异步拷贝：一次拷贝的基本粒度是 4/8/16 字节。当工作负载具备以下特征时，这一代机制会显得表达力不够：

• 需要搬运较大块数据（bulk），希望减少指令数量与控制开销
• 数据在 global 中是多维布局（2D/3D/更高维），而 shared 里需要的是某个 tile
• 每次 tile 搬运都要重复做大量地址计算、stride 处理、越界判断
• 希望把搬运视为一个可管理的后台事务，有清晰的 completion 语义

从 Hopper 架构（sm_90，compute capability 9.0）开始，CUDA 引入了 TMA（Tensor Memory Accelerator），把 bulk-asynchronous copy 和 bulk-tensor asynchronous copy 作为一等能力。

TMA 解决什么问题

Bulk vs Element-wise 核心区别

在理解 TMA 之前，先明确 bulk 与 element-wise 的核心区别：

Element-wise (cp.async)：每个线程发起自己的小块拷贝（4/8/16B），需要多条指令，适合细粒度控制。

Bulk (TMA)：一条指令传输整个 tile（可达百 KB），由硬件负责整个传输过程，减少指令开销。

TMA 的定位是：很多应用需要搬运大量数据且 global 中常是多维布局；把 sub-tiles 拷贝到 shared 需要地址计算，容易出错重复；TMA 旨在 offload 这些计算并提供高效搬运机制。

这不是说 cp.async 不好用，而是在不同规模和维度下各有适用场景：

TMA 的两条主线

在 CUDA 文档的分类里，TMA 相关能力可以组织成两条主线：

主线 1：bulk-asynchronous copies（1D contiguous）

• 面向一维连续数组的 bulk 拷贝
• 不需要 tensor map（descriptor）
• 可以用 pointer + size 参数在 device 侧执行

主线 2：bulk-tensor asynchronous copies（multi-dimensional tile）

• 面向多维张量的 tile 拷贝
• 需要一个 tensor map 来描述 global/shared 布局与访问属性
• 支持 1D 到 5D 的多维度

这个二分很重要：是否需要 descriptor、对齐/stride 的硬约束、以及 completion 的表达方式，都随1D vs multi-D而变化。

Tensor Map 是什么

PTX ISA 对 tensor map 给出明确定义：

• 128-byte 的 opaque object
• 可以位于 .const、.param（kernel 参数）或 .global 空间
• 用于描述 tensor 的属性与访问属性

术语统一：CUDA 文档中会出现 tensor map、tensor descriptor、CUtensorMap 等不同表述，它们指的是同一个对象。以下统一使用 tensor map。

Opaque (不透明) 含义

Opaque(不透明) 意味着：

• 我们不能直接读写它的内部字段
• 必须通过特定 API创建和使用它
• 内部结构可能随架构/驱动版本变化

CUtensorMap tensor_map;  // 128 字节的“黑盒子”//  错误：不能直接访问内部// tensor_map.base_addr = ptr;  // 编译错误！//  正确：通过 API 创建cuTensorMapEncodeTiled(&tensor_map, ...);

这种设计的好处是：NVIDIA 可以在不破坏 API 兼容性的情况下优化内部实现。

tensor map 的作用是让 tensor copy / tensor prefetch 指令能够理解多维布局 + 访问模式：

• global tensor 的 base address、各维度 size、各维度 stride
• shared memory 的 box 大小（即 tile 大小）
• element type、swizzle mode、out-of-bounds 填充模式等

Host 侧创建 Tensor Map

tensor map 的创建通常发生在 host 侧，使用 CUDA Driver API。示意性伪代码如下（具体参数单位与返回码处理请以我们所用 CUDA/Driver API 版本文档为准）：

tensor_map: CUtensorMapdtype       = CU_TENSOR_MAP_DATA_TYPE_FLOAT32rank        = 2base_ptr    = d_matrix// fastest dimension first (row-major 2D: x then y)global_dim    = {cols, rows}// rank-1 entries; unit (bytes vs elements) is defined by the Driver API// common constraint: 16B alignment (see the Driver API + Programming Guide)global_stride = {stride0}box_dim        = {tile_cols, tile_rows}element_stride = {1, 1}cuTensorMapEncodeTiled(  &tensor_map,  dtype, rank,  base_ptr,  global_dim,  global_stride,  box_dim,  element_stride,  interleave, swizzle, l2_promotion, oob_fill)

几个要注意的点：

• 维度顺序：fastest moving dimension comes first。对于 row-major 的 2D 矩阵，column（x）在前，row（y）在后。

• stride 约束：常见约束是 stride（按文档定义的单位，通常在描述中以字节约束）需要满足 16B 对齐。实际行为（编码失败、返回错误、或后续使用时行为未定义）以我们所用 CUDA/Driver API 版本的官方文档为准。

• box_dim：对应 shared memory 中的 tile 大小，决定了每次 TMA 操作搬运的数据量。

• 最后四个参数的含义：

简化建议：如果我们刚开始使用 TMA，建议所有这些参数都使用 NONE 变体。这是最安全的默认值，性能优化可以后续再调整。

Tensor Map 传递到 Device

创建好 tensor map 后，需要让 device 代码能访问它。常用的三种方式：

方式 1（推荐）：作为 const grid_constant kernel 参数

__global__ voidkernel(const __grid_constant__ CUtensorMap tensor_map){// use tensor_map directly}// launchkernel<<<grid, block>>>(tensor_map);

这是最推荐的方式。__grid_constant__ 告诉编译器这个参数在整个 grid 执行期间不会被修改。

方式 2：复制到 device constant

__constant__ CUtensorMap d_tensor_map;// on hostcudaMemcpyToSymbol(d_tensor_map, &tensor_map, sizeof(CUtensorMap));

方式 3：放到 global memory

CUtensorMap* d_tensor_map;cudaMalloc(&d_tensor_map, sizeof(CUtensorMap));cudaMemcpy(d_tensor_map, &tensor_map, sizeof(CUtensorMap), cudaMemcpyHostToDevice);

如果放到 global memory，本身并不意味着“每个 thread block 使用前都必须做 fence”，需要区分写入来源与是否存在 device 端修改：

• 如果 tensor map 仅由 host 在 kernel launch 前初始化并写入，然后在 kernel 中只读使用，通常不需要额外的 device fence。
• 如果 tensor map 在 device 端被写入/修改，且随后会被其他线程/CTA 读取，则必须按文档要求插入正确的同步与 fence（并注意 async/generic proxy 相关语义）。

此外，这种方式往往更慢：tensor map 需要从 global memory 读取，可能增加额外的访存开销。

PTX 中的 tensormap.replace

如果需要在 device 侧修改 tensor map 的某些字段（比如修改 global_address 指向下一个 batch），可以使用 PTX 的 tensormap.replace 指令。

可替换的字段包括：

• .global_address、.rank
• .box_dim、.global_dim、.global_stride、.element_stride
• .elemtype、.interleave_layout、.swizzle_mode、.swizzle_atomicity、.fill_mode

特别注意：.rank 字段是 zero-based 编码。如果我们想表示 2D tensor，new_val 应该是 1（2-1=1），不是 2。

另一个需要注意的点：PTX 的 .elemtype 数值编码不对应 CUDA driver API 的 CUtensorMapDataType enum 值。不要把它们混用。

On-device encode pattern

CUDA 文档给出了一种 on-device 动态修改 tensor map 的推荐模式：

• Host 创建一个 template tensor map
• 把 template 传入 kernel
• Device 端用 tensormap.replace 修改需要变化的字段
• 写入 global memory
• 执行 fence
• 使用修改后的 tensor map

这种模式适合需要动态改变搬运目标的场景，但增加了复杂性。大多数情况下，host 侧创建 + __grid_constant__ 参数传递就够用了。

Tensormap prefetch

PTX 的 prefetch 指令支持 prefetch.const.tensormap [ptr]，用于把 tensor map 的 cache line 预取到后续使用所需的状态。这需要 sm_90 或更高。

prefetch.const.tensormap [tensor_map_ptr];// ... later use tensor_map_ptr in cp.async.bulk.tensor ...

这是一个优化手段，可以减少首次使用 tensor map 时的 latency。

TMA 传输模式

TMA 支持多种传输模式：1D contiguous、multi-dimensional tile、以及一些特殊的 load mode。

1D Bulk-asynchronous：无需 Tensor Map

1D contiguous bulk-asynchronous copy 是最简单的 TMA 形态：我们只需要提供 pointer + size，不需要 tensor map。

CUDA 提供了几种 API：

// Option 1: cuda::memcpy_async (conditional TMA)// Falls back to sync copy if alignment/size not metcuda::memcpy_async(group, dst_shared, src_global, size, barrier);// Option 2: cuda::device::memcpy_async_tx (always TMA)// UB if alignment/size requirements not metcuda::device::memcpy_async_tx(dst_shared, src_global, barrier, size);// Option 3: PTX intrinsic (always TMA)// UB if alignment/size requirements not metcuda::ptx::cp_async_bulk(    cuda::ptx::space_shared, dst_shared,    cuda::ptx::space_global, src_global,    size, barrier);

Option 1 检查对齐和 size 条件，不满足时 fallback 到同步拷贝。Option 2 和 3 总是使用 TMA，条件不满足则 UB。

1D bulk 对齐要求

TMA 1D bulk 完整示例（CUDA 12.0+，sm_90+）

前置条件：本示例假设 blockDim.x == 256，n 是 256 的整数倍，且 src 地址 16B 对齐。对于非对齐或 partial tile 情况，cuda::memcpy_async 会 fallback 到同步拷贝。

#include<cuda/barrier>#include<cooperative_groups.h>namespace cg = cooperative_groups;// 使用 TMA 1D bulk 拷贝的完整示例__global__ voidtma_1d_bulk_example(constfloat* __restrict__ src,float* __restrict__ dst,int n){// Step 1: 定义 barrier 和分配 shared memoryusingbarrier_t = cuda::barrier<cuda::thread_scope_block>;    __shared__ barrier_t barrier;// 注意：shared memory 必须 16B 对齐（1D bulk 要求）__shared__ alignas(16)float smem[256];    cg::thread_block block = cg::this_thread_block();int tid = threadIdx.x;int offset = blockIdx.x * 256;// Step 2: 初始化 barrierif (tid == 0) {        init(&barrier, block.size());    }    __syncthreads();// Step 3: 发起 TMA 1D bulk 拷贝// 使用 cuda::memcpy_async，会自动检查对齐条件// 如果对齐/size 不满足，会 fallback 到同步拷贝size_t copy_size = 256 * sizeof(float);  // 必须是 16B 的倍数    cuda::memcpy_async(block, smem, src + offset,                        cuda::aligned_size_t<16>(copy_size), barrier);// Step 4: 等待拷贝完成    barrier.arrive_and_wait();    __syncthreads();// Step 5: 使用 shared memory 中的数据if (offset + tid < n) {        dst[offset + tid] = smem[tid] * 2.0f;    }}

关键点：

• alignas(16) 确保 shared memory 16B 对齐
• cuda::aligned_size_t<16>(copy_size) 告诉编译器 size 是 16B 对齐的
• 1D bulk 不需要 tensor map，直接用 pointer + size

Multi-dimensional Bulk-tensor：需要 Tensor Map

当我们需要搬运多维 tensor 的 tile 时，必须使用 tensor map。这也是 TMA 最强大的形态。

PTX 指令 cp.async.bulk.tensor 支持 1D 到 5D：

cp.async.bulk.tensor.dim.dst.src.completion [dstMem], [tensorMap, tensorCoords], [mbar];

其中 .dim 可以是 .1d、.2d、.3d、.4d、.5d。

CUDA 侧对应的 API 是 cuda::ptx::cp_async_bulk_tensor。

Multi-dimensional 对齐要求（更严格）

需注意：shared 对齐要求从 16B 升至 128B。如果把 1D 的对齐假设直接迁移到 multi-dimensional，可能会遇到问题。

tensorCoords：坐标向量

cp.async.bulk.tensor 使用 tensorCoords 来指定要搬运的 tile 在 global tensor 中的位置。坐标是 tile 的 top-left corner（左上角）。

坐标向量的格式：

• 元素类型是 .s32（signed 32-bit integer）
• 长度取决于 .load_mode
• 对于大多数模式，长度等于 .dim（2D tensor 就是 2 个坐标）
• 某些特殊模式（如 .tile::scatter4、.tile::gather4）使用长度为 5 的向量

坐标遵循 fastest-dimension-first 的约定。对于 row-major 2D 矩阵：

• tensorCoords[0] 是 column index（x）
• tensorCoords[1] 是 row index（y）

具体示例：加载 1024x1024 矩阵的一个 tile

// 假设：// - 矩阵大小: 1024 行 x 1024 列 (row-major)// - Tile 大小: 32 行 x 32 列// - 我们要加载从 (col=256, row=512) 开始的 tile// tensor map 创建时：// global_dim = {1024, 1024}  // {cols, rows}，fastest first// box_dim = {32, 32}         // {tile_cols, tile_rows}// PTX 指令中：int32_t tensorCoords[2];tensorCoords[0] = 256;  // column index (x) - fastest dimensiontensorCoords[1] = 512;  // row index (y)// cp.async.bulk.tensor.2d ... [tensorMap, {256, 512}] ...// 这会加载:// - 列: 256 到 287 (32 列)// - 行: 512 到 543 (32 行)// 共 32x32 = 1024 个元素

负坐标示例（边界处理）：

// 加载从 (-8, 90) 开始的 32x32 tileint32_t tensorCoords[2] = {-8, 90};// tensor 是 100x100，tile 是 32x32：//// 列方向 (x): 从 -8 开始，共 32 列 → -8 到 23//   - 列 -8 到 -1 (8列) 越界 → zero-fill//   - 列 0 到 23 (24列) 有效//// 行方向 (y): 从 90 开始，共 32 行 → 90 到 121//   - 行 90 到 99 (10行) 有效//   - 行 100 到 121 (22行) 越界 → zero-fill

Out-of-bounds 处理

TMA 越界语义明确，不同于 cp.async。

Global → Shared 方向：

• tile 如果部分越界，对应的 shared 区域会被 zero-filled
• top-left indices（坐标）可以为负

允许负坐标是一个有用的特性：当我们处理 halo region 或 padding 时，可以统一用 tile-space 坐标，不需要在 kernel 里做特殊的边界判断。

示例: tensor 是 100x100, 加载从 (-8, 90) 开始的 32x32 tile列范围: -8 to 23 (32列)  ├── 越界: -8 to -1 (8列) → shared 中这部分填 0  └── 有效: 0 to 23 (24列) → 从 tensor 拷贝行范围: 90 to 121 (32行)  ├── 有效: 90 to 99 (10行) → 从 tensor 拷贝  └── 越界: 100 to 121 (22行) → shared 中这部分填 0

Shared → Global 方向：

• tile 可以部分越界
• 但 top-left indices 不能为负

注意：读可负，写不可负。

Load Mode

PTX 的 cp.async.bulk.tensor 支持多种 load mode，通过 .load_mode 指定：

大多数场景用默认的 .tile 模式就够了。scatter/gather 和 im2col 是针对特定访问模式的优化。

TMA 支持的方向

TMA 不仅支持 global → shared，还支持其他方向：

什么是 Cluster？

Cluster 是 Hopper 架构（sm_90）引入的新概念，是一组协作执行的 CTA（Thread Block）。

传统的 Grid/Block 层次中，CTA 之间无直接通信。Hopper 引入了 Cluster 层次：Grid 包含多个 Cluster，每个 Cluster 包含多个 CTA。Cluster 内的 CTA 可以直接访问彼此的 shared memory、使用 cluster barrier 同步、以及通过 TMA multicast 共享数据。

Cluster 的典型用途：

• 多个 CTA 需要相同的输入数据 → TMA multicast 一次发送到所有 CTA
• CTA 之间需要共享中间结果 → distributed shared memory

不同方向的 completion 机制不同。

PTX 示例

// 2D tile copy: global -> shared::cta// Uses mbarrier for completioncp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes    [shrd], [tensorMap, {x, y}], [mbar];// 3D tile copy with multicast to clustercp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster    [shrd], [tensorMap, {x, y, z}], [mbar], ctaMask;// Prefetch tensor data to L2cp.async.bulk.prefetch.tensor.2d.L2.global    [tensorMap, {x, y}];

cp.async.bulk.prefetch.tensor

TMA 还提供了 tensor prefetch 能力：

cp.async.bulk.prefetch.tensor.dim.L2.global [tensorMap, tensorCoords];

这是 non-blocking 的，可能发起 tensor 数据到 L2 的异步 prefetch。它不会把数据搬到 shared，只是预热 L2 cache。

使用场景：如果我们知道接下来会访问某个 tile，可以提前发起 prefetch，让后续的 cp.async.bulk.tensor 更快完成。

1D vs Multi-dimensional 对比

选择的基本原则：如果数据是一维连续的，用 1D bulk 更简单。如果需要从多维 tensor 中提取 tile，用 multi-dimensional bulk-tensor。

完成机制配合

如引言所述，异步操作必须显式等待完成才能访问结果。对于 TMA/bulk 操作，completion 机制比 cp.async 更复杂，并且因方向而异。

mbarrier 基础概念

在讨论 TMA 的 completion 机制之前，先理解 mbarrier 基础概念。

mbarrier 基础

mbarrier 是一个存储在 shared memory 中的同步对象，用于协调线程间的同步。它与普通的 __syncthreads() 有本质区别：

mbarrier 核心概念

• Phase（阶段）：mbarrier 在不同的 phase 之间循环。每个 phase 可以被 arrive 和 wait。

• Arrive（到达）：线程调用 arrive 表示我完成了某项工作。

• Wait（等待）：线程调用 wait 等待当前 phase 完成。

• Transaction Count（事务计数）：对于 TMA，mbarrier 可以追踪期望的字节数和已完成的字节数。

mbarrier 的 Phase 在 0 和 1 之间循环：每个 phase 内线程可以 arrive/wait，phase 完成后自动切换到下一个 phase，如此往复。

mbarrier 与 TMA 配合逻辑

// 概念性流程：// 告诉 barrier 我期望收到 N 字节mbarrier_expect_tx(mbar, expected_bytes);// 发起 TMA 操作，指定完成后通知 mbarcp_async_bulk_tensor(..., mbar);  // 完成后自动 arrive_tx(actual_bytes)// 等待 phase 完成（expected_bytes == completed_bytes）mbarrier_wait(mbar);

mbarrier 完整 CUDA C++ 示例

完整使用 cuda::barrier 配合异步拷贝的示例（CUDA 12.0+）：

#include<cuda/barrier>#include<cooperative_groups.h>namespace cg = cooperative_groups;__global__ voidmbarrier_example(constfloat* __restrict__ input,float* __restrict__ output,int n){// Step 1: 定义 barrier 类型和分配空间usingbarrier_t = cuda::barrier<cuda::thread_scope_block>;    __shared__ barrier_t barrier;__shared__ alignas(16)float smem[256];  // 对齐以匹配 cuda::memcpy_async    cg::thread_block block = cg::this_thread_block();int tid = threadIdx.x;// Step 2: 初始化 barrier（只需一个线程执行）if (tid == 0) {// 第二个参数是 expected count：期望多少个 arrive 调用        init(&barrier, block.size());    }    __syncthreads();  // 确保所有线程看到初始化完成// Step 3: 发起异步拷贝，绑定到 barrier// memcpy_async 完成时会自动减少 barrier 的 transaction count（而不是增加 arrive count）    cuda::memcpy_async(block, smem, input, sizeof(float) * 256, barrier);// Step 4: 等待 barrier 的当前 phase 完成// 这会阻塞直到异步拷贝完成    barrier.arrive_and_wait();// Step 5: 现在可以安全使用 smem 中的数据    __syncthreads();  // 确保所有线程同步if (tid < 256) {        output[tid] = smem[tid] * 2.0f;    }}

关键步骤解析：

初始化阶段:  init(&barrier, 256)  ← barrier 期望 256 个 arrive        │发起阶段:  memcpy_async(..., barrier)  ← 拷贝完成时更新 transaction status        │        │ 此时线程可以做其他工作（本例没有）        │等待阶段:  barrier.arrive_and_wait()  ← 每个线程 arrive 并等待        │                       当 256 个 arrive 都到达时，phase 完成        │使用阶段:  smem[tid] 可以安全读取

两类 Completion 机制

如 Async-group 机制一节所述，PTX 异步操作的 completion 分为 async-group 和 mbarrier 两类。对于 TMA/bulk 操作，还引入了 bulk async-group，它与 cp.async 的 async-group 是独立的。

方向决定 Completion 机制

不同方向的 TMA 操作使用不同的 completion 机制：

此差异为模型边界，非性能建议。

本质原因：状态存储位置不同

• mbarrier (Shared Memory)：完成状态存储在 shared memory 对象中（物理可见性覆盖整个 block）。就像挂在墙上的钟，block 内任何人都能看。
• Async-group (Register/Hardware State)：完成状态维护在发起线程的私有硬件计数器中。就像戴在手上的表，只有自己能看。

如果 Thread B 尝试等待 Thread A 发起的 shared -> global 拷贝，Thread B 的私有计数器是 0，指令会立即返回（认为已完成），导致严重的 data race。

mbarrier::complete_tx::bytes

对于 global → shared 方向，TMA 使用 .mbarrier::complete_tx::bytes 作为 completion 机制。当异步拷贝完成时：

• 对指定的 mbarrier 执行 complete-tx 操作
• completeCount 等于拷贝的字节数
• 通过 generic-proxy 访问 mbarrier operand（关于 proxy 的详细解释见后文 Proxy 小节）

// TMA with mbarrier completioncp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes    [shrd], [tensorMap, {x, y}], [mbar];

在 CUDA 侧，这对应 transaction barrier 的概念。mbarrier 会追踪期望的总字节数和已完成的字节数，当两者相等时，phase 完成。

Transaction Count 配合

使用 TMA 的 mbarrier completion 时，需要配合 transaction count：

如果使用 cuda::memcpy_async： 会自动处理 expect_tx

如果使用 cuda::device::memcpy_async_tx 或 cuda::ptx::cp_async_bulk： 需要显式调用 expect_tx

CUDA 12.x Transaction Barrier API

以下是 CUDA 12.x（libcu++ 2.0+）中的标准写法：

#include<cuda/barrier>// 定义 barrier 类型（支持 transaction count）usingbarrier_t = cuda::barrier<cuda::thread_scope_block>;__shared__ barrier_t barrier;// 初始化 barrier（在 block 内一个线程执行）if (threadIdx.x == 0) {    init(&barrier, blockDim.x);  // expected count = block size}__syncthreads();// 告诉 barrier 期望收到多少字节（CUDA 12.x）// barrier.expect_transaction(expected_bytes); // (API 随版本变化，此处为示意)// 等待 phase 完成barrier.arrive_and_wait();

// Example with explicit transaction count (概念性伪代码)auto& barrier = barriers[stage];// Pseudocode: the concrete API names vary by CUDA version.// 1) Tell the barrier how many bytes to expect for this phase.barrier_expect_tx(barrier, bytes);// 2) Issue the TMA operation that will report bytes on completion.issue_tma_copy(dst_shared, src_global, bytes, barrier);// 3) Wait until the barrier observes expected bytes == completed bytes.barrier_wait(barrier);

这里的关键是：barrier 需要知道总共期望多少字节，TMA 完成时会报告实际拷贝了多少字节，当两者匹配时 barrier 的 phase 才能前进。

Bulk Async-group

对于 shared → global 方向，使用 bulk async-group completion。概念流程如下：

• 发起 store 操作（使用 cp.async.bulk 的 shared→global 变体，指定 bulk_group completion）
• 调用 cp.async.bulk.commit_group 提交
• 调用 cp.async.bulk.wait_group 0 等待完成

具体的 PTX 指令语法与修饰符顺序，请参考 PTX ISA 文档中 cp.async.bulk 的相关章节。

cp.async.bulk.commit_group 和 cp.async.bulk.wait_group 与 cp.async 的版本是分开的：

• cp.async.commit_group / wait_group 用于非 bulk 的 cp.async
• cp.async.bulk.commit_group / wait_group 用于 bulk 操作

wait_group.read：只等读完成

cp.async.bulk.wait_group 有一个 .read 变体：

cp.async.bulk.wait_group.read 0;

默认的 wait_group 等待完整完成：读 tensormap、读 source、写 destination、并对执行线程可见。

.read 变体只等待到读完 tensormap + 读完 source locations。这意味着：

• 可以安全覆写 source 或推进下一轮
• 但 destination 的写入可能还未完成

使用场景：当我们只需要确保 source 已经读完，可以安全修改 source 时，用 .read 可以减少等待时间。

异步操作的 Ordering 语义

cp.async / cp.async.bulk 等产生的异步操作，其完成与可见性不自动跟随 issuing thread 的 program order。

这意味着：

• 不能把我按顺序写了两条 cp.async.bulk当作它们会按顺序完成的证明
• 异步操作提供的 ordering guarantee 比普通 load/store 更弱
• 只有通过 completion mechanism（async-group 或 mbarrier）才能建立可靠的 happens-before 关系

举例：cp.async 的内存操作不与其他 cp.async 或后续指令有序，除非通过 cp.async.commit_group / cp.async.wait_group（或文档规定的其他 completion 机制）建立同步关系。

隐式 fence 行为

PTX 文档指出 cp.async.bulk 等操作完成后带隐式 generic-async proxy fence，使结果在观察到完成时对 generic proxy 可见。

Proxy (代理)

PTX 内存模型把不同的内存访问方式归类为不同的 proxy。同一内存位置通过不同 proxy 访问时，可能需要显式同步：

不同的 proxy 可能走不同的硬件路径，因此一个 proxy 的写入可能不会立即对另一个 proxy 可见。当一个线程用 TMA 写入、另一个用普通 load 读取时，可能需要 fence.proxy.async。通常 mbarrier completion 的隐式 fence 已足够。

但这个隐式 fence 的排序只相对于同一条异步指令内部的内存操作，不传递性建立对 issuing thread 先前指令的排序。

如果需要跨 proxy 访问同一内存位置（比如一个线程用 TMA 写，另一个线程用普通 load 读），可能需要显式的 fence.proxy.async。

关键词识别表

为了在阅读 PTX 或低层代码时能识别关键点，这里列出与 completion 相关的关键词：

mbarrier 的 phase/parity、arrive/wait 细节请参考 PTX ISA 文档。

TMA 高级特性概览

TMA 还提供一些高级特性：Multicast、Swizzle、OOB 填充模式。这些特性可以进一步优化特定场景的性能，但也带来额外的约束。

Multicast

Multicast 允许一次 TMA 操作把数据从 global 复制到 cluster 内多个 CTA 的 shared memory。

cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster    [shrd], [tensorMap, {x, y}], [mbar], ctaMask;

ctaMask 是 16-bit 掩码，每个 bit 对应 cluster 内的相对 CTA 编号（不是全局 %ctaid）。数据会写入每个目标 CTA 的 shared memory（按 CTA-relative offset）。mbarrier 信号也会 multicast 到相同的 CTA。具体 bit 映射规则以 PTX ISA 对 multicast 的定义为准。

使用建议： multicast 特性通常针对特定目标（例如 sm_90a）做了优化；在其他目标上收益不确定，可能无收益甚至变慢。不要把它当作普适优化，需要在目标架构上实测。

Swizzle

Swizzle 是一种用于减轻 shared memory bank conflict 的技术。TMA 可以在 tensor map 中编码 swizzle mode，让 TMA engine 在写 shared 时自动 swizzle，写回 global 时自动 unswizzle。

Swizzle 的工作原理：

当 swizzle 启用时，数据写入 shared memory 的地址会被变换，目的是减少 bank conflict。TMA engine 在写 shared 时自动 swizzle，写回 global 时自动 unswizzle。

[!CAUTION]Swizzle 会改变 shared memory 中的物理布局！

启用 swizzle 后，shared memory 中的数据不再是直观的连续布局。如果 kernel 侧的读取模式与 swizzle 后的布局不匹配，会导致：

• 读取错误数据（地址偏移不符合预期）
• 性能退化（本想消除的 bank conflict 可能更严重）

只有当后续读 shared 的访问模式与 swizzle 规则匹配时，才能正确使用数据并获得收益。如果不确定，建议使用 CU_TENSOR_MAP_SWIZZLE_NONE。

Swizzle 的具体地址变换公式与约束条件，请参考 CUDA Programming Guide 中关于 TMA Swizzle 的章节（包括 Table 24/25）。这些约束包括：

• Swizzle mode（128B/64B/32B/NONE）决定变换粒度
• Inner dimension 与 repeat 相关的尺寸约束
• Shared/Global 对齐要求
• Swizzle mapping 粒度固定 16B

不满足约束条件时，指令可能被视为 invalid 或行为未定义，具体以文档为准。

在 tensor map 中设置 swizzle：

cuTensorMapEncodeTiled(    &tensor_map,// ... other params ...    CU_TENSOR_MAP_SWIZZLE_128B,  // swizzle mode// ...);

OOB 填充模式

当 TMA 操作越界时，如何填充越界部分？tensor map 的 fill_mode 字段控制这个行为。

PTX 的 tensormap.replace 可以修改 .fill_mode：

CUDA 的 cuTensorMapEncodeTiled 使用 CUtensorMapFloatOOBfill 参数。

需要注意的是：CUDA 文档描述的 global→shared 越界区域 zero-filled 与 PTX 的 fill_mode 设置可能不是完全一一对应的。在使用时应以具体 API 的文档为准，避免混用不同层级的假设。

cp_mask / byteMask

cp.async.bulk 支持 .cp_mask 修饰符，配合 byteMask 操作数：

cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.cp_mask    [dst], [src], size, [mbar], byteMask;

byteMask 是 16-bit 掩码，用于在 16B 粒度内进行选择性复制。具体 bit 的语义与映射规则请以 PTX ISA 文档为准。

这个特性用于需要选择性复制数据的场景。

cta_group

PTX 支持 cta_group::N 修饰符，用于指定操作涉及的 CTA group 大小：

cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes.cta_group::2    [shrd], [tensorMap, {x, y}], [mbar];

这与 cluster 内的协作有关。

Tensor Copy 限制

PTX 文档列出了 tensor copy instructions 的一系列限制，涉及：

• 特定数据类型与 swizzle/OOB-NaN fill 的组合
• Box-Size[0] 与 Tensor-Size[0] 的对齐要求
• tensorCoords 的对齐要求
• 特定架构（如 sm_103a/sm_120a）下的额外限制

这些限制比较细碎，实际使用时建议参考 PTX ISA 文档的具体条目。

何时使用这些高级特性

这些高级特性增加了复杂性。如果基础的 TMA 操作能满足需求，不必急于使用高级特性。