NVSHMEM 基础回顾：NVSHMEM 简介、对称内存与 PE 概念、基本的 put/get 操作
各位同学，咱们今天聊聊 NVSHMEM 的基础。说实话，我刚开始接触这东西的时候，也觉得有点绕。但别急，咱们一步步来。
NVSHMEM 是什么？
NVSHMEM，全称是 NVIDIA Shared Memory Model。说白了，它是一套让 GPU 之间可以直接通信的库。你想想看，传统做法里，GPU 之间要交换数据，得先拷回 CPU 内存，再拷过去。多麻烦啊！
NVSHMEM 就不一样了。它让 GPU 能直接读写另一个 GPU 的内存。就像两个邻居，不用通过邮局，直接隔墙喊话就行。
核心思想：NVSHMEM 把多个 GPU 的内存虚拟成一个统一的全局地址空间。每个 GPU 都能直接访问这个空间里的任何位置。
我在项目中遇到过这样一个场景：8 张 GPU 卡做并行计算，每张卡只处理自己那部分数据。但边界上的数据需要互相交换。用传统 MPI 的话，代码写起来又长又容易出错。换成 NVSHMEM 后，代码量直接砍了一半。
对称内存与 PE 概念
好，接下来咱们聊聊两个核心概念：对称内存和 PE。
什么是对称内存？
对称内存，就是每个 PE（Processing Element，处理单元）上都有一块大小相同、布局相同的内存区域。嗯，这里要注意：不仅仅是大小一样，连里面的数据结构布局都得一样。
举个例子：
// 每个 PE 上都会分配 1024 个 float
float *buf = (float*)nvshmem_malloc(1024 * sizeof(float));这行代码一执行，每个 PE 上都会有一块 1024 个 float 的空间。而且它们在虚拟地址空间中的偏移量是一样的。
我的经验：对称内存是 NVSHMEM 的基石。你想想看，如果每个 PE 上的内存布局不一样，那一个 PE 怎么知道另一个 PE 的数据放在哪？所以，设计数据结构时，一定要保证对称性。
PE 是什么？
PE 就是处理单元。在 NVSHMEM 的世界里，一个 PE 可以是一个 GPU，也可以是一个 CPU 线程。每个 PE 都有一个唯一的编号，从 0 到 npes-1。
我习惯把 PE 想象成一个个独立的小房间。每个房间里有自己的柜子（对称内存），房间之间有条走廊（NVSHMEM 通信通道）。
| 概念 | 类比 | 说明 |
|---|---|---|
| PE | 小房间 | 每个 GPU 或 CPU 线程 |
| 对称内存 | 房间里的柜子 | 每个 PE 上大小相同的内存 |
| my_pe | 房间号 | 当前 PE 的编号 |
| npes | 房间总数 | 所有 PE 的数量 |
基本的 put/get 操作
好了，有了对称内存和 PE 的概念，咱们来看看怎么实际干活。NVSHMEM 提供了两个最基本的操作：put 和 get。
put 操作：把数据放过去
put 操作，就是把当前 PE 的数据，写到另一个 PE 的对称内存里。
// 把本地 src 中的数据，写到 PE 1 的 dest 位置
float src[4] = {1.0, 2.0, 3.0, 4.0};
nvshmem_float_put(&dest[0], &src[0], 4, 1);这行代码的意思很直白：把本地 src 数组里的 4 个 float，写到 PE 1 的 dest 数组里。
避坑指南：我曾经犯过一个错误——忘了 put 操作是异步的。也就是说，nvshmem_float_put 返回时，数据可能还没真正到达目标 PE。如果你紧接着就去读目标 PE 上的数据，可能会读到旧值。记得用 nvshmem_barrier() 或 nvshmem_quiet() 来同步。
get 操作：把数据拿过来
get 操作正好相反，是从另一个 PE 的对称内存里，把数据读到当前 PE。
// 从 PE 2 的 src 位置，读 4 个 float 到本地 dest
float dest[4];
nvshmem_float_get(&dest[0], &src[0], 4, 2);get 操作也是异步的。嗯，这里要注意：get 操作虽然异步，但数据最终会到达本地内存。你可以在 get 之后调用 nvshmem_barrier() 确保数据到位。
知识体系总览
下面这张图，是我自己画的 NVSHMEM 基础知识结构。你看一眼，心里就有谱了。
一个小例子：两个 PE 之间的数据交换
咱们来看个完整的例子。假设有两个 PE，PE 0 和 PE 1。PE 0 想把一个数发给 PE 1。
#include <nvshmem.h>
#include <stdio.h>
int main(void) {
nvshmem_init();
int my_pe = nvshmem_my_pe();
int npes = nvshmem_n_pes();
// 分配对称内存
float *data = (float*)nvshmem_malloc(sizeof(float));
*data = 0.0f;
if (my_pe == 0) {
*data = 3.14f;
// PE 0 把数据发给 PE 1
nvshmem_float_put(data, data, 1, 1);
nvshmem_quiet(); // 确保 put 完成
}
nvshmem_barrier_all(); // 全局同步
if (my_pe == 1) {
printf("PE 1 收到数据: %f\n", *data);
}
nvshmem_free(data);
nvshmem_finalize();
return 0;
}这个例子很简单，但包含了 NVSHMEM 的核心流程：初始化、分配对称内存、put 操作、同步、收尾。
我的建议：刚开始学 NVSHMEM 时，先把这个例子跑通。然后试着改一改，比如改成 get 操作，或者改成多个 PE 之间的通信。动手实践是最好的学习方式。
好了，NVSHMEM 的基础就聊到这儿。对称内存和 PE 的概念，加上 put/get 操作，就是后续所有高级功能的地基。你把这些搞明白了，后面学光晕交换就轻松多了。