3. 第一个 NVSHMEM 程序：初始化与终止

好，咱们直接动手写代码。这一章的目标很简单——让你跑通第一个 NVSHMEM 程序。说白了，就是学会怎么让 GPU 上的 PE（处理单元）互相认识，然后打个招呼。

3.1 为什么需要初始化与终止？

NVSHMEM 跟 MPI 有点像。你想想看，多个 GPU 要协同工作，总得有个「握手」和「告别」的过程吧？nvshmem_init() 就是那个握手动作，它负责建立 PE 之间的通信上下文。而 nvshmem_finalize() 就是告别，释放资源。

我个人习惯在程序一开始就调用 nvshmem_init()，在最后调用 nvshmem_finalize()。中间才是咱们的业务逻辑。嗯，这里要注意：这两个函数必须成对出现，缺一个都不行。

3.2 获取 PE 数量和 ID

每个 PE 都需要知道两件事：

• 总共有多少个 PE？ —— 用 nvshmem_n_pes()
• 我是第几个？ —— 用 nvshmem_my_pe()

这两个函数返回的都是 int 类型。PE 的编号从 0 开始，到 n_pes - 1 结束。我在项目中遇到过一个问题：新手容易把 PE 编号和 GPU 索引搞混。其实它们不是一回事。PE 是逻辑概念，一个 GPU 上可以跑多个 PE，但通常我们让一个 PE 对应一个 GPU。

3.3 Hello World 示例

好了，直接上代码。这是最经典的 NVSHMEM Hello World：

#include <stdio.h>
#include <nvshmem.h>

int main(void) {
    // 1. 初始化
    nvshmem_init();

    // 2. 获取 PE 信息
    int my_pe = nvshmem_my_pe();
    int n_pes = nvshmem_n_pes();

    // 3. 每个 PE 打印自己的信息
    printf("Hello from PE %d of %d\n", my_pe, n_pes);

    // 4. 终止
    nvshmem_finalize();
    return 0;
}

这段代码很简单，对吧？但有几个细节我想强调一下：

• 头文件：必须包含 nvshmem.h。少了它，编译器会报错。
• 编译：不能用普通的 nvcc 直接编译。你需要用 NVSHMEM 提供的编译器包装器，比如 nvshmemcc。具体命令我后面会讲。
• 运行：需要用 nvshmrun 来启动，比如 nvshmrun -n 4 ./hello 表示启动 4 个 PE。

3.4 编译与运行

咱们来走一遍完整的流程。假设你的环境已经装好了 NVSHMEM（没装好的话，先去装）。

编译命令：

nvshmemcc -o hello hello.c

运行命令：

nvshmrun -n 4 ./hello

输出应该类似这样（顺序可能不同）：

Hello from PE 0 of 4
Hello from PE 1 of 4
Hello from PE 2 of 4
Hello from PE 3 of 4

你看，每个 PE 都独立打印了自己的信息。这就是 NVSHMEM 的基本工作方式——多个 PE 并行执行同一份代码，但通过 my_pe 来区分各自的身份。

3.5 知识体系图

下面这张图展示了本章的核心逻辑：

这张图把咱们这一章的内容串起来了。从初始化开始，到获取信息，再到打印输出，最后终止。每一步都有对应的函数，缺一不可。

3.6 常见问题与避坑

我总结几个新手容易踩的坑：

• 忘记调用 nvshmem_init()：直接调用 nvshmem_my_pe() 会返回 0，但这是未定义行为。我曾经见过有人因为这个原因，在单卡上调试没问题，一上多卡就崩。
• 在 nvshmem_finalize() 之后调用 NVSHMEM 函数：这会导致段错误。记住，finalize 之后 NVSHMEM 就「下班」了。
• PE 编号不是 GPU 编号：虽然默认情况下 PE 0 对应 GPU 0，但你可以通过环境变量重新映射。不要硬编码假设。

好了，第一个 NVSHMEM 程序就到这里。代码不多，但它是你后续所有并行编程的基础。把这个流程刻在脑子里：init → 干活 → finalize。

---