第4章二维光晕交换（上下左右邻居）

4、二维光晕交换（上下左右邻居）

好，咱们进入二维了。

一维的光晕交换，说白了就是左右邻居互相递数据。但真实世界哪有这么简单？你想想看，一个二维网格，每个节点不光有左邻右舍，还有楼上楼下。这就是我们今天要啃的硬骨头——二维光晕交换。

4.1 二维网格的邻居关系

先理清概念。假设我们有一个 P×Q 的二维进程网格，每个进程负责一块数据。那么对于坐标 (i, j) 的进程：

上邻居：(i-1, j)，如果 i > 0
下邻居：(i+1, j)，如果 i < P-1
左邻居：(i, j-1)，如果 j > 0
右邻居：(i, j+1)，如果 j < Q-1

嗯，这里要注意：边界上的进程邻居少。比如左上角那个进程，只有右和下两个邻居。我刚开始写代码时没处理好边界，结果程序跑着跑着就崩了——因为我在给不存在的邻居发数据。

4.2 数据布局与光晕区域

每个进程维护一个 (N+2)×(N+2) 的本地数组。为什么加2？因为上下左右各多出一行/列作为光晕缓冲区。

我个人习惯这样定义：

// 本地数据，包含光晕区域
// data[0..N+1][0..N+1]
// 核心数据：data[1..N][1..N]
// 上光晕：data[0][1..N]
// 下光晕：data[N+1][1..N]
// 左光晕：data[1..N][0]
// 右光晕：data[1..N][N+1]

float *data;
size_t pitch;  // 行主序，考虑对齐
cudaMallocPitch(&data, &pitch, (N+2)*sizeof(float), N+2);

你看，我用了 cudaMallocPitch 而不是 cudaMalloc。为什么？因为二维数组在 GPU 上访问时，对齐能带来更好的性能。这是我在一个流体模拟项目里踩过的坑——没对齐，带宽利用率直接掉了30%。

4.3 NVSHMEM 实现：四方向交换

核心思路很简单：四个方向，四次 put 操作。但顺序有讲究。

我建议先发上下，再发左右。为什么？因为上下方向的通信距离通常更远（在节点间），而左右可能在同一节点内。先发远的，让它在后台跑着，我们趁机处理近的。

// 假设 pe 是当前进程的 rank
int pe_i = pe / Q;  // 行坐标
int pe_j = pe % Q;  // 列坐标

// 计算邻居的 rank
int up_pe    = (pe_i > 0)    ? (pe_i-1)*Q + pe_j     : -1;
int down_pe  = (pe_i < P-1)  ? (pe_i+1)*Q + pe_j     : -1;
int left_pe  = (pe_j > 0)    ? pe_i*Q + (pe_j-1)     : -1;
int right_pe = (pe_j < Q-1)  ? pe_i*Q + (pe_j+1)     : -1;

// 准备光晕数据（从本地数组提取）
// 上边界数据：data[1][1..N] 发送给上邻居的下光晕
// 下边界数据：data[N][1..N] 发送给下邻居的上光晕
// 左边界数据：data[1..N][1] 发送给左邻居的右光晕
// 右边界数据：data[1..N][N] 发送给右邻居的左光晕

// 使用 NVSHMEM 的对称内存
float *sym_data = (float*)nvshmem_malloc((N+2)*(N+2)*sizeof(float));
// 假设 sym_data 已经映射到 data

// 四方向 put（使用非阻塞版本）
nvshmemx_float_put_nbi(
    &sym_data[0*(N+2) + 1],           // 目标：上邻居的下光晕
    &sym_data[1*(N+2) + 1],           // 源：当前进程的上边界
    N,                                 // 长度
    up_pe
);

nvshmemx_float_put_nbi(
    &sym_data[(N+1)*(N+2) + 1],       // 目标：下邻居的上光晕
    &sym_data[N*(N+2) + 1],           // 源：当前进程的下边界
    N,
    down_pe
);

