1. NVSHMEM 概述：什么是 NVSHMEM？

各位同学好，我是你们这门课的老朋友。今天咱们聊聊 NVSHMEM。说实话，我第一次接触这东西的时候，心里也犯嘀咕——这不就是个通信库吗？但用了一段时间后，我发现它跟传统的 MPI、SHMEM 还真不太一样。

NVSHMEM，全称是 NVIDIA SHMEM。它是个并行编程库，专门为 GPU 集群设计的。说白了，它让你能直接在 GPU 之间读写数据，不用经过 CPU 中转。你想想看，这省了多少事？

1.1 NVSHMEM 与 MPI、SHMEM 的关系

先说说 MPI。MPI 是消息传递接口，大家都很熟悉。我在项目中用过很多次 MPI，它的模型是「发送-接收」。你要给另一个进程发数据，得先调用 MPI_Send，对方调用 MPI_Recv。这就像寄快递——你得打包、填单子，对方得签收。

NVSHMEM 呢？它走的是「单边通信」路线。什么意思？就是你可以直接往别人的 GPU 内存里写数据，不需要对方配合。我记得第一次用这个特性时，感觉就像有了万能钥匙——直接开门进去放东西，不用敲门。

那 SHMEM 呢？SHMEM 是 Cray 公司搞出来的一个经典库，也是单边通信的鼻祖。NVSHMEM 其实就是 SHMEM 思想在 GPU 上的实现。不过，NVSHMEM 加了很多针对 GPU 的优化，比如流式通信、异步操作这些。

特性	MPI	SHMEM	NVSHMEM
通信模型	双边（Send/Recv）	单边（Put/Get）	单边（Put/Get）
硬件支持	CPU 网络	CPU 网络	GPU + NVLink + InfiniBand
同步方式	显式同步	隐式 + 显式	隐式 + 显式 + 流同步
典型延迟	微秒级	微秒级	亚微秒级（GPU 直通）

嗯，这里要注意：NVSHMEM 不是要取代 MPI。我个人的习惯是，在需要大规模数据交换时用 MPI，在需要细粒度、低延迟通信时用 NVSHMEM。两者可以混用，但你要清楚各自的适用场景。

1.2 NVSHMEM 的核心概念

要理解 NVSHMEM，你得先搞懂三个核心概念：PE、对称内存、对称堆。咱们一个一个说。

1.2.1 PE（Processing Element）

PE 就是处理单元。在 NVSHMEM 里，每个 GPU 上的每个线程块（或者线程）都可以是一个 PE。不过，更常见的做法是把每个 GPU 当作一个 PE。比如你有 4 块 GPU，那就是 4 个 PE，编号从 0 到 3。

每个 PE 有自己的私有内存，也有共享的对称内存。PE 之间通过 PE ID 来区分。你调用 nvshmem_my_pe() 就能拿到自己的 ID，调用 nvshmem_n_pes() 知道总共有多少个 PE。

1.2.2 对称内存

对称内存是 NVSHMEM 最核心的概念。什么叫对称？就是每个 PE 上都有一块内存，它们的大小、类型、甚至地址偏移都一模一样。你想想看，这就像每个 PE 都有一块相同布局的「格子」，你往自己的格子里写数据，其他 PE 也能看到。

对称内存是通过 nvshmem_malloc 分配的。举个例子：

// 每个 PE 都分配 1024 个 float
float *sym_buf = (float*)nvshmem_malloc(1024 * sizeof(float));

// PE 0 往自己的对称内存里写数据
if (nvshmem_my_pe() == 0) {
    sym_buf[0] = 3.14f;
}

// 所有 PE 都能读到这个值（但需要同步）
nvshmem_barrier_all();
printf("PE %d: sym_buf[0] = %f\n", nvshmem_my_pe(), sym_buf[0]);

这里有个坑，我曾经踩过：对称内存的地址在每个 PE 上虽然「逻辑上」是对称的，但物理地址完全不同。你不能直接拿 PE 0 的指针去 PE 1 上解引用。必须用 NVSHMEM 的 API 来访问。

1.2.3 对称堆

对称堆就是对称内存的「集合」。每个 PE 上都有一个对称堆，它包含了所有通过 nvshmem_malloc 分配的内存块。你可以把对称堆想象成一个大的内存池，所有 PE 都能往里面放东西，也能从里面拿东西。

对称堆的大小是在初始化时指定的。比如：

// 初始化 NVSHMEM，指定对称堆大小为 256MB
nvshmem_init();
nvshmemx_init_attr_t attr;
attr.heap_size = 256 * 1024 * 1024;  // 256 MB
nvshmemx_init_attr(NVSHMEMX_INIT_ATTR_HEAP_SIZE, &attr);

我个人建议，对称堆的大小要根据你的数据量来定。太小了不够用，太大了浪费显存。我一般会先估算一下每个 PE 需要多少数据，再乘以 1.2 作为安全余量。

1.3 为什么需要 NVSHMEM？

你可能会问：我用 CUDA 的 cudaMemcpy 不也能在 GPU 之间传数据吗？为什么还要学 NVSHMEM？

原因很简单：cudaMemcpy 是 CPU 发起的，它需要 CPU 参与每一次数据传输。而 NVSHMEM 是 GPU 发起的，GPU 上的线程可以直接触发通信。这意味着什么？意味着你可以把通信和计算重叠起来，让 GPU 一边算一边传，效率高得多。

我记得有个项目，用 cudaMemcpy 做数据交换，通信开销占了总时间的 40%。换成 NVSHMEM 后，通过流式通信和计算重叠，通信开销降到了 15% 以下。这就是 NVSHMEM 的价值所在。

1.4 本章小结

好了，咱们捋一下这章的核心内容：

• NVSHMEM 是 GPU 上的单边通信库，让 GPU 直接读写其他 GPU 的数据。
• 它跟 MPI 的区别是通信模型不同（单边 vs 双边），跟 SHMEM 的区别是硬件平台不同（GPU vs CPU）。
• PE 是处理单元，每个 GPU 通常是一个 PE。
• 对称内存 是所有 PE 上布局相同的内存区域，通过 nvshmem_malloc 分配。
• 对称堆 是对称内存的集合，在初始化时指定大小。

下一章咱们会深入 NVSHMEM 的初始化流程和基本 API。不过在那之前，我建议你先把这几个概念消化掉。尤其是对称内存，它是后面所有内容的基础。