博主未授权任何人或组织机构转载博主任何原创文章，感谢各位对原创的支持！
博主链接
本人就职于国际知名终端厂商，负责modem芯片研发。
在5G早期负责终端数据业务层、核心网相关的开发工作，目前牵头6G算力网络技术标准研究。
博客内容主要围绕：
5G/6G协议讲解
算力网络讲解（云计算，边缘计算，端计算）
高级C语言讲解
Rust语言讲解
NVSHMEM 内存模型
PE：处理单元（process entity）
对称内存
NVSHMEM 的内存分配 API nvshmem_malloc()，其工作方式有点类似于标准的cudaMalloc()，但cudaMalloc()会返回一个本地 GPU 的私有地址1。使用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 在程序执行期间2创建，然后用于后续的动态内存分配。
练习1
下面我们把cudaMalloc()的调用替换为nvshmem_malloc()的调用。我们仍然可以对分配在本地的数据使用atomicAdd()，这样每个 PE 上的对称对象副本就会得到与之前相同的结果。
其次，我们对所有 PE 的结果求和。这是一次联合操作，它是全局归约操作。在 NVSHMEM 中，我们可以使用 nvshmem_int_sum_reduce(team, dest, source, nreduce) 对对称对象的所有实例求和。
- source：是我们要求和的对称地址；
- destination：是储存结果的地方；
- nreduce：是要归约的元素个数（对我们而言只有一个，因为我们的数据是标量）；
- team：是要进行求和运算的一组 PE3（我们将使用默认组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