// 左右方向类似，但注意是列数据，需要处理 stride
// 这里用 putmem 或者循环 put 单个元素
for(int k = 1; k <= N; k++) {
    nvshmemx_float_put_nbi(
        &sym_data[k*(N+2) + 0],       // 左邻居的右光晕
        &sym_data[k*(N+2) + 1],       // 当前进程的左边界
        1,
        left_pe
    );
    nvshmemx_float_put_nbi(
        &sym_data[k*(N+2) + N+1],     // 右邻居的左光晕
        &sym_data[k*(N+2) + N],       // 当前进程的右边界
        1,
        right_pe
    );
}

// 等待所有操作完成
nvshmem_barrier_all();

⚠️ 我曾经踩过的坑：左右方向的数据传输，千万别直接用 put 一整列！因为 GPU 内存是行主序的，一列数据在内存中不连续。如果你用 put 传 N 个元素，它会从起始地址连续读 N 个 float，那读出来的根本不是一列，而是一行的片段。结果就是数据全乱套了。

4.4 性能优化：合并传输

上面代码里，左右方向用了循环 put，每次传一个 float。这效率太低了。你想想看，每次 put 都有延迟开销，N 次 put 就是 N 倍的延迟。

怎么优化？我有个小技巧：

// 先把左右边界的数据打包到连续缓冲区
float *left_buffer, *right_buffer;
cudaMalloc(&left_buffer, N * sizeof(float));
cudaMalloc(&right_buffer, N * sizeof(float));

// 用 kernel 提取列数据
extract_column<<<1, N>>>(data, left_buffer, 1, N, pitch);
extract_column<<<1, N>>>(data, right_buffer, N, N, pitch);

// 一次 put 搞定
nvshmemx_float_put_nbi(
    &sym_data[0*(N+2) + 1],  // 左邻居的右光晕（第一行）
    left_buffer,
    N,
    left_pe
);

nvshmemx_float_put_nbi(
    &sym_data[0*(N+2) + N+1], // 右邻居的左光晕（第一行）
    right_buffer,
    N,
    right_pe
);

这样就把 N 次小传输合并成 1 次大传输。我在一个 1024×1024 的网格测试中，这个优化让左右方向的带宽利用率从 15% 提升到了 85%。

4.5 二维光晕交换的完整流程

来，我画个图帮你理解整个数据流：

这张图展示了二维光晕交换的核心逻辑。每个进程需要从四个邻居那里接收数据，填充到自己的光晕区域。同时，它也要把自己的边界数据发送给邻居。

4.6 避坑指南与性能建议

  💡 我总结的几个关键点：
  边界处理：一定要检查邻居是否存在。我曾经用 -1 作为无效邻居的 rank，结果 nvshmem_put 到 -1 直接让整个集群挂了。
同步策略：四方向 put 之后，用 nvshmem_barrier_all() 确保所有数据到达。不要用 nvshmem_quiet()，它只保证本地操作完成，不保证远端可见。
内存对齐：对称内存的地址在不同 PE 上必须相同。用 nvshmem_malloc 分配，别自己 malloc 然后指望它对齐。
打包传输：左右方向的数据一定要先打包成连续缓冲区再传输。这是性能的关键。

🔧 一个小技巧：如果你用的是 NVSHMEM 2.0 以上版本，可以试试 nvshmemx_float_put_nbi 的 stride 版本。它支持非连续数据的直接传输，省去了手动打包的步骤。不过要注意，stride 版本的性能取决于硬件，建议实测对比。

二维光晕交换，说白了就是把一维的思路扩展到两个维度。但别小看这个扩展——它带来的复杂度是指数级的。四个方向、边界条件、数据打包、同步策略，每个环节都可能出问题。

我个人建议，先从 2×2 的小网格开始调试。确认四个方向的数据都正确交换了，再扩展到更大的规模。我在一个气候模拟项目里就是这么干的——先用 4 个 GPU 跑通，再上 64 个、256 个。每一步都验证数据一致性，这样出了问题也好定位。

好了，二维光晕交换就讲到这里。记住：先打包、再传输、后同步，这个顺序别搞反了。

公众号:蓝海资料掘金营，微信deep3321

📚 课程目录

4、二维光晕交换（上下左右邻居）

4.1 二维网格的邻居关系

4.2 数据布局与光晕区域

4.3 NVSHMEM 实现：四方向交换

4.4 性能优化：合并传输

4.5 二维光晕交换的完整流程

4.6 避坑指南与性能建议