博主未授权任何人或组织机构转载博主任何原创文章，感谢各位对原创的支持！
博主链接

本人就职于国际知名终端厂商，负责modem芯片研发。
在5G早期负责终端数据业务层、核心网相关的开发工作，目前牵头6G算力网络技术标准研究。

博客内容主要围绕：
       5G/6G协议讲解
       算力网络讲解（云计算，边缘计算，端计算）
       高级C语言讲解
       Rust语言讲解

利用蒙特卡罗法求解 𝜋 的近似值（NVSHMEM）

在这里插入图片描述

NVSHMEM

NVSHMEM是一个并行编程模型，用于在多个 NVIDIA GPU 之间进行高效和可扩展的通信。NVSHMEM 依托于 OpenSHMEM 构建而成，可为横跨多个 GPU 内存的数据提供全局地址空间，并可通过细粒度的由 GPU 发起的操作、由 CPU 发起的操作和 CUDA 流操作访问该空间。NVSHMEM 为许多应用提供了令人信服的多 GPU 编程模型，对于具有高密度 GPU 和复杂互连的现代 GPU 服务器（例如NVIDIA DGX A100 服务器的 NVIDIA NVSwitch）来说尤其有价值。

请添加图片描述

为什么使用NVSHMEM

传统上，涉及多服务器 GPU 的通信模式可能看起来如下所示：计算发生在 GPU 上，而通信在同步 GPU 后发生在 CPU 上（确保数据发送有效）。虽然这种方法很容易编程，但会在应用的关键路径上引入初始化通信或启动核函数的延迟。我们会丧失计算与通信重叠的能力。如果我们通过流水操作工作来重叠通信和计算，延迟确实可以部分地隐藏，但代价是让应用变得更加复杂。

请添加图片描述

相反，在使用 GPU 而不是 CPU 启动的通信模型中，我们直接利用 GPU 同时进行计算和通信。我们可以用这种方式编写细粒度的通信模式，并且可以通过 GPU 架构的本质来隐藏通信延迟（在 GPU 架构中，计算中的Warp可以继续进行，而其他的Warps则会停下来等待数据）。

请添加图片描述

启动NVSHMEM应用

与 MPI 一样，NVSHMEM 也是具有 SPMD 编程风格的示例之一。NVSHMEM 提供了一个启动脚本¹，其名为nvshmrun，可用于处理启动 𝑀 个进程。nvshmrun的参数是-np，也就是要启动的进程数，然后是应用程序的可执行文件，然后是该可执行文件的任何参数。每个独立进程又名为处理单元 (PE)，有一个唯一的（零索引的）数字标识符与之相关联²。

请添加图片描述

初始化及终止 NVSHMEM

作为主机端的核心需求，我们必须初始化并终止 NVSHMEM，将这两者作为程序中的第一项和最后一项。

nvshmem_init();
...
nvshmem_finalize();

获取处理单元的 ID

API 调用 nvshmem_my_pe() 返回每个 PE 的唯一数字 ID。

int my_pe = nvshmem_my_pe();
int device = my_pe;
cudaSetDevice(device);

在多节点环境中，您必须考虑到一个事实，即 CUDA 设备在每个节点中始终都是零索引的。在这种情况下，您将获得仅对该节点有意义的本地 PE 标识符。例如，如果我们使用两个节点，每个节点有四个 GPU，那么我们将要求工作启动程序在每个节点上运行四个任务（如nvshmrun -np 8 -ppn 4 -hosts hostname1,hostname2），然后完成³：

int my_pe_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE);
int device = my_pe_node;
cudaSetDevice(device);

编译 NVSHMEM 代码

编译看起来和以前相似，但我们现在需要为 NVSHMEM 指向相关的文件包含命令include和库目录（-I $NVSHMEM_HOME/include -L $NVSHMEM_HOME/lib -lnvshmem）以及 CUDA 驱动 API 中的链接（-lcuda）。我们还需要把#include <nvshmem.h>⁴ 和 #include <nvshmemx.h>⁵ 添加到代码中。最后，我们需要添加-rdc=true以启用 浮动设备代码，这是 NVSHMEM 的一项需求。

练习1：使用带有 MC π 代码的 NVSHMEM

#include <iostream>
#include <curand_kernel.h>

#include <nvshmem.h>
#include <nvshmemx.h>

