4. Hello World 与第一个程序：初始化与终止、PE查询、简单的PUT/GET操作、验证数据一致性
好，咱们直接上手写代码。
很多初学者学NVSHMEM，上来就被各种概念吓住了。其实没那么复杂。你想想看，它本质上就是一套让GPU之间能直接通信的API。今天我们就从最基础的Hello World开始，把整个流程跑通。
4.1 初始化与终止：别小看这两步
NVSHMEM的程序，开头和结尾是固定的套路。我个人习惯，写任何NVSHMEM程序，第一件事就是检查环境。
#include <nvshmem.h>
#include <stdio.h>
int main(void) {
int my_pe, n_pes;
// 初始化NVSHMEM运行时
nvshmem_init();
// 获取当前PE的ID和总PE数
my_pe = nvshmem_my_pe();
n_pes = nvshmem_n_pes();
printf("Hello from PE %d of %d\n", my_pe, n_pes);
// 终止NVSHMEM运行时
nvshmem_finalize();
return 0;
}这段代码看着简单，但有个坑。我曾经在项目里遇到过，有人把nvshmem_init()写在了CUDA设备初始化之后，结果程序跑起来直接段错误。为什么？因为NVSHMEM的初始化会接管一部分设备资源管理，顺序错了就会冲突。
nvshmem_init()必须在任何CUDA API调用之前执行。这是硬性规定，没得商量。
4.2 PE查询：你是谁？你在哪？
PE（Processing Element）是NVSHMEM里的基本执行单元。每个PE都有一个唯一的ID，从0到N-1。说白了，这就是GPU的“身份证号”。
我刚开始用的时候，总觉得查PE号是多此一举。后来做多卡通信调试，才发现这玩意儿太重要了。没有它，你根本不知道数据是从哪张卡发出来的。
| 函数 | 作用 | 返回值 |
|---|---|---|
nvshmem_my_pe() |
获取当前PE的ID | int (0 ~ N-1) |
nvshmem_n_pes() |
获取总PE数量 | int |
这里有个小技巧：nvshmem_n_pes()返回的是你启动程序时指定的PE总数。如果你用4张卡启动，它就返回4。如果你只用了1张卡，它就返回1。嗯，这个逻辑很直白。
4.3 简单的PUT/GET操作：数据搬家
PUT和GET是NVSHMEM最核心的操作。PUT是把数据从当前PE发到目标PE，GET是从目标PE把数据拉回来。你可以把它想象成快递：PUT是寄件，GET是取件。
来看一个PUT的例子：
#include <nvshmem.h>
#include <stdio.h>
int main(void) {
int my_pe, n_pes;
int *buf;
int source = 42;
int dest = 0;
nvshmem_init();
my_pe = nvshmem_my_pe();
n_pes = nvshmem_n_pes();
// 每个PE分配对称内存
buf = (int*)nvshmem_malloc(sizeof(int));
if (buf == NULL) {
printf("PE %d: malloc failed\n", my_pe);
nvshmem_finalize();
return -1;
}
// 初始化本地缓冲区
*buf = 0;
// PE 0 把数据发送给 PE 1
if (my_pe == 0) {
nvshmem_int_put(buf, &source, 1, 1); // 把source的值put到PE1的buf
printf("PE %d: PUT value %d to PE 1\n", my_pe, source);
}
// 等待所有PUT操作完成
nvshmem_barrier_all();
// PE 1 打印接收到的数据
if (my_pe == 1) {
printf("PE %d: Received value %d\n", my_pe, *buf);
}
nvshmem_free(buf);
nvshmem_finalize();
return 0;
}这段代码里，nvshmem_int_put的参数分别是：目标地址、源地址、元素个数、目标PE号。注意，目标地址是目标PE上的对称地址，不是本地地址。这个区别很重要，我刚开始就搞混过。
nvshmem_barrier_all()或者nvshmem_quiet()。
GET操作类似，只是方向相反：
// PE 1 从 PE 0 拉取数据
if (my_pe == 1) {
nvshmem_int_get(buf, &source, 1, 0); // 从PE0的source拉取数据到本地buf
printf("PE %d: GET value %d from PE 0\n", my_pe, *buf);
}4.4 验证数据一致性：别信直觉，信代码
数据一致性是分布式编程里最头疼的问题。你想想看，数据从PE 0发到PE 1，中间经过PCIe总线、NVLink、甚至网络。任何一个环节出问题，数据就可能对不上。
我常用的验证方法很简单：每个PE先把自己的数据发出去，然后所有PE互相检查。就像开会点名，每个人报数，最后核对总数。
// 验证数据一致性的简单方法
int errors = 0;
// PE 0 发送数据给所有其他PE
if (my_pe == 0) {
for (int i = 1; i < n_pes; i++) {
nvshmem_int_put(buf, &source, 1, i);
}
}
nvshmem_barrier_all();
// 每个PE检查收到的数据是否正确
if (my_pe != 0) {
if (*buf != source) {
printf("PE %d: Data mismatch! Expected %d, got %d\n",
my_pe, source, *buf);
errors++;
}
}
// 汇总错误信息
nvshmem_barrier_all();
if (my_pe == 0) {
printf("Total errors: %d\n", errors);
}这段代码里，我用nvshmem_barrier_all()做了两次同步。第一次确保所有PUT操作完成，第二次确保所有检查完成。少一个同步，结果就可能不准。
4.5 完整示例：Hello World 升级版
最后，给你一个完整的Hello World程序。它包含了初始化、PE查询、PUT/GET操作和数据一致性验证。你可以直接编译运行：
// hello_nvshmem.c
// 编译: nvcc -o hello_nvshmem hello_nvshmem.c -lnvshmem
#include <nvshmem.h>
#include <stdio.h>
int main(void) {
int my_pe, n_pes;
int *buf;
int send_val, recv_val;
int errors = 0;
nvshmem_init();
my_pe = nvshmem_my_pe();
n_pes = nvshmem_n_pes();
buf = (int*)nvshmem_malloc(sizeof(int));
*buf = 0;
// 每个PE准备自己的数据
send_val = my_pe * 100;
// PE 0 收集所有数据
if (my_pe == 0) {
for (int i = 1; i < n_pes; i++) {
nvshmem_int_put(buf, &send_val, 1, i);
}
}
nvshmem_barrier_all();
// 验证数据
if (my_pe != 0) {
recv_val = *buf;
if (recv_val != 0) { // 期望从PE0收到0
printf("PE %d: Error! Expected 0, got %d\n", my_pe, recv_val);
errors++;
}
}
nvshmem_barrier_all();
if (my_pe == 0) {
printf("Hello World from NVSHMEM!\n");
printf("Total PEs: %d, Errors: %d\n", n_pes, errors);
}
nvshmem_free(buf);
nvshmem_finalize();
return 0;
}运行这个程序，你会看到每个PE打印自己的ID，然后PE 0汇总结果。如果一切正常，errors应该是0。
nvshmem_barrier_all()没配对。记住，所有PE必须执行相同次数的同步操作，多一个少一个都不行。
好了，第一个程序就到这里。代码不多，但把NVSHMEM的核心流程都走了一遍。你可以在自己的机器上试试，改改数据，加加PE，感受一下多卡通信的威力。