How to Streamline Shared Memory Space With the NVSHMEM Communication Library

Akhil Langer, Seth Howell, Jim Dinan, Nvidia | GTC Spring 2023

目录

• 1. 通信模型

  • 1.1 CPU发起的通信
  • 1.2 GPU发起的通信

• 2. NVSHMEM 概述

  • 2.1 NVSHMEM 简介
  • 2.2 NVSHMEM 对称内存模型
  • 2.3 NVLink 通信优化
  • 2.4 纵向和横向扩展带宽
  • 2.5 线程级通信
  • 2.6 线程组通信
  • 2.7 核内同步
  • 2.8 流序操作
  • 2.9 集体核函数启动
  • 2.10 与CUDA图的互操作性
  • 2.11 与MPI/OpenSHMEM的互操作性

• 3. NVSHMEM 最新特性

  • 3.1 发布 NVSHMEM 2.9.0
  • 3.2 GPUDirect Async 内核发起通信
  • 3.3 GPU-发起 (IBGDA) 与代理通信性能
  • 3.4 标量 Put 合并
  • 3.5 二进制分发改进
  • 3.6 CMAKE 构建系统
  • 3.7 内联设备代码
  • 3.8 集合优化
    • 3.8.1 使用 LL 协议改进延迟
    • 3.8.2 拓扑感知优化

• 4. 性能案例研究

  • 4.1 NVSHMEM + cuFFT = cuFFTMp
  • 4.2 NVSHMEM 在 GROMACS 中的应用
  • 4.3 NVSHMEM 在高性能 Linpack (HPL) 中的应用

• 5. 即将推出的特性和结论

  • 5.1 NVSHMEM 中的 Grace Hopper 支持
  • 5.2 第四代 NVLink
  • 5.3 NVSHMEM 内存空间与上下文
  • 5.4 总结
  • 5.5 NVSHMEM 状态

---

1. 通信模型

1.1 CPU发起的通信

特点：

• 在GPU上进行计算。
• 由CPU发起通信。
• 在边界处进行同步。

普遍使用的模型，但存在问题：
* 在关键路径中存在卸载延迟（Offload latencies in critical path）。
* 通信未重叠（Communication is not overlapped）。
* 隐藏这些延迟会增加代码复杂性。
* 不隐藏延迟会限制强扩展性。

1.2 GPU发起的通信

特点：
* 在GPU上进行计算。
* 由GPU发起通信。

优点：
* 消除了卸载延迟（Eliminates offload latencies）。
* 计算和通信重叠（Compute and communication overlap）。
* 通过线程隐藏延迟（Latencies hidden by threading）。
* 更易于用内联通信表达算法（Easier to express algorithms with inline communication）。

目标： 在提升性能的同时，简化编程。

---

2. NVSHMEM 概述

2.1 NVSHMEM 简介

OpenSHMEM，为NVIDIA GPU集群上的最佳性能而适配

• 功能： 将集群中多个GPU的内存聚合到一个分布式全局地址空间中。

  • 通过put、get、原子API进行数据访问。
  • 提供集合通信API。

• 通信与CUDA执行模型集成：

  1. GPU核函数发起的（GPU kernel-initiated）操作。
  2. CUDA流/图上的操作（Operations on CUDA streams/graphs）。
  3. CPU发起的（CPU initiated）操作。

• 可与CPU OpenSHMEM或MPI库结合使用，用于主机内存通信。

2.2 NVSHMEM 对称内存模型

独立分区聚合到全局地址空间

• 对称对象在每个PE上以相同大小集体分配。

  • 对称内存： nvshmem_malloc(...)
  • 私有内存： cudaMalloc(...)

• 读： nvshmem_get(...)

• 写：： nvshmem_put(...)
• 原子操作： nvshmem_atomic_add(...)
• 刷新写操作： nvshmem_quiet()
• 顺序写操作： nvshmem_fence()
• 同步： nvshmem_barrier()
• 轮询： nvshmem_wait_until(...)

2.3 NVLink 通信优化

使用第四代NVLink无缝扩展至256个GPU

• NVSHMEM无缝扩展：

  • 使用NVLink和PCIe进行扩展（Scales-up）。
  • 使用InfiniBand、RoCE等进行扩展（Scales-out）。

• 内部使用CUDA IPC和cuMem API将对等PE的对称内存映射到虚拟地址空间。

  • nvshmem_put/get on device → load/store
  • nvshmem_put/get_on_stream → cudaMemcpyAsync
  • nvshmem_ptr direct pointer bypass → kernel direct load/store