inline void CUDA_CHECK (cudaError_t err) {
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA error: %s\n", cudaGetErrorString(err));
        exit(-1);
    }
}

#define N 1024*1024

__global__ void calculate_pi(int* hits) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;

    // 初始化随机数状态（网格中的每个线程不得重复）
    int seed = 0;
    int offset = 0;
    curandState_t curand_state;
    curand_init(seed, idx, offset, &curand_state);

    // 在 (0.0, 1.0] 内生成随机坐标
    float x = curand_uniform(&curand_state);
    float y = curand_uniform(&curand_state);

    // 如果这一点在圈内，增加点击计数器
    if (x * x + y * y <= 1.0f) {
        atomicAdd(hits, 1);
    }
}


int main(int argc, char** argv) {
    // 初始化 NVSHMEM
    nvshmem_init();

    // 获取 NVSHMEM 处理元素 ID
    int my_pe = nvshmem_my_pe();

    // 每个 PE（任意）选择与其 ID 对应的 GPU
    int device = my_pe;
    CUDA_CHECK(cudaSetDevice(device));

    // 分配主机和设备值
    int* hits;
    hits = (int*) malloc(sizeof(int));

    int* d_hits;
    CUDA_CHECK(cudaMalloc((void**) &d_hits, sizeof(int)));

    // 初始化点击次数并复制到设备
    *hits = 0;
    CUDA_CHECK(cudaMemcpy(d_hits, hits, sizeof(int), cudaMemcpyHostToDevice));

    // 启动核函数进行计算
    int threads_per_block = 256;
    int blocks = (N + threads_per_block - 1) / threads_per_block;

    calculate_pi<<<blocks, threads_per_block>>>(d_hits);
    CUDA_CHECK(cudaDeviceSynchronize());

    // 将最终结果复制回主机
    CUDA_CHECK(cudaMemcpy(hits, d_hits, sizeof(int), cudaMemcpyDeviceToHost));

    // 计算 pi 的最终值
    float pi_est = (float) *hits / (float) (N) * 4.0f;

    // 打印结果
    std::cout << "Estimated value of pi on PE " << my_pe << " = " << pi_est << std::endl;
    std::cout << "Relative error on PE " << my_pe << " = " << std::abs((M_PI - pi_est) / pi_est) << std::endl;

    free(hits);
    CUDA_CHECK(cudaFree(d_hits));

    // 最终确定 nvshmem
    nvshmem_finalize();

    return 0;
}

编译和运行命令如下：

nvcc -x cu -arch=sm_70 -rdc=true -I $NVSHMEM_HOME/include -L $NVSHMEM_HOME/lib -lnvshmem -lcuda -o nvshmem_pi_step1 exercises/nvshmem_pi_step1.cpp
nvshmrun -np $NUM_DEVICES ./nvshmem_pi_step1

您运行了首个 NVSHMEM 程序，但很遗憾，我们所做的工作没什么意思，因为每个 PE 都执行同样的工作。（您可通过比较所有 PE 的输出来检查这句话是否正确。）理想情况下，我们希望将工作分散到不同的 PE 和 GPU 上。

练习2：跨 PE 分配工作

在本练习中，每个 GPU 将 𝑁 个样本点除以 PE 的数量 𝑀 。我们可以使用 API nvshmem_n_pes()来获得：

int n_pes = nvshmem_n_pes();

然后将 𝑁 除以 n_pes就行了。为了让PE的工作更有意思，我们执行一个额外的步骤，即为每个 PE 选择各自唯一的随机数的种子，这样可以让每个 GPU 做不同的工作：

int seed = nvshmem_my_pe();

#include <iostream>
#include <curand_kernel.h>

#include <nvshmem.h>
#include <nvshmemx.h>

inline void CUDA_CHECK (cudaError_t err) {
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA error: %s\n", cudaGetErrorString(err));
        exit(-1);
    }
}

#define N 1024*1024

__global__ void calculate_pi(int* hits, int seed) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;

    // 初始化随机数状态（网格中的每个线程不得重复）
    int offset = 0;
    curandState_t curand_state;
    curand_init(seed, idx, offset, &curand_state);

    // 在 (0.0, 1.0] 内生成随机坐标
    float x = curand_uniform(&curand_state);
    float y = curand_uniform(&curand_state);

    // 如果这一点在圈内，增加点击计数器
    if (x * x + y * y <= 1.0f) {
        atomicAdd(hits, 1);
    }
}


