GPUNvidia CUDA 编程高级教程——NVSHMEM 内存模型

Posted 2023-01-16 从善若水

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

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


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

NVSHMEM 内存模型

PE：处理单元（process entity）

对称内存

       NVSHMEM 的内存分配 API nvshmem_malloc()，其工作方式有点类似于标准的cudaMalloc()，但cudaMalloc()会返回一个本地 GPU 的私有地址¹。使用nvshmem_malloc()分配的对象称为对称数据对象。每个对称数据对象在所有 PE 上都有一个名称、类型和大小相同的对应数据对象。由nvshmem_malloc()返回的指针对应的虚拟地址称为对称地址。在 NVSHMEM 通信例程中使用对称地址对其他 PE 进行远程访问是合法的（对称地址也可以直接用于对 PE 本地内存的访问）。我们可以像操作普通本地地址一样操作虚拟地址。如要使用 NVSHMEM API 访问远程 PE 上的对称数据对象副本，我们可以像通常那样以指针作为存储索引，并使用远程目标 PE 中的相应位置。例如，

       如果我们执行了下面的语句：

int* a = (int*) nvshmem_malloc(sizeof(int));

那么我们既可以在本地 PE 上进行本地内存访问，也可以在远程 PE 上进行远程内存访问，来获取a[0]的值。理解这个操作的一种思考方法是，给定 M 个 PE，我们将长度为M的数组里的数据元素均匀地分配到所有 PE 上，这样每个 PE 只有一个元素。由于在本例中，对称数据对象的长度为 1，我们在任何 PE 上只需访问a[0]。

       在 NVSHMEM 中，对称数据对象的动态内存分配来自一个名为对称堆（symmetric heap）的特殊内存区域，由 NVSHMEM 在程序执行期间²创建，然后用于后续的动态内存分配。

练习1

       下面我们把cudaMalloc()的调用替换为nvshmem_malloc()的调用。我们仍然可以对分配在本地的数据使用atomicAdd()，这样每个 PE 上的对称对象副本就会得到与之前相同的结果。

       其次，我们对所有 PE 的结果求和。这是一次联合操作，它是全局归约操作。在 NVSHMEM 中，我们可以使用 nvshmem_int_sum_reduce(team, dest, source, nreduce) 对对称对象的所有实例求和。

• source：是我们要求和的对称地址；
• destination：是储存结果的地方；
• nreduce：是要归约的元素个数（对我们而言只有一个，因为我们的数据是标量）；
• team：是要进行求和运算的一组 PE³（我们将使用默认组NVSHMEM_TEAM_WORLD，这是所有 PE 的集合）;

总而言之，我们要做的是：

// 累积所有 PE 的结果
int* d_hits_total = (int*) nvshmem_malloc(sizeof(int));
nvshmem_int_sum_reduce(NVSHMEM_TEAM_WORLD, d_hits_total, d_hits, 1);

现在，所有的 PE 都有计数的总和，所以我们要做的第三个更改就是只需要在单个 PE 上打印结果。按照惯例，我们通常在 PE0 上进行打印。

if (my_pe == 0) 
    // 将最终结果复制回主机
    ...

    // 计算 pi 的最终值
    ...

    // 打印结果
    ...

完整代码如下（file name：nvshmem_pi_step3.cpp）：

#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 = (int*) malloc(sizeof(int));
    int* d_hits = (int*) nvshmem_malloc(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());

    // 累积所有 PE 的结果
    int* d_hits_total = (int*) nvshmem_malloc(sizeof(int));
    nvshmem_int_sum_reduce(NVSHMEM_TEAM_WORLD, d_hits_total, d_hits, 1);

    if (my_pe == 0) 
        // 将最终结果复制回主机
        CUDA_CHECK(cudaMemcpy(hits, d_hits_total, sizeof(int), cudaMemcpyDeviceToHost));

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

        // 打印结果
        std::cout << "Estimated value of pi averaged over all PEs = " << pi_est << std::endl;
        std::cout << "Relative error averaged over all PEs = " << std::abs((M_PI - pi_est) / pi_est) << std::endl;
    

    free(hits);
    nvshmem_free(d_hits);
    nvshmem_free(d_hits_total);

    // 最终确定 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_step3 exercises/nvshmem_pi_step3.cpp
nvshmrun -np $NUM_DEVICES ./nvshmem_pi_step3

结果如下：

Estimated value of pi averaged over all PEs = 3.14072
Relative error averaged over all PEs = 0.000277734

---

---

1. 例外情况在于，在使用 NVLink 连接 GPU 的系统中，可以使用 CUDA IPC 机制 让 GPU 直接访问彼此的内存。 ↩︎

2. 对称堆的默认大小是 1GB，可通过环境变量 NVSHMEM_SYMMETRIC_SIZE 加以控制。 ↩︎

3. 在 OpenSHMEM 1.5 规范的基础上，使用team指定涉及多个 PE 组的操作是 NVSHMEM 2.0 的新功能。 ↩︎