2.4 纵向和横向扩展带宽

理论对等体双向带宽 (GB/s)

• PCI Gen4 x16：64
• PCI Gen5 x16：128
• 第二代NVLink (DGX V100)：2400
• 第三代NVLink (DGX A100)：2400
• 第四代NVLink (DGX H100)：3600
• 第四代NVLink (32 DGX H100)：57600

理论网络注入 + 弹出带宽每GPU (GB/s)
* 100Gb ConnectX-5 (DGX V100)：25
* 200Gb ConnectX-6 (DGX A100)：50
* 400Gb ConnectX-7 (DGX H100)：100

2.5 线程级通信

• 允许细粒度通信和重叠。
• 在DGX系统上高效映射到NVLink网络。

_global_ void stencil_single_step(float *u, float *v, ...) {
    int ix = get_ix(blockIdx, blockDim, threadIdx);
    int iy = get_iy(blockIdx, blockDim, threadIdx);

    compute(u, v, ix, iy);
    // Thread-level data communication API
    if (iy == 1)
        nvshmem_float_p(u+(ny*1)*nx+ix, u[nx+ix], top_pe);
    if (iy == ny)
        nvshmem_float_p(u+ix, u[ny*nx+ix], bottom_pe);
}

for (int iter = 0; iter < N; iter++) {
    swap(u, v);
    stencil_single_step<<<..., stream>>>(u, v, ...);
    nvshmem_barrier_all_on_stream(stream);
}

2.6 线程组通信

• NVSHMEM操作可以由块/warp中的所有线程发出。
• 通过IB等网络进行更高效的数据传输。
• 仍然允许warp间/块间重叠。

_global_ void stencil_single_step(float *u, float *v, ...) {
    int ix = get_ix(blockIdx, blockDim, threadIdx);
    int iy = get_iy(blockIdx, blockDim, threadIdx);

    compute(u, v, ix, iy);
    // Thread block-level communication API
    int boffset = get_block_offet(blockIdx, blockDim);
    if (blockIdx.y == 0)
        nvshmemx_float_put_nbi_block(u+(ny*1)*nx+boffset, u+nx+boffset, blockDim.x, top_pe);
    if (blockIdx.y == (blockDim.y-1))
        nvshmemx_float_put_nbi_block(u+boffset, u+ny*nx+boffset, blockDim.x, bottom_pe);
}

for (int iter = 0; iter < N; iter++) {
    swap(u, v);
    stencil_single_step<<<..., stream>>>(u, v, ...);
    nvshmem_barrier_all_on_stream(stream);
}

2.7 核内同步

• 在核函数内跨PE进行集体或点对点同步。
• 将应用程序的更大一部分卸载到同一个CUDA核函数中。

_global_ void stencil_multi_step(float *u, float *v, int N, int *sync, ...) {
    int ix = get_ix(blockIdx, blockDim, threadIdx);
    int iy = get_iy(blockIdx, blockDim, threadIdx);

    for (int iter = 0; iter < N; iter++) {
        swap(u, v);
        // Thread block-level data exchange (assume even/odd iter buffering)
        int boffset = get_block_offet(blockIdx, blockDim);
        if (blockIdx.y == 0)
            nvshmemx_float_put_nbi_block(u+(ny*1)*nx+boffset, u+nx+boffset, blockDim.x, top_pe);
        if (!((blockIdx.y == (blockDim.y-1))))
            nvshmemx_float_put_nbi_block(u+boffset, u+ny*nx+boffset, blockDim.x, bottom_pe);

        this_grid.sync();
        if ((!itid) nvshmem_barrier(); // Be aware of synchronization costs. Best synchronization approach is application dependent!
        this_grid.sync();
    }
}

更多详情：https://github.com/NVIDIA/multi-gpu-programming-models

2.8 流序操作

NVSHMEM CPU发起的在CUDA流和图上入队的操作

• 并非总是将所有通信或同步移入CUDA核函数是最佳选择。
• Inter-CTA同步（例如grid.sync()）的延迟可能比核函数启动延迟更长。
• 允许混合细粒度通信和粗粒度同步。

2.9 集体核函数启动

确保使用设备侧核间同步时的进度

NVSHMEM 用法	CUDA 核函数启动
设备发起的通信	执行配置语法 <<<...,>>> 或启动API
设备发起的同步	nvshmemx_collective_launch

• CUDA的吞吐量计算模型允许（并鼓励）比GPU能容纳的更大的网格。
• 核间同步要求生产者和消费者线程并发执行。
• 集体启动通过CUDA协作启动和1PE/GPU的要求来保证共驻。

2.10 与CUDA图的互操作性

• CUDA Graphs (CUDA图)：

  • 在CUDA 10.0中引入。
  • 减少启动开销。
  • 工作流优化。

• NVSHMEM可与Graphs组合使用：

  • 流上操作作为节点。
  • CUDA核函数可以使用NVSHMEM。

2.11 与MPI/OpenSHMEM的互操作性

通过基于属性的初始化例程启用

MPI：
使用MPI_COMM_WORLD初始化NVSHMEM。

MPI_Init(&argc, &argv);

MPI_Comm mpi_comm = MPI_COMM_WORLD;
nvshmemx_init_attr_t attr;
attr.mpi_comm = &mpi_comm;
nvshmemx_init_attr(NVSHMEMX_INIT_WITH_MPI_COMM, &attr);

mype_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE);
CUDA_CHECK(cudaSetDevice(mype_node));