int main(int argc, char** argv) {
    // 初始化 NVSHMEM
    nvshmem_init();

    // 获取 NVSHMEM 处理元素 ID 和 PE 数量
    int my_pe = nvshmem_my_pe();
    int n_pes = nvshmem_n_pes();

    // 每个 PE（任意）选择与其 ID 对应的 GPU
    int device = my_pe;
    CUDA_CHECK(cudaSetDevice(device));

    // 分配主机和设备值
    int* hits;
    hits = (int*) malloc(sizeof(int));

    int* d_hits;
    CUDA_CHECK(cudaMalloc((void**) &d_hits, sizeof(int)));

    // 初始化点击次数并复制到设备
    *hits = 0;
    CUDA_CHECK(cudaMemcpy(d_hits, hits, sizeof(int), cudaMemcpyHostToDevice));

    // 启动核函数进行计算
    int threads_per_block = 256;
    int blocks = (N / n_pes + threads_per_block - 1) / threads_per_block;

    int seed = my_pe;
    calculate_pi<<<blocks, threads_per_block>>>(d_hits, seed);
    CUDA_CHECK(cudaDeviceSynchronize());

    // 将最终结果复制回主机
    CUDA_CHECK(cudaMemcpy(hits, d_hits, sizeof(int), cudaMemcpyDeviceToHost));

    // 计算 pi 的最终值
    float pi_est = (float) *hits / (float) (N / n_pes) * 4.0f;

    // 打印结果
    std::cout << "Estimated value of pi on PE " << my_pe << " = " << pi_est << std::endl;
    std::cout << "Relative error on PE " << my_pe << " = " << std::abs((M_PI - pi_est) / pi_est) << std::endl;

    free(hits);
    CUDA_CHECK(cudaFree(d_hits));

    // 最终确定 nvshmem
    nvshmem_finalize();

    return 0;
}

编译和运行命令如下：

nvcc -x cu -arch=sm_70 -rdc=true -I $NVSHMEM_HOME/include -L $NVSHMEM_HOME/lib -lnvshmem -lcuda -o nvshmem_pi_step2 exercises/nvshmem_pi_step2.cpp
nvshmrun -np $NUM_DEVICES ./nvshmem_pi_step2

在这里插入图片描述

nvshmrun本质上是一个到 Hydra 流程管理器的符号链接。虽然我们演示了 NVSHMEM 在独立运行的计算机中的使用情况，但 NVSHMEM 可与 MPI 作业启动环境兼容，比如它也可以与 Slurm 一起使用。如果使用 MPI 或 OpenSHMEM 启动作业，则相关代码修改如下所示。对于 MPI，我们首先初始化 MPI，然后在 MPI 上引导 NVSHMEM 初始化。关闭时，我们要先终止 NVSHMEM，然后终止 MPI。

int main() {
    MPI_Init(&argc, &argv);

    nvshmemx_init_attr_t attr;
    MPI_Comm comm = MPI_COMM_WORLD;
    attr.mpi_comm = &comm;

    nvshmemx_init_attr(NVSHMEMX_INIT_WITH_MPI_COMM, &attr);

    // ...

    nvshmem_finalize();

    MPI_Finalize();

    return 0;
}

在 OpenSHMEM 作业中，我们会改用以下做法

int main() {
    shmem_init();

    nvshmemx_init_attr_t attr;
    nvshmemx_init_attr(NVSHMEMX_INIT_WITH_SHMEM, &attr);

    // ...

    nvshmem_finalize();

    shmem_finalize();

    return 0;
}

↩︎

NVSHMEM 作为 OpenSHMEM 的实现，有许多术语都与 OpenSHMEM 相通（如 PE），并具有非常相似的 API。熟悉 MPI 的读者会发现 PE 类似于 MPI rank。 ↩︎
API nvshmem_team_my_pe() 是 NVSHMEM 2.0 中的新功能。请查看这篇博客，了解更多信息。 ↩︎
nvshmem.h提供符合 OpenSHMEM 标准的 API，类似于 nvshmem_*。 ↩︎
nvshmemx.h提供 NVIDIA 专用的扩展程序，类似于nvshmemx_*。 ↩︎