Advanced Performance Optimization in CUDA [S62192]

Igor Terentyev*, NVIDIA DevTech Compute
GPU Technology Conference / March 18th, 2024
* With Guillaume Thomas-Collignon & Athena Elafrou

目录 (Table of Contents)

  1. 议程 (Agenda
  2. 术语 (Nomenclature
  3. 协作网格阵列 (Cooperative Grid Arrays
  4. 内存模型 (Memory Model
  5. 异步屏障 (Asynchronous Barriers
  6. 异步数据拷贝 (Asynchronous Data Copies
  7. 可压缩内存 (Compressible Memory
  8. CUDA Graphs
  9. 总结
  10. 核心性能优化技术

议程 (Agenda)

  • 协作网格阵列 (Cooperative Grid Arrays)
  • 内存模型 (Memory Model)
  • 异步屏障 (Asynchronous Barriers)
  • 异步数据拷贝 (Asynchronous Data Copies)
  • 可压缩内存 (Compressible Memory)
  • CUDA 图 (CUDA Graphs)

术语 (Nomenclature)

  • CTA (Cooperative Thread Array) == 线程块 (Thread Block)
  • CGA (Cooperative Grid Array) == 线程块集群 (Thread Block Cluster)

代码片段:

namespace cg = cooperative_groups;

协作网格阵列 (Cooperative Grid Arrays)

简介

增加的共享内存(SMEM)可以在许多算法中节省对全局内存(GMEM)的访问。下图展示了从 Kepler 到 H100 架构,每个流式多处理器(SM)的共享内存大小的演变,以及 H100 架构上引入的分布式共享内存(Distributed SMEM)概念,其在一个大小为16的集群中可达到 3648 KB。

Page 5, SMEM size evolution from Kepler to H100, and the concept of Distributed SMEM.
Page 5, SMEM size evolution from Kepler to H100, and the concept of Distributed SMEM.

属性

引入了新的层级结构:线程(T) < 线程块(CTA) < 协作网格阵列(CGA) < 网格(GRID)。

集群 (Clusters):
* 保证支持最多 8 个 CTA / CGA。
* 如需超过8个(在 Hopper 架构上最多16个),需要显式选择加入。
* CTA 可以在逻辑上组织成 3D 块:
* 线性排名: cg.this_cluster().block_rank()
* 3D 排名: cg.this_cluster().block_index()

  • 集群内的 CTA 是协同调度(co-scheduled)的。

线程同步 (Thread synchronization):
* CTA 内同步: CTA → cg.this_block().sync()
* CGA 内同步: CGA → cg.this_cluster().sync()

潜在的资源利用不足:
* 当集群大小 > 2 时,可能会导致未使用的 SM。

Page 6, Diagram illustrating the structure of Cooperative Grid Arrays (CGA).
Page 6, Diagram illustrating the structure of Cooperative Grid Arrays (CGA).

分布式共享内存 (DSMEM) 属性

  • 分布式共享内存 (DSMEM): 集群内的共享内存可被所有 CTA 访问,形成 DSMEM。

    • 例如,H100 的总 DSMEM 大小约为 3.6 MB。
  • DSMEM 允许的操作:

    • 访问: 加载(LD) / 存储(ST) / 原子操作(ATOM)
    • 异步拷贝到远程 DSMEM
    • 多播 (Multicasting)
  • 注意: 远程 DSMEM 比本地 SMEM 慢:延迟低,但带宽相对较低。

  • 远程 DSMEM 必须进行映射:
__shared__ T smem;
auto dsmem_ptr = 
  cg::this_cluster().map_shared_rank(&smem, rank);
  • 重要: 在访问 DSMEM 时,必须确保远程 CTA 仍然存活!(相应地使用 cluster.sync())
Page 7, Diagram showing Distributed Shared Memory (DSMEM) within a CGA.
Page 7, Diagram showing Distributed Shared Memory (DSMEM) within a CGA.

启动支持 CGA 的核函数 (Launching kernels with CGA support)

编译时 CGA (Compile-time CGA):

constexpr int X = 2;
constexpr int Y = 2;
constexpr int Z = 2;
__cluster_dims__(X, Y, Z) void kernel();
dim3 grid = {...};
assert(grid.x % X == 0 && grid.y % Y == 0 && grid.z % Z == 0); //否则启动失败
kernel<<<grid_size,...>>>();

运行时 CGA (Run-time CGA):

__global__ void kernel(); // 和通常一样
cudaLaunchAttribute attribute[1];
attribute[0].id = cudaLaunchAttributeClusterDimension;
attribute[0].val.clusterDim.x = X;
attribute[0].val.clusterDim.y = Y;
attribute[0].val.clusterDim.z = Z;
assert(grid.x % attribute[0].val.clusterDim.x == 0 && ...); // 否则启动失败
cudaLaunchConfig_t config = {0};
config.attrs = attribute;
config.numAttrs = 1;
cudaLaunchKernelEx(&config, kernel, ...);

内存模型 (Memory Model)

为什么需要内存模型?

  • 为什么我们需要内存(一致性)模型?
  • 在多线程世界中,内存操作的行为可能很“奇怪”。

弱内存模型 (Weak memory model)

CUDA 使用弱内存模型。这意味着由一个线程执行的写操作,可能不会以相同的顺序被其他线程观察到。

考虑以下代码:
* 线程 0:

data = 42;
flag = 1;
  • 线程 1:
while (!flag) {}
assert(data == 42); // 不保证成立!!

在上述例子中,assert 语句不保证会成功。

弱内存模型与未定义行为

更进一步,上述情况不仅仅是“不保证”,而是未定义行为 (Undefined Behavior, UB)!

  • CUDA 使用弱内存模型:一个线程的写操作不保证被其他线程按序观察到。
  • 对同一内存的并发访问,且至少有一次是“写”操作,这构成了一次冲突 (CONFLICT),会导致未定义行为 (UNDEFINED BEHAVIOR)
  • 编译器假定代码中没有冲突,程序员有责任进行同步。

线程同步的要求

线程之间的同步需要:
1. 强顺序关系 (Strong ordering): "happens before" 关系。
2. 达到正确的相干点 (Point of Coherency)。

Page 14, Diagram illustrating a memory consistency issue between two SMs.
Page 14, Diagram illustrating a memory consistency issue between two SMs.

通用同步模式

使用共享内存的通用同步

一个非常常见的通信模式如下所示:

// 定义一个共享内存缓冲区
__shared__ T buf[1024];

// 线程向共享内存写入数据
buf[ind1] = value1;

// 同步点(Barrier)
// __syncthreads() 确保在此之前的所有写入操作对块内的所有线程可见
__syncthreads();

// 线程读取由其他线程写入的值
value2 = buf[ind2];

这种模式能正常工作,因为它满足两个关键条件:
1. 强顺序性(Strong ordering):建立了“先行发生”(happens before)关系。
2. 达到正确的相干点(Point of Coherency)

使用全局内存的通用同步

一个不太常见的通信模式是使用全局内存:

__global__ void kernel (T* buf)
{
  // 线程向全局内存写入数据
  buf[ind1] = value1;

  // 同步点(Barrier)
  // __syncthreads() 确保在此之前的所有写入操作
  // 对块内所有线程可见,包括对全局内存的写入!
  __syncthreads();

  // 线程读取同一块内其他线程写入的值
  value2 = buf[ind2];
}

__syncthreads() 同样可以同步块内线程对全局内存的访问。

当需要进行全局同步时,可以使用协作组(Cooperative Groups):

__global__ void kernel (T* buf)
{
  // 线程向全局内存写入数据
  buf[ind1] = value1;

  // 同步点(Barrier)
  // cg::this_grid().sync() 确保在此之前的所有写入操作
  // 对网格(grid)内所有线程可见,包括对全局内存的写入!
  cg::this_grid().sync();

  // 线程读取同一块内其他线程写入的值
  value2 = buf[ind2];
}
// 注意:这需要使用协作启动(cooperative launch)

这引出了一个问题:cg::this_grid().sync() 背后究竟发生了什么?

GPU 内存层级与相干点

GPU 内存层级结构图
Page 20: GPU内存层级结构图

上图展示了典型的 GPU 内存层级结构。相干点(Point of Coherency, POC)取决于线程通信的层级。数据在不同层级间(如 SM 内的 L1/共享内存、L2 缓存、DRAM)流动。当 SM 内的线程通信时,数据可能仅需在 L1/共享内存层级保持一致。

SM 间通信路径
Page 21: SM间通信路径示例

当不同 SM 上的线程需要通信时,它们不能直接访问对方的 L1 缓存。数据必须通过更高层级的缓存(如 L2 缓存)来同步,以确保可见性。

cg::this_grid().sync() 的工作机制是:
内存栅栏(Memory fence) + 屏障(barrier, 设备级别) + L1 缓存失效(L1 cache invalidate)

相干点层级
Page 23: 相干点(Point of Coherency)的层级

关于相干点(POC),你需要了解的关键层级包括:
- Self (thread):线程自身
- Block:线程块
- Cluster:线程簇
- Device:设备
- System:系统

细粒度通信与内存排序

问题:生产者/消费者模型

考虑一个生产者/消费者模型,其中线程运行在不同的 GPU 上:

// 生产者
__global__ void producer (T* data, int* flag)
{
  *data = value;   // 写入数据
  *flag = 1;       // 然后设置标志
}

// 消费者
__global__ void consumer (T* data, int* flag)
{
  while (*flag == 0) {}  // 等待标志
  value = *data;         // 然后读取数据
}

这种方式可行吗?答案是:不行。
- 原因1:缺乏强顺序性。编译器或硬件可能会重排 *data = value;*flag = 1; 的执行顺序。
- 原因2:未达到正确的相干点。一个 GPU 上的写入可能对另一个 GPU 不可见。

尝试1:添加内存栅栏

// 生产者
__threadfence_system(); // 内存栅栏
*flag = 1;

// 消费者
while (*flag == 0) {}
__threadfence_system(); // 内存栅栏
value = *data;

可行吗?答案仍然是:不行。
- __threadfence_system() 确保了强顺序性(“先行发生”关系)。
- 但它仍然没有达到正确的相干点。

尝试2:使用 volatile

// 生产者
__global__ void producer (T* data, volatile int* flag)
{
  *data = value;
  __threadfence_system();
  *flag = 1;
}

// 消费者
__global__ void consumer (T* data, volatile int* flag)
{
  while (*flag == 0) {}
  __threadfence_system();
  value = *data;
}

在实践中可行吗?是的。
- 满足强顺序性。
- volatile 关键字强制数据达到相干点。

但这种做法正确吗?不正确。
- volatile 不保证原子性(atomicity)。
- 当至少有一个是写操作时,对同一内存位置的并发访问会产生冲突,导致未定义行为(UNDEFINED BEHAVIOR)

尝试3:使用 CUDA 原子操作

// 生产者
__threadfence_system();             // 内存栅栏
atomicExch_system((int*)flag, 1);   // 然后设置标志

// 消费者
while (atomicAdd_system(flag, 0) == 0) {} // 等待标志
__threadfence_system();                   // 内存栅栏
value = *data;                            // 然后读取数据

可行吗?是的。正确吗?是的。
- 满足强顺序性。
- 达到正确的相干点。

但是...
- 这种实现并非最优雅或最高效的。
- 无法与 CPU 进行兼容的原子操作通信。

尝试4:使用 C++ 原子操作

// 生产者
// cuda::atomic<int, thread_scope_system>* flag
*data = value;              // 写入数据
*flag = 1;                  // 然后设置标志
flag->notify_all();

// 消费者
// cuda::atomic<int, thread_scope_system>* flag
flag->wait(0);              // 等待标志
value = *data;              // 然后读取数据

可行吗?是的。正确吗?是的,很优雅。 C++ 原子操作可以与 CPU 通信。
- C++ 原子操作内置了内存栅栏。
- 满足强顺序性。
- 达到正确的相干点。

但是...
- 这仍然不是最高效的方式。

内存顺序(Memory Order)与 Acquire / Release 语义

<blockquote>

"内存顺序指定了内存访问(包括常规的非原子访问)如何围绕一个原子操作进行排序。"

</blockquote>

std::atomic 的默认行为是顺序一致性(sequentially-consistent)
- 它不仅在原子操作之间强制执行严格的顺序,而且在所有内存操作之间也强制执行严格的顺序。
- 可以将其理解为“对硬件来说执行起来较慢”的方式。

为了优化性能,我们可以使用更宽松的内存顺序。在生产者/消费者模型中,关键的内存排序要求是:
- 生产者端:必须保证在设置 flag 之前的所有写操作(如 *data = value;)不会被重排序到设置 flag 之后。
- 消费者端:必须保证在读取 flag 之后的所有读操作(如 value = *data;)不会被重排序到读取 flag 之前。

Page 31: 生产者/消费者模型的内存排序问题
Page 31: 生产者/消费者模型的内存排序问题

围绕原子操作的细粒度排序:
- memory_model::release:
- 当前线程中的读或写操作不能被移动到 release 操作之后。
- 使当前线程的写操作对其他线程(使用相同原子变量)可见。

  • memory_model::acquire:
    • 当前线程中的读或写操作不能被移动到 acquire 操作之前。
    • 使其他线程(使用相同原子变量)的写操作对当前线程可见。

细粒度通信的优化

通过使用 memory_order_releasememory_order_acquire 语义,可以正确、优雅且高效地解决前述的生产者/消费者问题。

  • 生产者:使用 flag->write(1, memory_order_release)。这确保了在写入 flag 之前的所有内存写入(即 *data = value)都对其他线程可见。
  • 消费者:使用 flag->wait(0, memory_order_acquire)。这确保了在读取 flag 之后的所有内存读取(即 value = *data)都能看到生产者 release 之前写入的数据。

这种方法被描述为 正确、优雅且高效

Page 33: 使用 acquire/release 语义优化的生产者/消费者模型
Page 33: 使用 acquire/release 语义优化的生产者/消费者模型

内存模型:累积性 (Cumulativity)

内存操作的可见性是可传递和累积的。在一个同步链中,由 release 操作发布的数据,可以被后续的 acquire 操作获取。

如下例所示:
1. 设备1 - 线程1: 将 42 写入 *data,然后通过 release 操作将 flag_dev 设置为 1
2. 设备1 - 线程2: 通过 acquire 操作等待 flag_dev,断言 *data 的值为 42,然后通过 release 操作将 flag_sys 设置为 1
3. 设备2 - 线程3: 通过 acquire 操作等待 flag_sys。由于内存可见性的累积效应,线程3也能够断言 *data 的值为 42,即使它同步的 flag_sys 是由线程2设置的。

Page 34: 内存模型累积性示例
Page 34: 内存模型累积性示例

异步屏障 (Asynchronous Barriers)

异步屏障简介

CUDA屏障在指定范围内提供:
- 同步点(包括异步内存事务)。
- 内存排序。

传统 __syncthreads() 与异步屏障的对比

  • __syncthreads():

    • 阻塞所有线程,直到所有线程都到达。
    • __syncthreads() 之前的写操作对 __syncthreads() 之后的所有线程可见。
  • 异步屏障 (barrier.arrive() / barrier.wait()):

    • barrier.arrive(): 标记到达(非阻塞)。

      • arrive 之前的写操作在 wait 之后可见。
    • 可以在 arrivewait 之间执行独立操作(例如,对其他内存的操作)。

    • barrier.wait(): 阻塞直到所有线程都标记到达。
      • arrive 之前的写操作在 wait 之后可见。
Page 36: 异步屏障工作流程图
Page 36: 异步屏障工作流程图

异步屏障详解

Barrier 对象的构成

  • 追踪量 (Tracked quantities):

    • expected_arrival_count (>=0): 期望到达的线程数。
    • transaction_count (>=0) (可选): 事务计数。
  • 状态 (State):

    • arrival_count: 到达计数。
    • phase (0/1): 阶段。
  • 方法 (Methods):

    • arrive: 到达。
    • wait: 等待。

初始化

  • expected_arrival_count: 由用户提供的值。
  • arrival_count: 初始化为 expected_arrival_count
  • phase: 初始化为 0。
  • transaction_count (Hopper++): 由用户提供的值递增,由硬件在内存事务完成时递减(允许同步异步复制)。

操作细节
- Arrive (单个原子非阻塞操作):
1. arrival_count-- (或 -=n)。
2. 如果 arrival_count == 0
- phase ^= 1 当事务计数达到0时
- arrival_count := expected_arrival_count (屏障重置)。

  • Wait (等待阶段翻转的两种方式):

    • 基于令牌 (Token-based):

      • auto token = bar.arrive(); bar.wait(std::move(token));
      • 比显式追踪更昂贵。
    • 显式追踪阶段 (Explicitly tracking the phase):

      • while (!mbarrier_try_wait_parity(bar, curr_phase, 1000)) { curr_phase ^= 1; }
      • 需要显式地维护阶段。
Page 38: 异步屏障的 Arrive 和 Wait 操作
Page 38: 异步屏障的 Arrive 和 Wait 操作

异步屏障总结

下表总结了不同作用域的屏障特性。在Hopper架构上,Cluster 级别的屏障是硬件加速的。

Scope Synchronizes Typical Memory Location Initialization Arrive Wait
Block CTA Local SMEM Single-thread initialization (followed by CTA sync) ALLOWED ALLOWED
Cluster CGA Local SMEM (owner CTA)
Remote DSMEM (after mapping)
Single-thread initialization (followed by CGA sync) ALLOWED ALLOWED
Device GPU GMEM Single-thread initialization (followed by device sync<sup>#</sup>) ALLOWED ALLOWED
System NODE Accessible by all actors:
• GMEM<sup>*</sup>
• Uniform (managed)
Single-thread initialization ALLOWED<sup>*</sup> ALLOWED<sup>*</sup>

注脚:

  • #: 需要协作启动 (cooperative launch)。
  • *: 要访问非所属GPU的内存,必须启用对等访问 (peer access)。
Page 39: 异步屏障特性总结表
Page 39: 异步屏障特性总结表

异步屏障的位置

在生产者/消费者模型中,屏障对象的位置会影响性能。

代码逻辑:
- 生产者: 生产数据 -> 在数据屏障上arrive -> 等待复制屏障。
- 消费者: 等待数据屏障 -> 使用数据 -> 在复制屏障上arrive

实验数据 (在4路 Grace-Hopper 上测试):

  • GPU-GPU 通信: 将屏障放置在消费者(等待方)的内存中时,性能最佳(0.032 vs 0.065)。
  • CPU-GPU 通信: 将屏障放置在GPU(线程数最多的一方)上时,性能最佳(0.27 vs 3.39)。

建议:
- 将屏障放置在等待发生的地方
- 将屏障放置在对大多数线程来说是本地的内存中(例如GPU)。

Page 40: 屏障位置对性能的影响
Page 40: 屏障位置对性能的影响

异步数据拷贝 (Asynchronous Data Copies)

同步数据复制剖析

从全局内存到共享内存的同步复制 (smem[sind] = gmem[gind];)
- 涉及两个步骤:
1. 将数据从全局内存复制到寄存器。
2. 将数据从寄存器复制到共享内存。

  • 这会导致:
    • 长 scoreboard stall:线程等待数据从内存到达。
    • 浪费寄存器和L1带宽。
Page 42: 从全局内存到共享内存的同步复制流程
Page 42: 从全局内存到共享内存的同步复制流程

从共享内存到共享内存的同步复制 (smem[sind1] = smem[sind2];)

  • 同样涉及通过寄存器的两个步骤。
  • 导致 短 scoreboard stall
Page 43: 从共享内存到共享内存的同步复制流程
Page 43: 从共享内存到共享内存的同步复制流程

同步复制导致在途字节数 (bytes in flight) 较少

在典型的循环中,数据加载和计算是串行执行的。这意味着在计算阶段无法为下一次迭代预取数据,限制了内存带宽的利用。

Page 44: 同步复制的串行执行模式
Page 44: 同步复制的串行执行模式

异步数据复制的优势

异步复制允许数据加载和计算操作重叠执行,从而提高效率。

  • 增加在途字节数: 在执行当前迭代的计算时,可以开始加载下一次迭代所需的数据。
Page 45: 异步复制的重叠执行模式
Page 45: 异步复制的重叠执行模式

其他好处:
- 寄存器旁路 (Register bypass): 数据可以直接从内存传输到共享内存,无需经过寄存器。
- 减少L1流量 (Less L1 traffic)
- 更少的MIO压力 (Less MIO pressure) (指令更少)。
- 更好的依赖管理 (Better dependency management)

异步拷贝总结

下图总结了 Ampere 和 Hopper 架构上的异步数据拷贝操作。

Page 46 - 异步数据拷贝操作总结表
Page 46 - 异步数据拷贝操作总结表

表格内容详解:

  • Ampere 架构:

    • GMEM → SMEM (全局内存到共享内存): 使用 ldsm 的异步版本,通过线程局部 (Thread-local) 的 SMEM barrier 进行同步。C++ 接口为 __pipeline_memcpy_async,对应的 SASS 指令为 LDGSTS
  • Hopper 架构:

    • 引入了 TMA (Tensor Memory Accelerator) 来处理块拷贝。
    • GMEM → (D)SMEM (全局内存到[分布式]共享内存): 1D-5D 块拷贝,通过 (D)SMEM barrier 同步。
    • SMEM → GMEM (共享内存到全局内存): 1D 块拷贝,线程局部 (Thread-local) 同步。
    • SMEM → SMEM (共享内存到共享内存): 1D 块拷贝,通过 SMEM barrier 同步。
    • SMEM → (D)SMEM (共享内存到[分布式]共享内存): 通过 (D)SMEM barrier 同步。
    • 这些操作在 C++ 中通过 cuda::device::experimental::cp_async_bulk_* 系列函数实现。

架构图对比:

  • Ampere (LDGSTS): 数据从全局内存直接拷贝到共享内存。
  • Hopper (TMA): 数据从全局内存拷贝到集群 (Cluster) 内的分布式共享内存 (Distributed SMEM)。

Ampere 架构下的 LDGSTS

Page 47 - Ampere (LDGSTS) 异步拷贝
Page 47 - Ampere (LDGSTS) 异步拷贝

上图聚焦于 Ampere 架构中的 LDGSTS 指令,用于实现从全局内存 (GMEM) 到共享内存 (SMEM) 的异步数据拷贝。

  • 操作: smem[sind] := gmem[gind] 的异步版本。
  • 同步: 线程局部 (Thread-local)。
  • C++: __pipeline_memcpy_async
  • SASS: LDGSTS

LDGSTS 剖析

Page 48 - LDGSTS 数据流与特性
Page 48 - LDGSTS 数据流与特性

__pipeline_memcpy_async(&smem[sind], &gmem[gind], sizeof(T))smem[sind] = gmem[gind] 的异步版本。

特性:

  1. 绕过寄存器 (Bypassing registers): 数据直接从 L2 缓存移动到共享内存,不经过寄存器。这可以带来更高的占用率 (occupancy) 并减少寄存器溢出 (spillage)。
  2. 两种模式:
    • L1 旁路模式 (L1 bypass mode): 当 sizeof(T) 和对齐方式为 16B 时触发。优点是不会污染 L1 缓存。
    • L1 访问模式 (L1 access mode): 当 sizeof(T) 和对齐方式为 4B 或 8B 时触发。

数据流: 数据从全局内存加载,经过 L2 缓存,然后根据模式的不同,或直接写入共享内存(L1 旁路),或通过 L1 缓存写入共享内存。操作的完成由记分板 (Scoreboard) 追踪。

Hopper 注意事项: 当 CTA (Cooperative Thread Array) 是 CGA (Cooperative Group Array) 的一部分时,共享内存必须是本地 SMEM。

LDGSTS 的拷贝与同步机制

异步拷贝的同步通过 __pipeline_commit()__pipeline_wait_prior(N) 两个原语 (primitives) 实现。

Page 49 - LDGSTS 同步原语
Page 49 - LDGSTS 同步原语
  • __pipeline_memcpy_async(...): 发出异步加载指令。
  • __pipeline_commit(): (非阻塞) 为之前发出的加载操作创建批次计数依赖屏障 (batch-counting dependency barrier)。当数据被写入共享内存后,该屏障对执行线程解除。
  • __pipeline_wait_prior(N): 阻塞,直到除了最后 N 个 commit 之外的所有 commit 都准备就绪。
Page 50 - __pipeline_wait_prior(N) 示例
Page 50 - __pipeline_wait_prior(N) 示例

上图展示了 __pipeline_wait_prior(N) 的不同用法:
- __pipeline_wait_prior(2): 等待,直到第一批 commit 完成。
- __pipeline_wait_prior(0): 等待所有 commit 完成。

在条件分支中使用 commit 的注意事项

__pipeline_commit 表现为全线程束 (full-warp) 级别的指令,不受线程束分化 (warp divergence) 的影响。

Page 51 - __pipeline_commit 在分化代码中的行为
Page 51 - __pipeline_commit 在分化代码中的行为
  • 如果在条件分支中调用 __pipeline_commit(),即使只有一个分支被执行,其行为也等同于提交了两次(或在完全分化的情况下最多32次)。
  • 后果: 这可能导致对第一批次的过度等待 (overwait)。
  • 建议: 不要在有分化的代码中进行提交 (commit) 操作。

LDGSTS 应用案例: 在条件代码中批处理加载

考虑一个常见的模式,如 Halo-Center-Halo 交换,其中不同的线程根据其 ID 从不同的内存区域加载数据。

基线实现(同步)

Page 52 - 条件加载的基线实现(同步)
Page 52 - 条件加载的基线实现(同步)

在同步实现中,if-else 结构会导致编译器生成多个分支。每个分支中的加载 (LDG) 和存储 (STS) 操作都会导致流水线停顿 (STALL)。

理想情况

Page 53 - 条件加载的理想编译器行为
Page 53 - 条件加载的理想编译器行为

理想情况下,编译器可以使用谓词 (predicates) 将所有加载操作同时发出(ALL IN FLIGHT),从而避免停顿。但编译器可能不会这样做。

使用 LDGSTS 优化

Page 54 - 使用 LDGSTS 优化条件加载
Page 54 - 使用 LDGSTS 优化条件加载

通过使用 __pipeline_memcpy_async,可以将条件加载转换为异步操作。
- 编译器生成 LDGSTS 指令,这些指令不会导致停顿。
- 所有异步加载通过一个 __pipeline_commit() 提交,形成一个批次。
- DEPBAR (Dependency Barrier) 指令会产生一次停顿,等待所有异步加载完成。这远优于多次停顿。

完整代码示例与性能

Page 55 - 条件加载的完整代码示例
Page 55 - 条件加载的完整代码示例

上图是一个更完整的代码示例,增加了一个维度 ty。注意,条件判断基于 ty,而 ty 对于同一线程束内的所有线程是相同的,因此不会产生线程束内部分化。

H200 实验结果

Page 56 - H200 实验结果(Un-staged)
Page 56 - H200 实验结果(Un-staged)
  • 相对于同步拷贝的加速比:

    • Un-staged (非暂存): 1.3倍
  • 内存吞吐量:

    • 同步拷贝 (Sync copy): 46%
    • 非暂存 (Un-staged): 67%
Page 57 - H200 实验结果(Two-stage)
Page 57 - H200 实验结果(Two-stage)

通过实现两阶段(Two-stage,即双缓冲)流水线,性能得到进一步提升。

  • 相对于同步拷贝的加速比:

    • Two-stage (两阶段): 1.7倍
  • 内存吞吐量:

    • 两阶段 (Two-stage): 84%

Nsight Compute (NCU) 分析

基线(同步拷贝)

Page 58 - NCU 分析:同步拷贝
Page 58 - NCU 分析:同步拷贝

NCU 性能分析显示,同步版本 (STS 指令) 存在大量的长记分板停顿 (Stall Long Scoreboard),这是主要的性能瓶颈。

优化后(异步预取)

Page 59 - NCU 分析:异步预取
Page 59 - NCU 分析:异步预取

使用 LDGSTS 的异步版本显著改善了性能:

  • Stall Barrier 从 23 周期减少到 9 周期。
  • Stall Long Scoreboard 从 25 周期锐减到 2 周期。
    这证实了异步加载有效地隐藏了内存延迟。

LDGSTS 应用案例: 暂存/预取 (Staging/Prefetching)

LDGSTS 非常适合实现软件流水线(或称双缓冲),以重叠计算和内存访问。

Page 60 - LDGSTS 应用于暂存/预取
Page 60 - LDGSTS 应用于暂存/预取

模式说明:
1. 在主循环开始前,预取第一个数据块。
2. 在循环的每次迭代中:
- 开始预取下一个 (i+1) 数据块。
- 等待上一个 (i) 数据块加载完成。
- 使用已加载的数据块 (i) 进行计算。

  1. 循环结束后,处理最后一个数据块。

这种模式通过让数据获取和计算并行执行,有效地隐藏了内存访问延迟,是提高 GPU 核函数性能的常用技术。

为了隐藏数据从全局内存(Global Memory)移动到共享内存(Shared Memory)的延迟,可以采用流水线(pipelining)的方式,在计算当前数据块的同时,异步预取下一个数据块。

Page 61
Page 61

上图展示了一个双缓冲(double-buffering)的实现。代码中定义了一个大小为2的共享内存数组 smem[2]
- 在主循环开始前,预取第一个数据块(i=0)。
- 在循环体内部(i 从 0 到 n-1),首先等待前一个异步拷贝操作完成 (__pipeline_wait_prior(1))。然后,使用已在共享内存中的数据 smem[is] 进行计算。在计算的同时,发起下一个数据块的异步拷贝 (__pipeline_memcpy_async)。

流水线过程如下:
1. 迭代 i=0: 从全局内存获取第0阶段的数据到共享内存,然后进行第0阶段的计算。
2. 迭代 i=1: 在进行第1阶段计算的同时,从全局内存获取第1阶段的数据。
3. ...
4. 迭代 i=n-1: 进行最后阶段的计算,此时没有新的预取操作。

这种方法通过重叠数据传输和计算,有效隐藏了内存访问延迟。

Page 62
Page 62

可以将缓冲区的数量从2个增加到3个(三缓冲),如上图代码所示,smem 数组大小变为3,索引通过模3运算 (is + 1) % 3 进行更新。

Page 63
Page 63

使用更多的流水线阶段(例如,预取多个未来的迭代所需的数据)可以增加“在途字节数”(Bytes in flight),从而更好地利用内存带宽并隐藏更长的延迟。

编写 LDGSTS 的方式

异步数据拷贝功能可以通过多种编程接口实现:

Page 64
Page 64
  1. 原生指令 (Primitives): 使用 <cuda_pipeline.h> 中的底层接口。

    • __pipeline_memcpy_async(...): 对应 LDGSTS 指令。
    • __pipeline_commit(): 对应 LDGDEPBAR 指令。
    • __pipeline_wait_prior(N): 对应 DEPBAR.LE SB,N 指令。
  2. libcu++: 使用 <cuda/pipeline.h> 中提供的更高级、更安全的 C++ 封装。

    • cuda::make_pipeline(): 创建流水线对象。
    • pipeline.producer_acquire() / pipeline.producer_commit(): 生产者侧操作。
    • cuda::memcpy_async(...): 异步拷贝。
    • pipeline.consumer_wait_prior<N>() / pipeline.consumer_release(): 消费者侧操作。
  3. cooperative_groups: 使用 <cooperative_groups/memcpy_async.h> 中的接口,通常与线程块协作相关。

    • cg::memcpy_async(...): 异步拷贝,对应 LDGSTS + LDGDEPBAR
    • cg::wait_prior<N>(block): 等待操作完成,对应 DEPBAR.LE SB,N

LDGSTS 与屏障 (Barrier)

异步拷贝可与屏障结合,实现复杂的生产者/消费者同步模式。

生产者/消费者示例:
在一个循环的每次迭代中:

  1. 生产者线程(Producer threads)将全局内存(GMEM)的一个区域拷贝到共享内存(SMEM)。
  2. 消费者线程(Consumer threads)在拷贝完成后访问共享内存。
Page 65
Page 65

上图展示了迭代0时,生产者将数据从GMEM拷贝到SMEM,消费者从SMEM读取数据的过程。

Page 66
Page 66

上图展示了后续迭代(迭代1)的相同过程。

以下代码展示了如何使用 mbarrier 实现该模式:

Page 68
Page 68

代码解释:
- __shared__ mbarrier_t bar;: 声明一个共享内存屏障。
- mbarrier_init(&bar, N_PRODUCERS);: 初始化屏障,设置需要到达的生产者线程数量。
- 生产者逻辑 (is_producer):
- __pipeline_memcpy_async(...): 发起从全局内存到共享内存的异步拷贝。
- __pipeline_arrive_on(&bar);: 关键操作。此操作将屏障的到达(arrive)与异步拷贝的完成绑定。当硬件完成 memcpy 后,会自动增加屏障的计数值。这避免了在软件中显式等待拷贝完成再进行同步。

  • 消费者逻辑 (else):

    • while (!__mbarrier_try_wait_parity(&bar, phase, 1000)): 消费者线程在此循环等待,直到所有生产者都已完成其数据拷贝(即屏障计数值达到预期)。phase 用于在多次迭代中交替屏障状态。
    • 等待成功后,消费者可以安全地使用已加载到共享内存中的数据。
  • phase ^= 1;: 翻转阶段位,为下一次迭代的同步做准备。

异步数据拷贝:Hopper 架构

Hopper 架构引入了更强大的异步拷贝能力,包括 SM(流多处理器)之间的直接通信。

STAS (Store Asynchronous)

STAS 指令支持从寄存器(REG)到分布式共享内存(DSMEM)的异步拷贝。

Page 69
Page 69

Hopper 异步拷贝特性表 (部分)

硬件 方向 操作 同步 C++ SASS
Hopper REG -> DSMEM dsmem[sind] := var 的异步版本 DSMEM barrier cuda::ptx::st_async STAS

集群 (Cluster) 与分布式共享内存 (Distributed SMEM)
- 一个集群由多个 SM 组成。
- DSMEM 是集群内 SM 共享的内存空间,允许一个 SM 直接读写另一个 SM 的共享内存,实现了高效的 SM 间通信。

Hopper 集群异步拷贝示例 (REG -> DSMEM)

Page 70
Page 70

示例场景:
在一个循环中,集群内的每个线程块(CTA)将其数据拷贝到下一个 CTA 的共享内存中(例如,CTA #i 拷贝到 CTA #i+1 的 SMEM)。最后一个 CTA 将数据拷贝回第一个 CTA,形成一个环形通信模式。

Page 71
Page 71

同步机制:
每个 CTA 上使用两个屏障来实现同步:
- bar: 用于通知消费者 CTA,来自生产者 CTA 的数据已经拷贝完成。
- bar_done: 用于通知生产者 CTA,消费者 CTA 已经准备好接收新数据(即已处理完上一批数据)。

代码实现:

Page 74
Page 74

代码详解:
1. 设置 (__global__ 函数体前半部分):
- __cluster_dims__(8, 1, 1): 定义一个包含8个CTA的集群。
- cuda::barrier: 声明两个屏障 barbar_done
- cluster.sync(): 同步集群内所有线程块。
- 通过 cluster.block_rank() 获取当前CTA的ID (rk),并计算出环形通信中前一个 (rk_prev) 和后一个 (rk_next) CTA的ID。
- cluster.map_shared_rank(...): 获取邻居CTA的共享内存 (smem_next) 和屏障 (bar_next, bar_prev) 的句柄。

  1. 主循环 (for 循环):

    • st_async(&smem_next[...]): 生产者CTA中的线程执行异步存储,将数据写入下一个CTA的共享内存。
    • mbarrier_arrive_expect_tx(...): 生产者在 bar_next 上发出信号,通知消费者它预期会发送 tx_count 数量的事务。这是一个非阻塞操作。
    • while (!barrier_try_wait_parity(bar, ...)): 消费者CTA等待自己的 bar,直到前一个CTA的数据完全到达。
    • 使用数据: r = smem[threadIdx.x],消费者读取数据。
    • mbarrier_arrive(sem_release, ..., bar_prev): 消费者处理完数据后,在 bar_prev 上发出信号,通知其生产者(即前一个CTA),它已经准备好接收下一批数据。
    • while (!barrier_try_wait_parity(bar_done, ...)): 生产者在开始下一次 st_async 前,等待自己的 bar_done,确保其消费者(即下一个CTA)已经准备就绪。

    重要: mbarrier 操作中的 scopespace 参数必须正确匹配。例如,当在一个远程屏障上发信号时,需要使用 scope_cluster,并且 space 参数必须与屏障所在的内存空间(如 space_shared)匹配。

UTMA (Unified Tensor Memory Accelerator)

UTMA 是 Hopper 架构中的一个硬件单元,用于加速在全局内存和共享内存之间的大块张量(Tensor)数据的移动。

Page 75
Page 75

Hopper 异步拷贝特性表 (UTMA 部分)

硬件 方向 操作 同步 C++ SASS
Hopper (TMA) GMEM -> (D)SMEM 1D-5D 块拷贝 (Uniform) (D)SMEM barrier cuda::device::experimental:: cp_async_bulk_tensor_... UTMALDG
Hopper (TMA) SMEM -> GMEM 1D-5D 块拷贝 (Uniform) Thread-local cuda::device::experimental:: cp_async_bulk_tensor_... UTMASTG
  • 操作: UTMA 支持高达5维的张量数据块的统一拷贝。
  • C++ 接口: 通过 cp_async_bulk 系列函数调用。
  • 硬件加速: TMA(Tensor Memory Accelerator)硬件单元负责执行这些拷贝操作,将 SM 从繁重的数据搬运任务中解放出来,使其可以专注于计算。
  • 应用场景: 非常适合需要在大规模全局内存张量和共享内存之间进行高效数据交换的算法,如深度学习中的算子。

TMA (Tensor Memory Accelerator)

TMA nD - 属性

张量拷贝功能 (非详尽列表):
* 单一线程可以发起 1D-5D 的块拷贝。
* 目标模式:
* Tile - 源布局被保留。
* Im2col - 源拷贝框中的元素被重排成列。

  • 越界填充模式:

    • 零 (Zeros)
    • 非数值 (NaNs)
  • 共享内存 (SMEM) Swizzling

约束 (非详尽列表):
* SMEM 地址必须 128B 对齐。
* GMEM 地址必须 16B 对齐。
* 步长 (Strides) 必须是 16B 的倍数。
* 拷贝框 (Copy-box) 的快速维度大小必须是 16B 的倍数。
* 拷贝框的起始地址必须是 16B 对齐。

张量拷贝几何形状由 CUtensorMap 描述符描述:
* 可以是一个 __grid_constant__ 内核参数或 __constant__ 对象 (64B 对齐)。

Page 76
Page 76

TMA nD - 张量描述符

张量描述符是常量,包含以下信息:
* nD (维度, 1, 2, 3, 4, 或 5)
* base_pointer (GMEM 中的基地址)
* tensor_stride[nD - 1] (张量步长)
* tensor_size[nD] (张量大小)
* box_size[nD] (要拷贝的数据框大小)
* element_stride[nD] (元素步长)

下图展示了在全局内存(GMEM)中的一个张量,以及定义其几何形状的参数。左下角显示了将被拷贝到共享内存(SMEM/DSMEM)的数据框。

Page 77: TMA nD 张量描述符示意图
Page 77: TMA nD 张量描述符示意图

TMA nD (GMEM->SMEM) - 张量描述符

从全局内存 (GMEM) 向共享内存 (SMEM) 拷贝数据时,指令参数是变量,包括:
* dst_pointer (SMEM/DSMEM 中的目标地址)
* coords[nD] (拷贝框在张量中的起始坐标)
* 张量描述符
* 同步屏障 (Synchronization barrier)

下图演示了如何根据坐标 coords 从 GMEM 的张量中定位一个数据框,并将其拷贝到 SMEM 的 dst_pointer 位置。

Page 78: 从 GMEM 到 SMEM 的异步数据拷贝
Page 78: 从 GMEM 到 SMEM 的异步数据拷贝

TMA 支持对越界内存的访问进行填充。如下图所示,当拷贝的数据框 (box) 的一部分超出了张量定义的范围时(黄色区域),TMA 会自动使用预设的模式(如零或NaN)填充这部分数据,而不会导致访存错误。

Page 79: TMA 越界读取处理示例
Page 79: TMA 越界读取处理示例

通过设置 element_stride,TMA 能够支持非连续的数据拷贝,例如加载交错的数据。element_stride[nD] 必须为1或2。下图展示了拷贝时跳过部分元素(白色条纹),只读取指定元素(蓝色条纹)的情况。

Page 80: TMA 使用 element_stride 进行非连续数据拷贝
Page 80: TMA 使用 element_stride 进行非连续数据拷贝

TMA nD (SMEM->GMEM) - 张量描述符

从共享内存 (SMEM) 向全局内存 (GMEM) 拷贝数据时,指令参数同样是变量,但源和目标相反:
* src_pointer (SMEM/DSMEM 中的源地址)
* coords[nD] (拷贝框在目标张量中的起始坐标)
* 张量描述符
* 同步屏障 -> 线程本地同步 (thread-local sync)

下图演示了从 SMEM 的 src_pointer 位置读取一个数据框,并根据坐标 coords 将其写入 GMEM 的张量中。

Page 81: 从 SMEM 到 GMEM 的异步数据拷贝
Page 81: 从 SMEM 到 GMEM 的异步数据拷贝

与 GMEM 到 SMEM 的拷贝不同,从 SMEM 到 GMEM 的拷贝有一个重要约束:坐标 coords 必须大于等于0。不允许将数据写入张量范围之外的负坐标区域。

Page 82: SMEM 到 GMEM 拷贝的坐标约束
Page 82: SMEM 到 GMEM 拷贝的坐标约束

TMA nD 与 LDGSTS 对比

以一个 X 型 16 阶有限差分 (FD) 模板为例进行说明:
* 线程块大小: 32x8
* 共享内存大小: 48x24 (32x8 中心区域 + 每边 8 点的光环区域)

LDGSTS (伪代码):
使用 LDGSTS 时,每个线程需要根据自己的位置计算偏移量,并执行多个条件加载指令来获取所需的数据。这会导致:
* MIO (Memory Input/Output) 压力
* 寄存器压力
* 谓词压力

TMA (伪代码):
使用 TMA 时,只需由单个线程 (tid == 0) 发起一个无条件的 async_load_tma 操作,即可加载整个线程块所需的数据(包括光环区域)。这种方式代码更简洁,效率更高。

对于更高阶的模板,LDGSTS 的问题会更严重。例如,12点光环区域(如图中的#3和#5区域)无法被一个线程块完全覆盖,需要每个区域执行两次加载操作。

Page 84: TMA 与 LDGSTS 在模板计算中的对比
Page 84: TMA 与 LDGSTS 在模板计算中的对比

TMA 细节

TMA 编程模型是 Uniform 的:
一个 warp 中所有活跃的线程都以“相同的参数”执行 TMA 操作。
高级代码 async_load_tma(args_per_thread); 会被编译器展开为在一个循环中为每个活跃线程发布指令。

错误与非优化的代码:
* 错误代码 (竞争条件): async_load_tma(single_thread_args); 如果所有线程都用相同的参数(例如,为整个块设计的参数)调用,会导致竞争。
* 非优化代码: if (tid == 0) async_load_tma(single_thread_args); 尽管代码意图是让单个线程执行,但编译器无法确定 tid == 0 对于整个 warp 只有一个线程为真,因此仍然会生成一个检查每个线程的剥离循环 (peeling loop)。

更优代码:
应该使用 cooperative group 的 invoke_one 来确保只有一个线程执行 TMA 操作,这是最高效且正确的方式。

if (tid == 0)
    cg::invoke_one(coalesced_threads(), 
                   async_load_tma(single_thread_args));
Page 87: TMA 编程模型的正确与错误示例
Page 87: TMA 编程模型的正确与错误示例

TMA nD 主要步骤总结

GMEM -> SMEM 拷贝流程:

  1. INIT (初始化):

    • 由单个线程初始化 mbarrier 对象。
  2. FIRE (触发):

    • 同步线程。
    • 由单个线程使用 invoke_one 调用 cp_async_bulk_tensor_2d_global_to_shared 来发起异步拷贝。
    • 通过 mbarrier_arrive_expect_tx 到达屏障并预期事务。
  3. WAIT FOR COMPLETION (等待完成):

    • 在循环中调用 mbarrier_try_wait_parity 来检查 mbarrier,直到拷贝完成。
Page 88: TMA nD 从 GMEM 到 SMEM 的主要步骤总结
Page 88: TMA nD 从 GMEM 到 SMEM 的主要步骤总结

SMEM -> GMEM 拷贝流程:

  1. FIRE (触发):

    • 在所有线程使用完 SMEM 后,进行同步。
    • 由单个线程使用 invoke_one 发起多个 cp_async_bulk_tensor_2d_shared_to_global 操作。
  2. COMMIT GROUP (提交组):

    • 调用 cp_async_bulk_commit_group() 提交一批异步操作。
  3. WAIT FOR COMPLETION (等待完成):

    • 调用 cp_async_bulk_wait_group<N>() 等待除了最后 N 组之外的所有组完成,以实现计算与通信的重叠。
Page 89: TMA nD 从 SMEM 到 GMEM 的主要步骤总结
Page 89: TMA nD 从 SMEM 到 GMEM 的主要步骤总结

SMEM 存储体冲突

对于一个给定的二维数组 array[NX][y][x],传统的线性索引计算方式为:ind = y * NX + x

NX 是 SMEM 存储体数量(例如32)的倍数时,一个 warp 中的所有线程(通常是32个)访问同一行不同列的元素时,它们的内存地址 y * 32 + x 经过模32运算后会映射到同一个存储体。这会导致严重的32路存储体冲突,使得并行访问退化为串行访问,极大降低性能。

下图标示了当 NX=32 时,同一行(例如 y=0)的连续32个元素都映射到索引为0到31的 SMEM 存储体。当一个 warp 访问这一行时,所有32个线程都试图访问不同的存储体,因此无冲突。但如果 warp 访问同一列的不同行(图中垂直方向),所有访问都会命中同一个存储体,导致冲突。

Page 90: SMEM 存储体冲突示例
Page 90: SMEM 存储体冲突示例
SMEM 存储体冲突 - 填充 (Padding)

对于给定的二维数组 array[NX][NX]index [y][x] 通过以下方式映射到线性索引 ind

  • 传统方式ind := y * NX + x

NX 是存储体数量(例如32)的倍数时,沿 y 轴(列)的连续访问会命中同一个 SMEM 存储体,导致存储体冲突。如下图所示,当 NX=32 且数据类型大小为4字节时,ind % 32 的计算结果显示,访问同一列的所有元素(x相同,y不同)都会请求同一个存储体(图中黄色高亮部分),造成序列化访问。

为了解决这个问题,可以采用填充(Padding)的方法,将数组的维度从 NX 增加到 NX+1

  • 填充方式ind := y * (NX+1) + x

通过填充,同一列的连续访问会被分散到不同的存储体中,从而避免冲突。下图标示了使用 NX+1 填充后,列访问不再产生冲突(图中绿色高亮部分)。

Page 91 - SMEM 存储体冲突 - 填充
Page 91 - SMEM 存储体冲突 - 填充
SMEM 存储体冲突 - Swizzling

Swizzling 是另一种避免 SMEM 存储体冲突的技术。

  • 传统方式ind := y * NX + x
  • Swizzling 方式ind := y * NX + (y ^ x)

Swizzling 通过对列索引 x 和行索引 y 进行异或操作来改变内存访问模式,从而将原本会导致冲突的列访问分散到不同的存储体中。如下图所示,经过 Swizzling 处理后,列访问不再有存储体冲突。

Page 92 - SMEM 存储体冲突 - Swizzling
Page 92 - SMEM 存储体冲突 - Swizzling

TMA nD-SMEM Swizzling

Tensor Memory Accelerator (TMA) 在128字节的段(segment)内对8个16字节的块(chunk)进行 swizzle 操作。

约束条件:
* NX * sizeof(T) == SWIZZLE_SIZE,其中 T 是数组 array[][NX] 的类型。
* 允许的 SWIZZLE_SIZE 值为 32、64、128。

计算 Swizzled 索引的步骤:
给定 T 类型数组 array[][NX] 中的 [y][x] 索引:
1. 计算16字节块在128字节段内的索引:
i16 := (y * NX + x) * sizeof(T) / 16
y16 := i16 / 8
x16 := i16 % 8
2. 计算 swizzled 后的16字节块索引:
x16_swz := y16 ^ x16
3. 计算最终的 swizzled 索引:
x_swz := x16_swz * 16 / sizeof(T) % NX + x % (16 / sizeof(T))

下表展示了不同 SWIZZLE_SIZE 下16字节块的 swizzled 索引。

Page 93 - TMA nD-SMEM Swizzling 机制
Page 93 - TMA nD-SMEM Swizzling 机制

分段排序性能

在一个分段排序(Segmented sort)任务中,对不同的异步数据复制方法进行了性能评估。

  • 任务: 对 4194304 个长度为 128 的整数序列(约 2GB)进行排序。
  • 算法: Pair-wise sorting network 算法。
  • 平台: H200 实验。

下表对比了四种 SMEM 复制方法的性能:

Page 94 - 分段排序性能对比表
Page 94 - 分段排序性能对比表

结论: TMA + swizzle 表现最佳,其耗时最短(1.07 ms),带宽最高(4066 GB/s),达到了理论峰值(SOL)的 87%。

异步复制 - UBLKCP

Hopper 架构引入了 TMA(Tensor Memory Accelerator)来支持高效的异步数据复制,其中包括 UBLKCP (Uniform Block Copy) 指令。

下表总结了不同的异步复制操作,其中高亮部分为 Hopper TMA 的能力:

Page 95 - 异步复制操作概览及UBLKCP
Page 95 - 异步复制操作概览及UBLKCP
  • Hopper TMA: 支持全局内存(GMEM)与共享内存(SMEM)之间,以及SMEM与SMEM之间的1D块复制(Uniform)。
  • C++ API: cuda::device::experimental::cp_async_bulk_global_to_shared, cp_async_bulk_shared_to_global
  • PTX 指令: cp.async_bulk.*
  • 可用性: 最新的 GitHub CCCL 库已提供支持。

下图展示了 Hopper 架构中 TMA 在全局内存和分布式共享内存(Cluster 内的 Shared Memory)之间传输数据的路径。

Page 95 - Hopper TMA 数据路径示意图
Page 95 - Hopper TMA 数据路径示意图

TMA - 属性

UBLKCPmemcpy 的异步版本,其主要属性如下:

  • 由单个线程发起。
  • 使用源/目标指针和复制大小作为参数(不使用 Tensor descriptor,因此不支持 bound-box、stride、fill、swizzle 等高级功能)。
  • 对齐要求:

    • GMEM/SMEM 地址需 16B 对齐。
    • 复制大小需是 16B 的倍数。
  • 同步机制: 与 TMA nD 使用相同的同步机制(barrierthread-local)。

    • mbarrier_arrive_expect_tx(sem_release, scope_cta, space_shared, &bar, SIZE);
    • cp_async_bulk_commit_group();
    • cp_async_bulk_wait_group_read<N>();

可压缩内存 (Compressible Memory)

Page 97 - 可压缩内存章节封面
Page 97 - 可压缩内存章节封面

简介

可压缩内存是指对全局内存(GMEM)进行硬件压缩。

  • 数据在发送到 GMEM 之前被压缩。
  • 数据从 GMEM 读取后被解压缩。

下图展示了数据在 SM 和 GMEM 之间流经 L2 缓存时进行压缩和解压缩的过程。

Page 98 - 可压缩内存工作原理
Page 98 - 可压缩内存工作原理

快速事实

  • 透明性: 对用户完全透明(只需更改 GMEM 分配调用)。
  • 效益: "压缩带宽"(GMEM <-> L2 之间),而不是压缩存储空间(GMEM 或 L2 本身)。
  • 粒度: 128B(缓存行粒度)。
  • 压缩: 无损压缩(压缩比为 2x、4x,如果无法压缩则为 1x)。压缩是自动的,并非所有数据都能被压缩。
  • 可用性: 在 Ada 和 Hopper 架构上可用。
  • 潜在风险: 可能会降低性能,因为压缩缓存未命中(compression cache misses)会导致巨大的性能惩罚。

Saxpy 示例 - Kernels

以下通过三种不同的 SAXPY 内核实现来展示访存局部性对可压缩内存性能的影响。

  1. "No-loop" kernel: 每个线程处理一个元素,访存具有良好的局部性。
    Page 100 - "No-loop" kernel 代码及访存模式

  2. "Grid-stride loop" kernel: 每个线程以网格大小为步长处理多个元素,导致访存分散,局部性差。随着计算波次(waves)的增加,访存局部性会恶化。
    Page 101 - "Grid-stride loop" kernel 代码及访存模式

  3. "Block-stride loop" kernel: 每个线程块处理一块连续的数据,局部性较好。随着计算波次的增加,访存局部性得到改善。
    Page 102 - "Block-stride loop" kernel 代码及访存模式

Saxpy 示例 - 性能

实验使用高度可压缩的数据 x[i] = y[i] = {i, i, i, i} 来测试性能。

下表展示了在 H100 平台上,针对 1.3GB 和 13GB 数组,不同内核在启用和禁用内存压缩时的性能对比。

Page 103 - Saxpy 性能对比表
Page 103 - Saxpy 性能对比表

结论: 局部性至关重要!
* 对于访存局部性好的内核("No-loop", "Block-loop"),可压缩内存显著提升了有效带宽(例如,No-loop 带宽从 3.71 TB/s 提升到 5.77 TB/s)。
* 对于访存局部性差的内核("Grid-loop"),启用压缩反而导致性能严重下降(例如,带宽从 3.53 TB/s 降至 0.86 TB/s)。

NCU 分析

使用 NVIDIA Nsight Compute (NCU) 可以观察内存压缩的效果。

下图是 NCU 的内存图表分析,显示了:
* L2 压缩(L2 Compression)比率达到了 4.00。
* 从 L2 缓存到系统内存的流量,虽然原始数据大小为 1.64 GB,但实际传输的数据量仅为 322.65 MB。这与 1.3GB 的数组被 4 倍压缩的结果(1.3GB / 4 = 322.65 MB)相符。

Page 104 - NCU 内存压缩分析
Page 104 - NCU 内存压缩分析

地震波示例

可压缩内存在实际应用中也很有用,例如逆时偏移(RTM)中的三维有限差分波传播计算。

  • 算法: 内核在一个三维网格上执行,z 轴是迭代最慢的轴。
  • 计算模式: 每个线程块(TB)处理一个 xy 平面上的瓦片(tile),并在 z 方向上循环。x 和 y 方向的偏导数通过 SMEM 计算,而 z 方向的偏导数通过寄存器队列(REG queue)计算。
  • 访存特性: 这种沿 z 轴的模板计算具有良好的空间局部性,适合利用可压缩内存来提升性能。
Page 105 - RTM 地震波传播计算示例
Page 105 - RTM 地震波传播计算示例

地震波示例 - 改善局部性

该示例为RTM(逆时偏移)中的3D有限差分波传播问题。

代码结构如下:
- 定义块大小 BZ = 64。
- 定义3D的块(block)和网格(grid)维度。
- 内核函数中的z循环是沿最慢的轴进行的。

Page 106, 3D数据块和z轴循环示意图
算法详情请参见 Paulius Micikevicius 的 "3D Finite Difference Computation on GPUs using CUDA"。

为了改善局部性,可压缩内存被引入。在传统的内存布局中,不同波前切片(wave slices)在z轴方向上相距较远,导致局部性差。通过数据压缩,可以将这些切片在内存中更紧凑地排列,从而提高局部性。局部性可以近似地表示为:局部性 ≈ (波前切片间的z轴距离) * z轴步长

Page 107, 改善局部性前后内存布局对比
Page 107, 改善局部性前后内存布局对比

地震波示例 - L40s 性能

问题描述:
- 8阶TTI RTM(倾斜横向各向同性介质的逆时偏移),单遍(single pass)内核。
- 计算域大小:800x800x800(9个卷 @ 2.2 GB / 卷)。
- 线程块大小:32x16。

性能测试结果 (L40s MCells/sec):

下表展示了在不同z-blocking(z轴分块)大小下,使用和不使用内存压缩的性能对比。结果显示,启用压缩后,尤其是在较小的z-blocking尺寸下,性能得到显著提升。例如,在32 z-blocking时,性能从未压缩的1594 MCells/sec提升至51497 MCells/sec。

Page 108, L40s性能数据表
Page 108, L40s性能数据表

性能会受到波场可压缩性的影响。波场从零开始传播,随着时间的推移,波场变得越来越复杂,可压缩性随之降低,从而导致性能下降。下图展示了不同时间步长下的性能变化,以及波场随时间演变的形态。

Page 109, 不同时间步长下的性能数据及波场演变
Page 109, 不同时间步长下的性能数据及波场演变

分配可压缩内存

通过虚拟内存管理API(CUDA驱动)可以分配可压缩内存。主要步骤如下:
1. 设置分配属性: 在 CUmemAllocationProp 中设置 compressionTypeCU_MEM_ALLOCATION_COMP_GENERIC
2. 对齐分配大小: 将分配大小向上舍入到硬件支持的粒度。
3. 保留虚拟地址范围: 使用 cuMemAddressReserve
4. 分配物理内存: 使用 cuMemCreate
5. 映射内存: 使用 cuMemMap 将物理内存映射到虚拟地址范围。
6. 设置访问权限: 使用 cuMemSetAccess 使内存可访问。

以下是相关API调用的示例代码:

Page 110, 分配可压缩内存的示例代码
Page 110, 分配可压缩内存的示例代码

CUDA Graphs

Page 111, CUDA Graphs 标题页
Page 111, CUDA Graphs 标题页

减少启动开销

CUDA Graphs通过将一系列CUDA操作(如内核启动、内存拷贝)捕获到一个图中,然后一次性启动整个图来减少内核的启动开销。这避免了CPU为每个内核单独发起的启动延迟,使得GPU可以背靠背地执行内核,提高了效率,尤其适用于包含许多小型内核的工作负载。

Page 112, CUDA Graphs 与传统内核启动方式的对比
Page 112, CUDA Graphs 与传统内核启动方式的对比

简单示例 - 捕获

一个典型的多流(multi-stream)工作流通常使用事件(events)来同步不同流之间的操作。

Page 113, 使用多流和事件的CUDA代码及其依赖关系图
Page 113, 使用多流和事件的CUDA代码及其依赖关系图

这个工作流可以通过 CUDA Graphs 进行捕获。将上述代码序列放置在 cudaStreamBeginCapturecudaStreamEndCapture 调用之间,CUDA运行时会进行一次“空运行”(dry-run),记录下所有的操作及其依赖关系,并将其构建成一个图(graph)。

Page 114, 将多流操作捕获为CUDA Graph
Page 114, 将多流操作捕获为CUDA Graph

使用步骤

使用CUDA Graphs通常分为三个步骤:
1. 捕获图 (Capture graph):类似于编写代码。使用 cudaStreamBeginCapturecudaStreamEndCapture 将一系列CUDA操作捕获到一个 cudaGraph_t 对象中。
2. 创建可执行图 (Create executable graph):类似于编译代码。使用 cudaGraphInstantiate 将捕获的图实例化为一个 cudaGraphExec_t 对象。此步骤会进行优化和验证。
3. 启动可执行图 (Launch executable graph):类似于执行编译后的代码。使用 cudaGraphLaunch 将实例化的图提交到一个流中执行。

图(Graph)封装的内容:
- 流依赖关系和并发性:
- 即使通过 cudaGraphLaunch 将图启动到单个流中,图中的分支(branches)也可以并发执行。
- 节点可以继承流的优先级(需要设置 cudaGraphInstantiateFlagUseNodePriority 标志)。
- 用于捕获的流在捕获后可以被销毁。

  • 内核参数:
    • 如何在运行时更改内核参数是一个需要考虑的问题。
Page 115, CUDA Graphs 使用步骤
Page 115, CUDA Graphs 使用步骤

图的修改

对于常见的时间步进模式(time-stepping pattern),即在一个循环中反复执行相同的内核序列,有几种方法可以处理图的更新:

  1. 捕获奇偶迭代图: 为奇数和偶数次迭代分别捕获图,并交替启动。
  2. 始终捕获并更新: 使用 cudaGraphExecUpdate 来更新图。
  3. 显式图构建API: 使用 cudaStreamGetCaptureInfo, cudaGraphAddKernelNodecudaStreamUpdateCaptureDependencies 等API在每次迭代中手动更新图的节点。

通常,拓扑结构的改变需要重新实例化图。但一些微小的拓扑变化,例如启用/禁用某个节点(使用 cudaGraphNodeSetEnabled),可能不需要重新实例化。

Page 116, CUDA Graphs 修改方法
Page 116, CUDA Graphs 修改方法

其他捕获限制

以下是一些关于CUDA Graph捕获的限制(非详尽列表):
- 默认旧式流 (Default legacy stream) 无法被捕获。应将库调用从 libraryCall(cudaStreamDefault) 更改为 libraryCall(stream)
- 同步调用 (Synchronous calls) 无法被捕获。例如,应使用 cudaMemcpyAsync 替代 cudaMemcpy
- 不能同步: 无法捕获 cudaStreamSynchronizecudaDeviceSynchronize 等同步操作。
- 主机逻辑 (Host logic): 无法直接捕获主机端逻辑。应使用 cudaLaunchHostFunc 将其包装成一个可以在流中执行的节点。

总结来说,任何可以放入流(stream)中的操作都可以被捕获

Page 117, CUDA Graphs 捕获限制
Page 117, CUDA Graphs 捕获限制

示例 - 算法

以下是一个基于地震道处理的工作流,该工作流由多个步骤组成,非常适合使用CUDA Graphs进行优化:
1. 正向 cuFFT
2. 将频率延拓到多个道(traces)
3. 对多个道进行逆向 cuFFT(由于大小可变,无法批处理)
4. 归约(Reductions)操作(使用 CUB 库),并组合奇偶部分
5. 最终合并

这个工作流可以被捕获为一个复杂的CUDA Graph,其中包含了并行和串行的依赖关系,并通过迭代启动来处理数据。

Page 118, 地震道处理算法的CUDA Graph示意图
Page 118, 地震道处理算法的CUDA Graph示意图

示例 - 代码

下图展示了如何将上述地震道处理算法捕获为CUDA Graph。代码中包含了对 cufftcub 库的调用。这些库调用可以被无缝地捕获到图中,成为子图(sub-graphs),无需额外的开发工作。

Page 119, 地震道处理算法的CUDA Graph捕获代码
Page 119, 地震道处理算法的CUDA Graph捕获代码

示例 - 性能

性能测试背景:
- 输入道:1024个采样点
- 输出道:100个,范围在128-8192之间

结果:
在H200上的测试表明,使用CUDA Graphs带来了5倍的加速

下图直观地对比了使用和不使用CUDA Graphs时的GPU执行情况:
- 使用Graphs: 内核执行紧凑且高效。
- 不使用Graphs: 由于CPU启动开销,内核执行被“串行化”,即使使用多个流也几乎没有效果。

Page 120, 使用与不使用CUDA Graphs的性能对比图
Page 120, 使用与不使用CUDA Graphs的性能对比图

Nsight systems

  • 要在图内部进行跟踪,请使用 nsys--cuda-graph-trace node 选项。
  • 注意:这比默认设置的性能分析开销更高。
  • 默认选项是 --cuda-graph-trace graph
Page 121: Nsight Systems 中 CUDA Graphs 的跟踪视图,展示了使用 node 选项可以观察到图内部的详细执行情况。
Page 121: Nsight Systems 中 CUDA Graphs 的跟踪视图,展示了使用 node 选项可以观察到图内部的详细执行情况。

进一步阅读

CUDA graphs API 的内容非常丰富,我们刚才看到的只是冰山一角。

请参阅本次 GTC 的相关演讲:
- Jiri Kraus 的 "Multi GPU Programming Models for HPC and AI" [S61339]
- 以及配套示例:https://github.com/NVIDIA/multi-gpu-programming-models/tree/master/nccl_graphs

  • Jiqun Tu 和 Ellery Russell 的 "Accelerating Drug Discovery: Optimizing Dynamic GPU Workflows with CUDA Graphs, Mapped Memory, C++ Coroutines, and More" [S61156]

相关资源:
- GTS'23 [S51211]: "CUDA Graphs 101"
- GTC'21 [S32082]: "Effortless CUDA Graphs"
- 技术博客: "Constructing CUDA Graphs with Dynamic Parameters"
- 技术博客: "Getting Started with CUDA Graphs"
- CUDA 示例:
- simpleCudaGraphs
- jacobiCudaGraphs
- graphMemoryNodes
- graphMemoryFootprint

总结

总结与要点

  • 内存模型:

    • 每个 CUDA 程序员都应该了解的内存知识...
  • 异步内存拷贝与屏障 (barriers):

    • 增加在途字节数(bytes in flight)+ 减少压力(寄存器、谓词、MIO 等)。
    • 需要一些编程工作。
  • 可压缩内存:

    • 提高有效全局内存(GMEM)带宽(如果数据是可压缩的)。
    • 可以是一种非常容易实现的优化:只需替换 cudaMalloc
    • 危险: 对访问模式敏感(可能导致性能下降);可能需要修改核函数(kernel)才能有效。
  • CUDA graphs:

    • 适用于执行时间短的核函数。
    • 可以是一种容易实现的优化:流捕获(stream capturing)只需几行代码。

核心性能优化技术

GTC'24 的核心性能优化技术演讲列表

  • CUDA 编程与性能优化入门 [S62191]
  • CUDA 高级性能优化 [S62192]
  • Grace CPU 超级芯片的性能优化 [S62275]
  • Grace Hopper 超级芯片架构与深度学习应用性能优化 [S61159]
  • 用于 HPC 和 AI 的多 GPU 编程模型 [S61339]
  • 更多数据,更快速度:Python 和 C++ 中的 GPU 内存管理最佳实践 [S62550]
  • 利用 Grace Hopper 的能力加速向量数据库搜索 [S62339]
  • 从零到极致:通过逐步优化将服务吞吐量提升数十倍 [S62410]