OpenSHMEM：
使用SHMEM默认上下文初始化NVSHMEM。

shmem_init();

nvshmemx_init_attr_t attr;
nvshmemx_init_attr(NVSHMEMX_INIT_WITH_SHMEM, &attr);

mype_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE);
CUDA_CHECK(cudaSetDevice(mype_node));

---

3. NVSHMEM 最新特性

3.1 发布 NVSHMEM 2.9.0

自 NVSHMEM 2.5.0 / GTC Spring 2022 以来新增的特性：

• InfiniBand GPUDirect Async (IBGDA) 传输

• HPC SDK 集成

  • 库版本控制 (Library versioning)
  • CUDA 次要版本兼容性支持 (CUDA minor version compatibility support)
  • 模块化传输、引导支持 (Modular transport, bootstrap support)

• CMake 构建系统 (CMake build system)

• 集合优化 (Collective Optimizations)
• 内联设备代码 (Inlined Device Code)

• DMABUF 支持

  • 与上游 ibverbs 兼容 (Compatibility with upstream ibverbs)
  • 无需 nv_peer_mem 模块 (nv_peer_mem module not needed)

• Slingshot-11 支持

  • Beta 版支持 (Beta support)
  • 稳定性与性能的持续改进工作 (Ongoing work on stability and performance)

3.2 GPUDirect Async 内核发起通信

与 Mellanox NICs 更紧密的集成

• CPU 代理线程不再用于内核发起的同步操作。

• 内核可以直接执行：

  • 1. 写入内存数据
  • 2-4. 提交工作请求至 NIC
  • 5-8. 执行通信
  • 1. 检查操作完成

• 并行度显著提升，从而提高了吞吐量。

• 与 Mellanox Stack 集成：
  • 支持 NVIDIA Mellanox 网络
  • 直接在 Mellanox 软件栈上运行
  • 通过利用原生原子操作避免了对 GDRCopy 的需求

3.3 GPU-发起 (IBGDA) 与代理通信性能

在 2kB 和 16kB 传输大小下的全带宽

• IBGDA (左图) 性能随着线程块 (CTAs) 的数量扩展；代理 (右图) 在 8 CTAs 时达到扩展极限。
• 使用 128 CTAs 时，IBGDA 块 Put 吞吐量提高了 5.75 倍；使用 1 CTA 时，IBRC 块 Put 吞吐量提高了 1.8 倍。
• 权衡：当 CTAs 并行通信时，IBGDA 提供高吞吐量；代理具有较低的操作切换开销。
• 测试节点：Juwels-Booster Node: 2x AMD EPYC 7402, 4x NVIDIA A100 40GB, 4x NVIDIA Mellanox ConnectX-6 200 Gbps HCAs。

3.4 标量 Put 合并

将多个 NVSHMEM "p" 操作合并为一个工作请求
* IBGDA 合并 (IBGDA Coalescing) 在 128 个 CTA 时达到 1935 MOPS 的消息速率，远高于 IBRC 的 180 MOPS 和 IBGDA 的 1.7 MOPS。

3.5 二进制分发改进

改进了可移植性和可用性
* HPC_SDK 在多个平台上分发 NVSHMEM 及其他兼容软件包。
* HPC_SDK 中使用的 NVSHMEM 二进制包的近期更改通常也很有用。
* NVSHMEM 现在支持 CUDA 次要版本兼容性。
* 编译器限制：所有静态链接的部分必须使用最旧的编译器版本构建。

• 模块化传输从主要的 NVSHMEM 库中移除了特定于传输的要求。

  • 用于构建传输的文件现已包含在软件包中。

• 模块化引导接口已得到改进。

3.6 CMAKE 构建系统

• CMake 支持构建所有 NVSHMEM 组件 (源码安装、软件包、测试、性能测试、示例)。
• 替代 Makefiles，支持所有 NVSHMEM_* 环境变量。
• 也可以通过 CMake 原生方式提供选项。
• 使用 CMake 命令行选项 "CUDA_ARCHITECTURES" 更改配置的架构，而非 NVCC_GENCODE。
• 生成的 NVSHMEMConfig.cmake 文件允许轻松地将目标集成到现有 CMake 项目中。
• 自 v2.9 起，NVSHMEM 正式弃用 Makefiles，转而支持 CMake 构建系统。

3.7 内联设备代码

改进了延迟
* 设备 API 实现：
* P2P 内存复制 (内联)
* 远程通信 (未内联)
* 集合 API 实现 (未内联)

• 此功能使所有设备 API 实现都能够内联。
• 通过 NVSHMEM 构建时宏启用：NVSHMEM_ENABLE_ALL_DEVICE_INLINING = 0/1。
• 所有设备 API 实现已移至 NVSHMEM 头文件。
• 通过 __forceinline__ 限定符启用内联。
• 优点：
  • 无函数调用开销。
  • CUDA 编译器可以全面了解内核代码，从而改进编译器优化、寄存器使用等。

3.8 集合优化

3.8.1 使用 LL 协议改进延迟

• 典型的消息传递方法：nvshmem_put + nvshmem_fence + nvshmem_signal。
• LL (低延迟) 协议：
  • 将 4 字节数据与 4 字节标志交织。
  • 消息大小加倍但同步成本低。
  • 适用于中小型消息的低延迟。
  • 乒乓延迟最多提高 75%。

3.8.2 拓扑感知优化

• 拓扑感知算法：

  • 形成子团队：team_node 和 team_same_mype_node。

• 广播：

  • 向节点领导者广播。
  • 节点内广播。

• 规约 (Reduce)：

  • 扁平化的 All-to-all：

    • 每个 PE 将数据发送给其他所有 PE。
    • 每个 PE 执行所有数据的本地规约。

  • 拓扑感知 All-to-all：

    • 在节点内执行扁平化的 all-to-all 规约。
    • 跨 team_same_mype_node 执行扁平化的 all-to-all 规约。

• 性能提高 95%。

---

4. 性能案例研究

4.1 NVSHMEM + cuFFT = cuFFTMp

• 多节点多 GPU：使用 NVIDIA cuFFTMp FFTs 进行规模化计算。

  • 更多信息请访问：https://developer.nvidia.com/blog/multi-node-multi-gpu-using-nvidia-cufftmp-ffts-at-scale

• JAX + cuFFTMp：

  • JAX-only、JAX+cuFFTMp 和 cuFFTMp 在 Selene 上进行 3D C2C FP32 FFT 的强扩展性测试 (2022 年 12 月)。
  • JAX 最新进展 [S51956] - Frederic 等人。

• GROMACS + cuFFTMp：

  • STMV 强扩展性 (1M 原子)
  • BenchPEP-h 强扩展性 (12M 原子)
  • NVIDIA 数学库最新进展 [S51176] - Harun Bayraktar 等人。

4.2 NVSHMEM 在 GROMACS 中的应用

• 如前一幻灯片所述，cuFFTMp 自动使用 NVSHMEM。
• 在执行分布式 FFTs 时，PME GPU 之间存在全对全通信模式。

• GROMACS 中还存在其他可能受益于 NVSHMEM 的通信模式：

  • PP (粒子-粒子) Halo 交换
  • PME (粒子网格 Ewald) Halo 交换
  • PP-PME 通信

• 原型版本使用 NVSHMEM 进行 PP halo 交换通信。

• 为了评估性能，使用了反应场 (RF) 案例。
• 作为 PME 远距离力计算的替代方案。
• 在扩展到多个节点时，对 PP Halo 交换特别敏感。
• 在扩展到多个节点时，PP halo 交换具有显著的性能优势。
• 参考：分子动力学前沿 CUDA 技术 [S51110] - Alan Gray 等人。

4.3 NVSHMEM 在高性能 Linpack (HPL) 中的应用

• 传统 HPL：

  • GPU 用于有限的操作，例如 GEMM。
  • 因子分解 (LU 分解) 和通信由 CPU 发起和处理。

• HPL + NVSHMEM：

  • 无 CPU 版本。
  • 所有计算和通信都在 GPU 上完成。
  • 矩阵因子分解和 GEMM 在不同的内核中并行完成。
  • 因子分解需要频繁的细粒度通信。
  • HPL 实现变为无 CPU，因此无需针对不同类型的 CPU 进行优化。
  • 性能提升高达 8%。

---

5. 即将推出的特性和结论

5.1 NVSHMEM 中的 Grace Hopper 支持

本节详细介绍了 Grace Hopper 超级芯片架构及其在 NVSHMEM 中的支持。

• NVIDIA Grace Hopper 超级芯片

  • 

  • Grace Hopper 超级芯片支持 Grace Arm CPU 和 Coherent Hopper GPU。
  • Grace CPU: 配备 2 个 LPDDR5X，总计 512 GB 内存，带宽为 546 GB/s。
  • Hopper GPU: 配备 2 个 HBM3，总计 96 GB 内存，带宽为 3000 GB/s。
  • NVLink C2C: 连接 CPU 和 GPU，提供 900 GB/s 带宽。
  • 高速 I/O: 包含 4x 16x PCIe-5 接口，提供 512 GB/s 带宽。
  • NVLink 网络: 包含 18x NVLink 4 接口，提供 900 GB/s 带宽，可连接多达 256 个 GPU。
  • 架构亮点包括：高速网络 I/O、大容量主内存、高 CPU/GPU 带宽、硬件一致性以及第四代 NVLink。

5.2 第四代 NVLink

第四代 NVLink 进一步增强了互联性能。

• 每个 NVSwitch 拥有 64 个 NVLink 端口。
• 端口配备 NVIDIA Scalable Hierarchical Aggregation Reduction Protocol (SHARP)™ 引擎。
• 提供网络内规约 (in-network reductions) 和多播加速 (multicast acceleration) 功能。
• 优化 NVSHMEM Collectives:
  • 可将大型 all-reduce 操作的带宽加倍。
  • 小消息集合操作的延迟更低。
  • 屏障延迟降低。
    

5.3 NVSHMEM 内存空间与上下文

本节阐述了 NVSHMEM 内存空间和上下文，以实现对 Grace CPU 内存的访问并改进通信管理。

• 内存空间 (Memory Spaces)

  • 用户定义与团队关联的对称内存。
  • 包括新的对称堆、注册用户缓冲区和继承的对称空间。
  • 优势: 增加了灵活性，提供了对额外内存类型的访问。

• 上下文 (Contexts)

  • 管理一组 RMA (远程内存访问) 和 AMO (原子内存操作)。
  • Quiet 和 Fence 操作被限定在特定上下文内。
  • 优势: 线程/SM 隔离减少了竞争（例如，IBGDA QP 分配给 SM），实现流水线非阻塞通信。
    
  • 图示展示了 nvshmem_team_split 创建 Team，其中包含内存空间参数 team_config_t 和 Memory Space。Memory Space 包含 Symmetric Heap 和 Registered User Buffer(s)。
  • nvshmem_ctx_create 创建 Context，与 Resources 通过 ctx_options_t 关联。Context 处理 NVSHMEM RMA and AMO Operations Quiet and Fence。

5.4 总结

NVSHMEM 是一个用于 NVIDIA GPU 集群的 PGAS (Partitioned Global Address Space) 库。

• NVSHMEM 可无缝扩展，支持：

  • 使用 NVLink 连接 GPU 的节点级 GPU 编程。
  • 通过 InfiniBand 或 RoCE 连接的多节点 GPU 集群。

• NVSHMEM 提供 Stream/Graph、GPU 内核启动和 CPU 启动的 API。

• 与 CUDA 编程模型的集成可以提高 GPU 集群的性能并简化扩展。
  

5.5 NVSHMEM 状态

NVSHMEM 的最新发布信息和关键特性。

• 最新版本: NVSHMEM 2.9.0，于 2023 年 3 月发布。
• 可在 https://developer.nvidia.com/nvshmem 下载。
• NVSHMEM 2.8 作为 NVIDIA HPC SDK 发布的一部分提供，容器可通过 NGC 获取。

• 新特性 (New Features):

  • Hopper GPU 支持。
  • 优化的 GDA-KI 网络传输。
  • 设备代码内联。
  • CMake 构建系统。
  • 改进的集合操作性能。

• 实现特性 (Implementation Features):

  • NVLink 和 PCIe P2P 支持。
  • InfiniBand、RoCE 和 Slingshot (beta) 支持。
  • X86 和 Power9 支持。
  • 可与 MPI 和 OpenSHMEM 互操作。