核心要义NVLS：
< 1.> 每个参与集合通信贡献一块物理显存，登记进本地，也登记进交换机（nvswitch）；
< 2.> 实际上发起集合通信的主角，是交换机（nvswitch）。
当要做 sum 集合通信时，实际上是某一个 gpu 发出了一条特别的 gpu 指令，发送到 交换机；
交换机解析这个信息报后就知道要把众 gpu 的登记进来的显存的对应位置的数据一并 ld 到交换机本地（通过交换结构）；
然后使用本地的cpu中的算力 sum 起来。最后将 sum 结果广播送回各个 GPU。
1. NVLS 的在网计算具体计算情景
NVLS（NVLink SHARP）是第三代 NVSwitch（Hopper/NVLink4）起在交换机 ASIC 里内置的归约与组播引擎，把原来由 GPU SM 执行的"加法/归约"计算下沉到网络里。它支持的具体计算情景可以分几层来看：
一、交换机硬件本身能算的两类操作
- 在网归约（in-network reduction）：多个 GPU 把数据发到 NVSwitch，数据在交换机端口汇合时直接完成算术运算（如求和），聚合结果再下发。这一点和 IB 网络上的 SHARP 思路一致，只是把计算点从 IB 交换机搬到了 NVSwitch。
- 硬件组播（multicast）：交换机把一份数据包复制分发给所有订阅的 GPU，用于广播/聚合阶段。
触发方式是 GPU 侧发出 multimem.ld_reduce（拉取归约）和 multimem.st（推送组播）这类内存语义的 PTX 指令，由 GPU 驱动、NVSwitch 执行。
二、对应到集合通信（NCCL）的计算情景
- AllReduce ：NVLS 最核心、也是最主要的场景。被分解为 reduce-scatter + all-gather 两阶段：第一阶段各 GPU 用
multimem.ld_reduce拉取所有参与者的数据、在交换机端口完成归约；第二阶段用multimem.st把归约结果推给交换机，由交换机组播给所有 GPU。典型受益者是大模型训练里的数据并行梯度同步 和张量并行每层 attention/MLP 之后的 all-reduce------在网归约既省 NVLink 带宽，又把 GPU SM 从通信计算中解放出来，方便通信与计算 overlap。 - ReduceScatter 单独加速：NCCL 中 NVLS 除 AllReduce 外还支持 ReduceScatter 和 AllGather------这正好对应 FSDP/DeepSpeed ZeRO 的梯度分片同步模式。
- AllGather / Broadcast 走组播：不需要算术计算，用硬件 multicast 一次发送、多点接收，省掉环形转发。
- 机内 + 机间的混合情景 ：NCCL 的
NVLS算法是"机内 NVLink SHARP 归约 + 机间 CollNet（IB SHARP）"，NVLSTree则是"机内 NVLS + 机间树形 fan-out"。 - 用户自定义 kernel：通过 NCCL Device API 的 multimem 语义、CUDA multicast 对象、PyTorch SymmetricMemory 或 NVSHMEM，开发者可以在自己的 CUDA/Triton kernel 里直接调用 NVLS 的归约/组播内存语义，而不必走标准 NCCL 集合调用。
三、支持的运算与数据类型
交换机做的是结合性/交换性的逐元素归约运算 ：主要是 sum、min、max（SHARP 家族还支持 and/or/xor 等位运算）。数据类型覆盖 FP32、FP16、BF16、FP64、INT8/32/64 等；第五代 NVSwitch（GB300/NVL72）的 SHARP-v3 进一步扩展到 FP8、FP16 等低精度归约。另一个工程上的好处：归约顺序由交换机固定，浮点 all-reduce 结果确定性更强。
四、不适合的情景
NVLS 只对"可分解为归约 + 广播"的模式有效。**点对点通信、all-to-all（如 MoE expert 并行的 dispatch/combine）**无法映射到交换机归约/组播上，NVLS 帮不上忙。另外它只能在同一 NVLink 域（NVSwitch clique）内工作，跨节点仍需叠加 CollNet/Tree；小消息上 NVLS 也可能因为建组开销而不如 Ring。
简单总结：NVLS 的"在网计算"本质上就两种原语------归约（sum/min/max 等）+ 组播 ，落到业务上就是 AllReduce、ReduceScatter、AllGather/Broadcast 这几类集合通信的加速，最典型的计算情景是大模型训练/推理中的梯度同步（DP/FSDP）和张量并行逐层 all-reduce。
2. NVLS 通信过程解析
下面以一台 8×H100（NVLink4 + 第三代 NVSwitch 域内）的机器、对一份 FP32 梯度做 ncclAllReduce(op=ncclSum) 为例，把 NVLS 的全过程拆开。一次 AllReduce 从逻辑上就是 Reduce-Scatter（交换机上做 sum）+ AllGather（交换机上做组播） ，NCCL 的 NVLS 算法正是用两条 multimem 指令分别对应这两段：multimem.ld_reduce 负责在网归约，multimem.st 负责在网组播。
阶段 0：一次性建组（初始化时发生，不在每次通信里）
NVLS 的前提是有一块所有 GPU 都映射的组播内存对象（Multicast Object, MC），建立过程大致是：
- NCCL 初始化时探测拓扑，确认参与通信的所有 rank 都在同一个 NVLink 域（NVSwitch clique）内。
- 本机 localRank 0 调用
cuMulticastCreate创建 MC 对象------它本身不含物理内存，只是 NVSwitch 上登记的一个"成员集合"；随后把句柄导出成 shareable handle，通过 bootstrap 广播给本机其他 rank，其他 rank 导入得到自己的 mcHandle。 - 每个 rank 调
cuMulticastAddDevice把自己的 GPU 加进组播组。 - 每个 GPU 在本地分配一块普通物理显存（UC buffer），再通过
cuMulticastBindMem把它绑定为 MC 对象在同一偏移上的后备存储------于是同一个 MC 地址在每个成员 GPU 上各有一份物理副本。 - 每个进程把 MC 句柄映射进自己的虚拟地址空间，得到
mcptr。此后 GPU 对mcptr的普通读，读的是自己本地那份副本 ；而multimem.ld_reduce/multimem.st这两条特殊指令，才会触发交换机跨副本的归约/组播语义。
这一步是理解后面的关键：MC 地址 = 8 份分散在各 GPU 上的物理副本 + NVSwitch 上的一个归约/组播引擎。
阶段 1：通信发起与算法选择
应用调用 ncclAllReduce(sendbuff, recvbuff, count, ncclFloat32, ncclSum, comm, stream)。NCCL 检查：消息在 NVLink 域内、NVLS 已启用、消息尺寸超过阈值，于是选中 NVLS 算法，启动一个只占用很少 CTA 的通信 kernel（这也是为什么 NVLS 对 SM 的占用远低于 Ring------加法不再由 GPU 线程做）。
阶段 2：数据分片与"属主"划分
设梯度总大小为 m。NCCL 把它切成 8 个 slice，slice i 的归约属主是 rank i；每个 slice 再细分成若干 chunk，在多个 channel 上流水推进。分片属主制的作用和 Ring 一样：每块数据只需要一个 GPU 拿到最终归约结果，再广播出去即可。
以一个小例子说明：4 个 GPU（GPU0--3），向量 X = [x0, x1, x2, x3] 分 4 片，每个 GPU 上有一份本地值 X^(r)，目标是每个 GPU 都得到 S = X^(0)+X^(1)+X^(2)+X^(3)。
阶段 3：Reduce-Scatter ------ sum 在交换机里完成
对每个 chunk（比如 slice 0，属主 GPU0）：
- 写入本地副本：8 个 GPU 各自把 slice 0 对应的本地数据用普通 store 写进自己那份 UC 物理页（即本地 MC 副本的对应偏移）。
- 属主发起归约读 ：属主 GPU0 对该 MC 地址发出
multimem.ld_reduce.add.f32。这条指令不是普通读------它被路由到 NVSwitch，交换机据此从全部 8 个成员 GPU 的本地副本上读取同一偏移的数据 ，各路响应在交换机内部汇聚时，由 switch ASIC 里的归约引擎直接完成逐元素sum，把归约结果返回给 GPU0。也就是说 GPU0 发出 m/8（分 chunk 后每次一小段）的请求、收回的是已经加完的和，正如 NVLS 的数据流描述：GPU 提交multimem.ld_reduce后"发送 1/4m 数据给交换机，并从交换机拿回 1/4m 已归约数据"（以 4 卡为例）。这也是 HotChips/GTC 资料里 NVLS AllReduce 第一阶段的含义：各 GPU 把 chunk 发给 NVSwitch，归约后的结果散回对应属主。 - 结果写回 GPU0 本地 buffer 的 slice 0 区域。此刻只有 GPU0 持有 slice 0 的完整和
S[0]------这正是 Reduce-Scatter 的语义。
8 个 slice 由 8 个属主并行/流水地各自发起 ld_reduce，归约负载天然分散在所有属主上，NVLink 各端口流量均衡。顺带一提：交换机的归约顺序固定，所以浮点 sum 的逐位结果具有确定性。
阶段 4：AllGather ------ 组播由交换机完成
属主 GPU0 拿到 S[0] 后：
- 属主发起组播写 ：GPU0 把
S[0]用multimem.st写到 slice 0 对应的 MC 地址。这条写被 NVSwitch 复制，一次性落入全部 8 个成员 GPU 本地副本的同一偏移------等价于"发一份、交换机扇出 8 份"。 - 其他 GPU 随后从本地副本把数据读到
recvbuff（若通信 buffer 与用户 buffer 已注册合一，则直接就是最终结果，省掉拷贝）。
所有 slice 的广播同样流水进行。阶段 3 和阶段 4 在 chunk 粒度上交叠：slice 0 在广播时，slice 1 可能正在归约。
阶段 5：同步与收尾
- chunk 级同步靠 buffer 内的 flag/到达计数完成（确保属主发起
ld_reduce前所有成员已写完本地副本）。 - 全部 chunk 完成后 kernel 退出，stream 上的 AllReduce 完成。若上层是 FSDP/DDP，下一步就是 optimizer 或下一层的计算。
流量账与关键特征
设每 GPU 数据量 m、N=8：归约阶段每个 GPU 的数据被读取约 (N-1)/N·m，广播阶段写回约 (N-1)/N·m，线上总字节数与 Ring 相同（≈2m）------NVLS 的赢点不在字节数，而在于：
- 加法从 SM 挪进交换机 ：Ring 里每个中间 GPU 要用 SM 做
recvReduceSend，NVLS 下 GPU 只做内存读写，SM 释放给计算，便于通信/计算 overlap； - 扇出/扇入在交换机一跳完成：没有逐跳转发，延迟和中间节点开销更低；
- 归约顺序固定，结果确定。
也正因为它本质上是"归约 + 组播"两个原语，NVLS 只对 AllReduce/ReduceScatter/AllGather 这类模式有效------点对点、All-to-All（MoE dispatch）映射不上去。另外要注意 NVLS 对 ReduceScatter 的支持曾有版本差异：早期 NCCL 版本会拒绝 NVLS 的 reduce-scatter（报 "no algorithm/protocol available"），需要较新版本或走 AllReduce 路径；而在 NCCL 的算法表里，NVLS 主要面向 AllReduce，并借硬件能力附带支持 ReduceScatter 与 AllGather。
最后补一句机间扩展：跨节点时 NCCL 用 NVLS（机内 NVLS + 机间 CollNet/IB SHARP）或 NVLSTree（机内 NVLS + 机间树）组合，机内这一段的过程与上面完全一致。
3. cu drv 中 NVLS 相关 API
这组函数（cuMulticastCreate / AddDevice / BindMem / BindAddr / Unbind / GetGranularity，CUDA 12.x 起随 Hopper/NVSwitch 引入）本质上是一套把 NVSwitch 的硬件组播/归约能力暴露给软件的"内存对象管理"API 。它不是发明了一种新的通信原语，而是给 multimem.ld_reduce / multimem.st 这两条 PTX 指令提供一个合法的、可寻址的"操作对象"。支持与否可以用 CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED 查询------它要求 NVSwitch 交换机组网，没有 NVLink fabric 的卡（如 RTX 4090）报不支持，驱动再新也没用。
一、为什么需要这套 API：引入目的
要理解设计动机，先看硬件侧的事实：NVSwitch 里的归约引擎和组播扇出逻辑，需要知道三件事------哪些 GPU 是一个组（成员表）、每个成员的数据放在哪（地址映射表）、逻辑地址空间有多大 。这些状态以前完全藏在硬件/驱动内部，软件无法配置；而没有软件可配置的对象，multimem 指令就不知道该对"谁的哪些副本"做归约。
于是 CUDA 借鉴了已有的 VMM（Virtual Memory Management）三段式模型 （cuMemCreate 物理句柄 / cuMemMap 虚拟映射 / 句柄导出共享），新增了一类特殊对象：multicast object（MC 对象） ------一个本身不含物理内存 的逻辑实体，代表"一个组播组 + 一段逻辑地址空间"；各成员 GPU 再把自己的物理显存"挂载"到这段逻辑空间上，形成同一偏移、多份物理副本的结构。官方的定义就是：cuMulticastCreate 创建的对象"enables certain memory operations to be broadcast to a team of devices"，设备通过 AddDevice 加入，内存通过 BindMem/BindAddr 绑定，最后像普通 VMM 分配一样用 cuMemMap 映射进各进程的虚拟地址空间。
这个抽象带来一个关键的运行时分叉：GPU 对映射后的 MC 地址做普通 ld/st，访问的是自己本地那份副本；只有用 multimem.* 指令访问，地址才被识别为组播地址、路由到 NVSwitch 触发归约或组播 。实测语义对照就是：写多播地址=广播到所有设备，归约=触发硬件加速（multimem.red.global.add.s32 / multimem.ld.reduce.global.add.s32）。
二、逐函数解析：目标与背后原理
1. cuMulticastGetGranularity ------ 先问硬件"对齐单位是多少"
- 目标 ：查询给定属性下，MC 对象的
size、绑定偏移（mcOffset/memOffset）和映射地址必须是多少的整数倍；分CU_MULTICAST_GRANULARITY_MINIMUM（正确性下限）和RECOMMENDED（性能最优）两档。 - 原理 ：交换机的组播地址表和归约引擎的跟踪粒度是粗粒度的（页面级而非字节级），驱动把硬件约束向上透传为一个粒度参数，让应用在分配前就对齐，避免绑定失败或触发低效路径。所以正确顺序是先查粒度、再定 size。
2. cuMulticastCreate ------ 创建"逻辑实体"，向交换机登记组播组
- 目标 ：按
CUmulticastObjectProp（numDevices组大小、size逻辑空间大小、handleTypes共享句柄类型）创建一个 MC 对象，返回mcHandle。size限定了每个设备可绑定的内存总量，且必须是 MINIMUM 粒度的倍数。 - 原理 ：这一步在驱动/NVSwitch 侧分配了一个组播组的记录 （将来的成员端口掩码、地址转换表项都挂在这个对象上），并划定一段逻辑偏移空间
[0, size)。注意它不分配任何显存 ------这就是为什么后面每个 GPU 还要各自cuMemCreate分配物理内存再绑定。handleTypes决定它能否跨进程共享（同机 POSIX fd / 跨节点 NVLink 域用 FABRIC），共享走cuMemExportToShareableHandle/cuMemImportFromShareableHandle，销毁复用cuMemRelease------完全复用 VMM 的句柄机制。
3. cuMulticastAddDevice ------ 把 GPU 注册进组播组
- 目标 ：把一个 GPU 与 MC 对象关联。关联在对象生命周期内永久有效 ；且必须凑满
numDevices个成员后，任何Bind/cuMemMap才能返回（否则会阻塞）。 - 原理 ：交换机的组播扇出和归约归并都要知道完整的成员端口集合。归约尤其依赖成员固定------只有成员集确定，归约顺序才固定（这正是 NVLS 结果确定性的来源之一），所以 API 把"组未配齐就使用"直接做成阻塞，从机制上杜绝半成品组参与计算。
4. cuMulticastBindMem ------ 建立"同一逻辑偏移 → 各 GPU 物理副本"的映射
- 目标 ：把某个成员 GPU 上用
cuMemCreate分配的物理内存（memHandle+memOffset）绑定到 MC 对象的mcOffset处。约束：该 GPU 必须已 AddDevice；size/mcOffset/memOffset满足粒度对齐；可共享/导入的 MC 对象只能绑定可共享内存。 - 原理 ：这是整个机制的核心一步 。每个成员 GPU 都执行一次 Bind，把自己那块物理页登记为 MC 地址空间在
mcOffset处的"本地副本"。此后当某 GPU 对 MC 地址发出multimem.ld_reduce，交换机就能按这张映射表，从全部成员的对应偏移 并行拉取数据做归约；发出multimem.st时，则把数据扇出写进所有成员的对应偏移。换句话说，Bind 配置的就是交换机归约/组播引擎的地址翻译表。
5. cuMulticastBindAddr ------ 用"虚拟地址"版本的绑定
- 目标 ：与 BindMem 相同，但以已映射的虚拟地址
memptr指定内存；内存可来自cuMemCreate或cudaMallocAsync（流序分配器）。 - 原理：给上层更灵活的接入方式------比如想直接把 memory pool 里已分配好、已映射 VA 的 buffer 挂进组播组，而不必再走一遍句柄路线。底层做的事与 BindMem 一样，都是登记副本映射，只是寻址输入从"物理句柄+偏移"换成"VA"。
6. cuMulticastUnbind ------ 解绑
- 目标 ：解除某设备在
mcOffset处的副本绑定。mcOffset/size必须满足粒度，且必须与当初 Bind 时的取值精确一致，否则行为未定义。 - 原理 ：从交换机的映射表中摘除对应表项，物理内存本身归还应用管理（MC 对象和其逻辑空间仍在，直到
cuMemRelease）。
三、串起来看：NCCL 里的真实调用链
NCCL 的 NVLS 初始化（nvlsGroupCreate / nvlsGroupBindMem）就是这套 API 的标准用法，可以对照前面的流程看：
- rank0
cuMulticastCreate建组，把 handle 作为 shareable handle 通过 bootstrap 广播，其他 rank 导入拿到mcHandle； - 各 rank
cuMulticastAddDevice把本卡加入组； - 各 rank 本地
cuMemCreate（CU_MEM_ALLOCATION_TYPE_PINNED+ DEVICE）分配物理显存作为 UC buffer，cuMemMap+cuMemSetAccess映射成本进程 VA； - 最后
cuMulticastBindMem(mcHandle, 0, ucHandle, 0, size, 0)把 UC buffer 绑成 MC 偏移 0 处的副本。
之后再 cuMemAddressReserve + cuMemMap 把 MC 对象本身 映射进每个进程的地址空间得到 mcptr，通信 kernel 就对 mcptr 使用 multimem.ld_reduce / multimem.st 发起在网计算------也就是上一问里 AllReduce 的归约与组播两个阶段。
总结 ：Create 建组（交换机侧登记）、AddDevice 定员（成员端口掩码）、BindMem/BindAddr 落图（逻辑偏移→各 GPU 物理副本的映射表）、GetGranularity 守对齐（硬件表项粒度的约束透传）、Unbind/cuMemRelease 拆图销组。这套 API 的全部目的，就是把 NVSwitch 里原本不可编程的组播/归约状态，变成一组可显式配置、可跨进程共享的内存对象，从而让上层（NCCL、NVSHMEM、自定义 kernel）能用两条内存指令驱动在网计算。
4. sample multi_node_p2p 项目分析
项目地址：https://github.com/NVIDIA/multi-gpu-programming-models/tree/master/multi_node_p2p
同目录注释版文件：
jacobi.cpp、jacobi_kernels.cu、Makefile（代码未改动，仅加中文注释）
一、项目概述
NVIDIA "multi-gpu-programming-models" 系列示例中的一员：二维五点差分 Jacobi 迭代求解器（带正弦 Dirichlet 边界、L2 范数收敛判据）。同系列的其它版本分别用 MPI、NCCL、NVSHMEM 实现 GPU 间通信；multi_node_p2p 是最新、也最"底层"的版本------它不用任何通信库传输 GPU 数据，而是直接用 CUDA Driver API 在 MNNVL（Multi-Node NVLink，多节点 NVSwitch fabric）环境下实现两种机制：
| 机制 | 解决的问题 | 使用的 API | 对应文件位置 |
|---|---|---|---|
| Fabric 内存 + 对端映射（P2P 直写） | halo（边界行）交换 | cuMemCreate + CU_MEM_HANDLE_TYPE_FABRIC、cuMemExportToShareableHandle、MPI 交换句柄、cuMemImportFromShareableHandle、cuMemMap/cuMemSetAccess |
jacobi.cpp 的 allocate_fabric_mem / map_peers |
Multicast(MC) 对象 + multimem 指令 |
L2 范数的全员求和（Allreduce）+ 全员栅栏 | cuMulticastGetGranularity / Create / AddDevice / BindMem / Unbind，kernel 内 multimem.red / multimem.ld_reduce 内联 PTX |
jacobi.cpp 主函数 MC 段 + jacobi_kernels.cu 的 all_reduce_norm_barrier_kernel |
MPI 在本程序中只承担"带外控制面"：句柄交换、广播、栅栏，以及未开 -use_mc_red 时的标量 Allreduce。
二、运行流程总览
main()
├─ MPI_Init / 命令行解析 / local_rank 绑卡
├─ 前置检查: CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED
│ （-use_mc_red 时另查 MULTICAST_SUPPORTED，禁止 MPS 超订）
├─ single_gpu() 单 GPU 参考解 + 基线耗时
├─ 行向域分解（ny-2 行均分，高低块均衡）
├─ allocate_fabric_mem(a / a_new) VMM 五步分配可共享显存
├─ MPI 交换 fabric 句柄 → 导入 → map_peers()
│ 此后 a_new_top / a_new_bottom 即为邻居显存的本地视图
├─ [-use_mc_red] 建 MC 对象：
│ GetGranularity → rank0 Create → Bcast 句柄 → 导入
│ → AddDevice → Barrier → cuMemCreate(UC) → BindMem
│ → 映射 MC 视图与 UC 视图
├─ 主循环 while (l2_norm > tol && iter < iter_max):
│ 清零范数累加器
│ launch_jacobi_p2p_kernel() 计算 + 直写邻居 halo 行
│ 若本轮查范数:
│ MC 模式: launch_all_reduce_norm_barrier_kernel()
│ （multimem.red 栅栏 + multimem.ld_reduce 在网求和）
│ 传统模式: D2H 拷贝 + MPI_Allreduce
│ MPI_Barrier（防快 rank 覆盖慢 rank 的 halo / 保护 MC 读时序）
│ swap 双缓冲
├─ 拷回结果与参考解逐点比对（MPI_Allreduce 汇总正确性）
├─ rank0 输出加速比 / 并行效率
└─ 资源销毁（MC: Unbind/Unmap/Release；P2P: unmap_peers）三、机制 A：Fabric 内存与 halo 直写（逐点分析）
- 为什么用
CU_MEM_HANDLE_TYPE_FABRIC而不是 POSIX fd ：POSIX 句柄只能经 Unix domain socket 在本机进程间传递；本示例要用 MPI 在任意节点间 交换句柄，只有 FABRIC 类型编码了 NVLink fabric 内的全局可路由标识（jacobi.cpp注释中明说了这一点）。 - 分配侧五步法 ：
cuMemGetAllocationGranularity→cuMemCreate(PINNED, DEVICE)→cuMemExportToShareableHandle→cuMemAddressReserve→cuMemMap→cuMemSetAccess。这是 CUDA VMM 的标准动作，与cuMulticast*的映射阶段同源。 - 为什么所有 rank 按
chunk_size_high分配 ：cuMemMap映射对端内存时 size 必须一致，而各 rank 分到的行数可能差 1 行，统一按最大块分配保证映射合法。 - halo 交换的"push 模型" ：
jacobi_p2p_kernel中，算到本 rank 首行时顺手a_new_top[top_iy*nx+ix] = new_val（写入上方邻居的底部 halo），算到末行时写下方邻居的顶部 halo（行 0）。这是一次普通的全局 store，硬件经 NVLink fabric 直达邻居节点显存------无 memcpy、无通信库、无 staging buffer，且与计算融合在同一个 kernel 里。环形邻居关系（rank0 的 top 是最后一名）天然实现周期边界。 - 代价/同步 ：push 模型下没有"接收方确认"，所以每轮迭代后必须
MPI_Barrier，防止快 rank 把邻居还在读的 halo 行提前覆盖。这是用主机栅栏换通信简易性的取舍，也是示例性质代码可接受、生产代码需要更细粒度同步（如 NVSHMEM signal）的地方。
四、机制 B：MC 组播对象与在网归约（逐点分析）
对应 -use_mc_red 路径，用 8 字节的 real_int_pair { value, arrival_counter } 演示 NVSwitch 的在网计算能力：
- 建组 ：
cuMulticastCreate(numDevices=size, handleTypes=FABRIC)由 rank0 创建、广播 fabric 句柄、其余 rank 导入------跨节点共享因此必须 FABRIC 类型。 - 定员 ：每个 rank
cuMulticastAddDevice，随后MPI_Barrier确保组配齐（成员不满时 Bind/Map 会阻塞，这是 API 的硬约束）。 - 落图 ：各 rank
cuMemCreate一块本地物理显存（UC），cuMulticastBindMem(mc, 0, uc, 0, size)登记为 MC 偏移 0 的本地副本------这一步配置的就是交换机侧的"逻辑偏移 → 各 GPU 物理副本"映射表。 - 双视图映射 ：同一块物理内存映射两次------
mc_ptr（组播视图，供multimem.*使用）和uc_ptr（本地视图，供普通读写）。 - kernel 内三条指令 （
all_reduce_norm_barrier_kernel，单线程执行）：multimem.red.release.sys.global.add.u32 [mc->arrival_counter], 1------ 交换机把 +1 原子施加到所有成员副本 ，一次指令完成"系统级 release 栅栏 + 全员计数自增"（.release保证本 rank 先写入的 L2 部分和对全体可见）；fence.proxy.alias------ MC/UC 是同一块物理内存的两个虚拟别名，跨 proxy 访问需此 fence 定序；- UC 副本上
cuda::atomic_ref(system scope)自旋，等arrival_counter >= size*(iter+1)（计数器跨迭代累计、永不复位）； multimem.ld_reduce.relaxed.sys.global.add.f32 [mc->value]------ 一条 load，由 NVSwitch 对所有成员副本求和后返回，即 NVLS/SHARP 式在网归约；GPU 端开方即得全局 L2 范数。
- 保护性
MPI_Barrier：MC 归约只有加法语义、没有"读完成"时序保证，host 侧必须栅栏后才能允许下一轮清零value（代码注释明确说明）。 - 正确性细节 ：
expected_count用无符号计算，大迭代数 × 多 GPU 溢出时回绕比较仍良定义。
与 NCCL NVLS 的关系：NCCL 用同一硬件能力加速大规模 AllReduce/ReduceScatter；本示例展示的是同一原语在细粒度标量归约 + 栅栏场景的用法，且完全绕开通信库。
五、约束与适用环境
- 硬件：sm_90+（Hopper 起）且位于 NVSwitch fabric（MNNVL，如 GB200 NVL72 / 多节点 NVLink 域）；
CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED不满足时直接报错退出。 - 不允许 MPS 超额订阅（多进程共卡时 MC 语义不成立）。
- kernel 内
multimem代码用#if __CUDA_ARCH__ >= 900保护，低档架构编译出来该 kernel 为空操作------此时-use_mc_red的数值结果不可靠，属于示例的已知取舍。
六、总结
multi_node_p2p = "把 NVLink fabric 当作一块可全局寻址的内存"的教学示例：halo 交换 用 fabric 句柄 + VMM 映射实现零通信库的跨节点直写；收敛判据 用 MC 组播对象 + multimem.red/ld_reduce 让 NVSwitch 在网完成栅栏与求和。两者合起来，覆盖了 NVSwitch 在网计算的两种基本形态------数据面的 P2P 直访与控制面的在网归约/组播。
七、源码注释
7.1 cpp
cpp
/* ============================================================================
* jacobi.cpp ------ host 侧主程序（逐行中文注释版，原代码逐行保留、未做任何改动）
*
* MNNVL（多节点 NVLink fabric）环境下的多 GPU Jacobi 求解器。
* 不用通信库传 GPU 数据，使用两套 CUDA Driver API 机制：
* (A) Fabric 内存 + 对端映射 → halo 行 P2P 直写（kernel 内普通 store）
* (B) Multicast(MC) 对象 + multimem 指令 → L2 范数在网求和 + 全员栅栏
* （-use_mc_red 开启；否则退回 D2H 拷贝 + MPI_Allreduce）
* MPI 仅作控制面：句柄交换/广播、栅栏、以及非 MC 模式的标量 Allreduce。
* ==========================================================================*/
/* Copyright (c) 2017, 2024, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions
* are met:
* * Redistributions of source code must retain the above copyright
* notice, this list of conditions and the following disclaimer.
* * Redistributions in binary form must reproduce the above copyright
* notice, this list of conditions and the following disclaimer in the
* documentation and/or other materials provided with the distribution.
* * Neither the name of NVIDIA CORPORATION nor the names of its
* contributors may be used to endorse or promote products derived
* from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
* EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
* PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
* CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
* EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
* PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
* PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
* OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/
#include <algorithm> // std::find / std::min / std::swap
#include <cassert> // （本文件未直接使用，系列统一保留）
#include <cmath> // std::asin / std::sqrt / std::fabs
#include <cstdio> // printf / fprintf
#include <cstdlib> // exit
#include <iostream> // （istringstream 经由 <sstream>；此处系列统一保留）
#include <limits> // （本文件未直接使用，系列统一保留）
#include <sstream> // std::istringstream：命令行参数解析
#include <mpi.h> // MPI：控制面（句柄交换、栅栏、标量归约）
// ---- MPI 错误检查宏（多行续行结构，注释只能加在宏外） ----
// 逐行含义：执行 call 并保存状态码 → 非 MPI_SUCCESS 则把错误码翻译成可读串 →
// 打印调用文本(#call)、行号、文件、错误描述 → exit 终止。
#define MPI_CALL(call) \
{ \
int mpi_status = call; \
if (MPI_SUCCESS != mpi_status) { \
char mpi_error_string[MPI_MAX_ERROR_STRING]; \
int mpi_error_string_length = 0; \
MPI_Error_string(mpi_status, mpi_error_string, &mpi_error_string_length); \
if (NULL != mpi_error_string) \
fprintf(stderr, \
"ERROR: MPI call \"%s\" in line %d of file %s failed " \
"with %s " \
"(%d).\n", \
#call, __LINE__, __FILE__, mpi_error_string, mpi_status); \
else \
fprintf(stderr, \
"ERROR: MPI call \"%s\" in line %d of file %s failed " \
"with %d.\n", \
#call, __LINE__, __FILE__, mpi_status); \
exit(mpi_status); \
} \
}
#include <cuda.h> // CUDA Driver API：cuMem*（VMM）与 cuMulticast* 都在此头文件
// ---- CUDA Driver API 错误检查宏 ----
// 执行 call 得 CUresult → 非 CUDA_SUCCESS 则 cuGetErrorString 取错误串 → 打印并退出。
#define CUDA_CALL(call) \
{ \
CUresult cudaStatus = call; \
if (CUDA_SUCCESS != cudaStatus) { \
const char* error_string; \
cuGetErrorString(cudaStatus, &error_string); \
fprintf(stderr, \
"ERROR: CUDA Driver call \"%s\" in line %d of file %s failed " \
"with " \
"%s (%d).\n", \
#call, __LINE__, __FILE__, error_string, cudaStatus); \
exit(cudaStatus); \
} \
}
#include <cuda_runtime.h> // CUDA Runtime API（cudaMalloc/cudaMemcpy/kernel 启动等）
// ---- NVTX 时间线标注（编译加 -DUSE_NVTX 生效；Makefile 默认开启） ----
#ifdef USE_NVTX
#include <nvtx3/nvToolsExt.h> // NVTX v3 头文件
const uint32_t colors[] = {0x0000ff00, 0x000000ff, 0x00ffff00, 0x00ff00ff,
0x0000ffff, 0x00ff0000, 0x00ffffff}; // ARGB 调色板
const int num_colors = sizeof(colors) / sizeof(uint32_t); // 颜色数量
// PUSH_RANGE(name,cid)：压入一个带颜色的 NVTX 区间（nsys 时间线上可见）
// 逐行含义：取色 → 清零属性结构 → 填版本/大小/颜色类型/颜色/消息类型/文本 → 压栈
#define PUSH_RANGE(name, cid) \
{ \
int color_id = cid; \
color_id = color_id % num_colors; \
nvtxEventAttributes_t eventAttrib = {0}; \
eventAttrib.version = NVTX_VERSION; \
eventAttrib.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE; \
eventAttrib.colorType = NVTX_COLOR_ARGB; \
eventAttrib.color = colors[color_id]; \
eventAttrib.messageType = NVTX_MESSAGE_TYPE_ASCII; \
eventAttrib.message.ascii = name; \
nvtxRangePushEx(&eventAttrib); \
}
#define POP_RANGE nvtxRangePop(); // 弹出区间
#else
#define PUSH_RANGE(name, cid) // 未开 NVTX 时为空宏
#define POP_RANGE
#endif
// ---- CUDA Runtime 错误检查宏（与 .cu 中同款） ----
#define CUDA_RT_CALL(call) \
{ \
cudaError_t cudaStatus = call; \
if (cudaSuccess != cudaStatus) { \
fprintf(stderr, \
"ERROR: CUDA RT call \"%s\" in line %d of file %s failed " \
"with " \
"%s (%d).\n", \
#call, __LINE__, __FILE__, cudaGetErrorString(cudaStatus), cudaStatus); \
exit(cudaStatus); \
} \
}
// ---- 单/双精度编译期选择（必须与 .cu 保持一致） ----
#ifdef USE_DOUBLE
typedef double real; // 双精度
#define MPI_REAL_TYPE MPI_DOUBLE // 配套 MPI 类型
#else
typedef float real; // 默认单精度
#define MPI_REAL_TYPE MPI_FLOAT
#endif
// Group l2 norm and barrier counter for simplified handling of Multi Cast (MC) memory
// 【注】MC/UC 组播内存中的共享结构：value=本 rank L2 部分平方和；arrival_counter=栅栏计数
struct real_int_pair {
real value; // 待跨 GPU 求和的标量
unsigned int arrival_counter; // 已到达 GPU 计数（跨迭代累计，永不复位）
};
constexpr real tol = 1.0e-8; // 收敛阈值：L2 范数低于此值即停止迭代
const real PI = 2.0 * std::asin(1.0); // 由 asin(1)=pi/2 得到 pi（避免手写常量误差）
// ---- 由 jacobi_kernels.cu 提供的 4 个 launch 封装（host 侧前向声明） ----
// 边界初始化
void launch_initialize_boundaries(real* __restrict__ const a_new, real* __restrict__ const a,
const real pi, const int offset, const int nx, const int my_ny,
const int ny);
// 单 GPU 版 Jacobi 迭代
void launch_jacobi_kernel(real* __restrict__ const a_new, const real* __restrict__ const a,
real* __restrict__ const l2_norm, const int iy_start, const int iy_end,
const int nx, const bool calculate_norm, cudaStream_t stream);
// 多 GPU 版：计算 + 直写上/下邻居 halo 行
void launch_jacobi_p2p_kernel(real* __restrict__ const a_new, const real* __restrict__ const a,
real* __restrict__ const l2_norm, const int iy_start,
const int iy_end, const int nx, real* __restrict__ const a_new_top,
const int top_iy, real* __restrict__ const a_new_bottom,
const int bottom_iy, const bool calculate_norm, cudaStream_t stream);
// MC 在网栅栏 + 在网求和
void launch_all_reduce_norm_barrier_kernel(real* __restrict__ const l2_norm,
real_int_pair* __restrict__ partial_l2_norm_uc,
real_int_pair* __restrict__ partial_l2_norm_mc,
const int num_gpus, const int iter, cudaStream_t stream);
// 单 GPU 参考实现：返回耗时并把参考解写回 a_ref_h
double single_gpu(const int nx, const int ny, const int iter_max, real* const a_ref_h,
const int nccheck, const bool print);
// ---- 命令行解析：取 "-name value" 形式的参数值，缺省给 default_val ----
template <typename T>
T get_argval(char** begin, char** end, const std::string& arg, const T default_val) {
T argval = default_val; // 先填默认值
char** itr = std::find(begin, end, arg); // 在 argv 中找参数名
if (itr != end && ++itr != end) { // 找到且后面还有值
std::istringstream inbuf(*itr); // 用字符串流解析成 T
inbuf >> argval;
}
return argval;
}
// ---- 命令行解析：仅判断开关型参数（如 -csv）是否存在 ----
bool get_arg(char** begin, char** end, const std::string& arg) {
char** itr = std::find(begin, end, arg); // 找参数名
if (itr != end) {
return true; // 存在 → 开
}
return false;
}
// ---- 向上对齐：value 取整到 granularity 的整数倍（granularity 须为 2 的幂） ----
template <typename T>
T round_up(const T value, const T granularity) {
return (value + granularity - 1) & ~(granularity - 1); // 位运算版向上取整
}
// ---- 一份"可跨进程/跨节点共享的 fabric 显存"的完整账本 ----
struct fabric_mem {
CUdeviceptr ptr; // 本地分配映射后的虚拟地址
size_t size; // 用户请求的大小
size_t aligned_size; // 对齐到分配粒度后的大小
size_t granularity; // VMM 分配最小粒度（通常 2MB）
int device_id; // 所在 GPU 编号
CUmemGenericAllocationHandle generic_handle; // 本地 cuMemCreate 得到的物理内存句柄
CUmemFabricHandle fabric_handle; // 导出的 fabric 共享句柄（交 MPI 传输）
CUdeviceptr ptr_top; // 上方邻居缓冲映射进本进程的地址
CUmemGenericAllocationHandle generic_handle_top; // 上方邻居内存的导入句柄
CUdeviceptr ptr_bottom; // 下方邻居缓冲映射进本进程的地址
CUmemGenericAllocationHandle generic_handle_bottom; // 下方邻居内存的导入句柄
};
// ---- MC(组播)/UC(单播) 内存对的账本 ----
// 同一块本地物理显存被映射为两个视图：uc_ptr 普通读写只看本副本；
// mc_ptr 供 multimem.* 指令经 NVSwitch 对所有成员副本归约/组播。
struct uc_mc_pair {
CUdeviceptr uc_ptr; // UC 视图虚拟地址
CUdeviceptr mc_ptr; // MC 视图虚拟地址
CUmemGenericAllocationHandle mc_handle; // MC 对象句柄（rank0 创建，其余导入）
CUmemGenericAllocationHandle uc_handle; // 本地物理内存句柄
size_t mc_size; // MC 逻辑大小（按推荐粒度对齐）
size_t uc_size; // UC 物理分配大小（按最小粒度对齐）
};
/**
* Allocate sharable fabric memory:
* 1. Create a CUDA memory handle for the allocation (cuMemCreate)
* 2. Export handle to fabric handle for later exchange (cuMemExportToShareableHandle)
* 3. Reserve address range for allocation (cuMemAddressReserve)
* 4. Map allocation into address range (cuMemMap)
* 5. Make allocation accessible (cuMemSetAccess)
*
* 【注】分配一块可共享显存的 VMM 标准五步法。关键在 FABRIC 句柄类型：
* POSIX fd 句柄只能经 Unix socket 在本机进程间传递；本示例要靠 MPI 在任意
* 节点间传句柄（MNNVL 域），必须用编码了 fabric 全局标识的 FABRIC 类型。
*/
void allocate_fabric_mem(fabric_mem& fm, const size_t size, const int device_id) {
fm.device_id = device_id; // 记录目标 GPU
// It is required to use a socket in the Unix domain with
// CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR to exchange sharable handles between processes. To
// enable communcation of shareable handles via MPI we need CU_MEM_HANDLE_TYPE_FABRIC so using
// that here.
const CUmemAllocationHandleType handle_type = CU_MEM_HANDLE_TYPE_FABRIC; // 句柄类型：FABRIC
CUmemAllocationProp prop = {}; // 分配属性结构清零
prop.requestedHandleTypes = handle_type; // 要求分配支持 FABRIC 导出
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; // 物理页分配即固定（不可换出/迁移）
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; // 分配位置：显存
prop.location.id = fm.device_id; // 在哪块卡上
// 查询该属性下分配的最小粒度（通常 2MB），后续 size 按它对齐
CUDA_CALL(cuMemGetAllocationGranularity(&fm.granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM));
fm.size = size; // 记录请求大小
fm.aligned_size = round_up(fm.size, fm.granularity); // 实际分配按粒度向上取整
CUDA_CALL(cuMemCreate(&fm.generic_handle, fm.aligned_size, &prop, 0 /*flags*/)); // ①建物理句柄
// ②导出为 fabric 句柄------之后由 MPI 发给上/下邻居进程
CUDA_CALL(cuMemExportToShareableHandle(&fm.fabric_handle, fm.generic_handle, handle_type,
0 /*flags*/));
// ③在本进程 VA 空间保留一段连续地址（alignment 用 granularity）
CUDA_CALL(cuMemAddressReserve(&fm.ptr, fm.aligned_size, fm.granularity, 0 /*baseVA*/, 0 /*flags*/));
// ④把物理分配映射进保留的 VA 段（offset=0）
CUDA_CALL(cuMemMap(fm.ptr, fm.aligned_size, 0 /*offset*/, fm.generic_handle, 0 /*flags*/));
// ⑤授予本 GPU 对该映射的读写权限（VA 此时才真正可用）
CUmemAccessDesc desc = {};
desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; // 授权对象：某 GPU
desc.location.id = prop.location.id; // 即本卡
desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; // 读写权限
CUDA_CALL(cuMemSetAccess(fm.ptr, fm.aligned_size, &desc, 1 /*count*/));
}
// 释放 fabric 显存：解除映射 → 释放物理句柄 → 释放 VA 段
void free_fabric_mem(fabric_mem& fm) {
CUDA_CALL(cuMemUnmap(fm.ptr, fm.aligned_size)); // 解除 VA→物理 映射
CUDA_CALL(cuMemRelease(fm.generic_handle)); // 释放物理内存句柄
CUDA_CALL(cuMemAddressFree(fm.ptr, fm.aligned_size)); // 释放 VA 保留段
// Fabric handle does not hold any resources so doe not need to be freed
// 【注】fabric 句柄只是句柄的"传输编码"，本身不占资源，无需释放
}
/**
* Map peer memory for direct load store access:
* 1. Reserve address range for peer allocation (cuMemAddressReserve)
* 2. Map peer memory into address range (cuMemMap)
* 3. Make allocation accessible (cuMemSetAccess)
*
* 【注】把邻居（已导入句柄的）物理显存映射进本进程地址空间：此后本 GPU kernel
* 对 ptr_top/ptr_bottom 的普通 load/store 会经 NVLink fabric 直达对端显存。
* desc.location.id 仍填本卡 id------是"本 GPU 获得访问权"，而非在对方建窗口。
*/
void map_peers(fabric_mem& fm) {
// 为上方、下方邻居的映射各保留一段 VA
CUDA_CALL(cuMemAddressReserve(&fm.ptr_top, fm.aligned_size, fm.granularity, 0 /*baseVA*/,
0 /*flags*/));
CUDA_CALL(cuMemAddressReserve(&fm.ptr_bottom, fm.aligned_size, fm.granularity, 0 /*baseVA*/,
0 /*flags*/));
// 把两个邻居的物理内存分别映射进对应 VA 段
CUDA_CALL(cuMemMap(fm.ptr_top, fm.aligned_size, 0 /*offset*/, fm.generic_handle_top, 0 /*flags*/));
CUDA_CALL(cuMemMap(fm.ptr_bottom, fm.aligned_size, 0 /*offset*/, fm.generic_handle_bottom,
0 /*flags*/));
// 授予本 GPU 对两段映射的读写权限
CUmemAccessDesc desc = {};
desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
desc.location.id = fm.device_id;
desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
CUDA_CALL(cuMemSetAccess(fm.ptr_top, fm.aligned_size, &desc, 1 /*count*/));
CUDA_CALL(cuMemSetAccess(fm.ptr_bottom, fm.aligned_size, &desc, 1 /*count*/));
}
// 解除对端映射并释放导入句柄与 VA 段
void unmap_peers(fabric_mem& fm) {
CUDA_CALL(cuMemUnmap(fm.ptr_top, fm.aligned_size)); // 解除上方映射
CUDA_CALL(cuMemUnmap(fm.ptr_bottom, fm.aligned_size)); // 解除下方映射
CUDA_CALL(cuMemRelease(fm.generic_handle_top)); // 释放上方导入句柄
CUDA_CALL(cuMemRelease(fm.generic_handle_bottom)); // 释放下方导入句柄
CUDA_CALL(cuMemAddressFree(fm.ptr_top, fm.aligned_size)); // 释放上方 VA 段
CUDA_CALL(cuMemAddressFree(fm.ptr_bottom, fm.aligned_size)); // 释放下方 VA 段
}
/* ============================================================================
* main ------ 主流程（第一段：初始化、前置检查、参考解、域分解、P2P 映射、MC 建组）
* ==========================================================================*/
int main(int argc, char* argv[]) {
MPI_CALL(MPI_Init(&argc, &argv)); // MPI 初始化
int rank;
MPI_CALL(MPI_Comm_rank(MPI_COMM_WORLD, &rank)); // 本进程全局 rank
int size;
MPI_CALL(MPI_Comm_size(MPI_COMM_WORLD, &size)); // 进程总数（= GPU 数）
int num_devices = 0;
CUDA_RT_CALL(cudaGetDeviceCount(&num_devices)); // 本机可见 GPU 数
// ---- 命令行参数 ----
const int iter_max = get_argval<int>(argv, argv + argc, "-niter", 1000); // 迭代上限
const int nccheck = get_argval<int>(argv, argv + argc, "-nccheck", 1); // 每多少轮查一次范数
const int nx = get_argval<int>(argv, argv + argc, "-nx", 16384); // 网格列数
const int ny = get_argval<int>(argv, argv + argc, "-ny", 16384); // 网格行数
const bool csv = get_arg(argv, argv + argc, "-csv"); // csv 输出开关
const bool use_mc_red = get_arg(argv, argv + argc, "-use_mc_red"); // 启用 MC 在网归约
// ---- 划分本机共享内存通信子，得到机内 rank（用于绑卡） ----
int local_rank = -1;
int local_size = 1;
{
MPI_Comm local_comm;
MPI_CALL(MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, rank, MPI_INFO_NULL,
&local_comm)); // 按"同机"分组
MPI_CALL(MPI_Comm_rank(local_comm, &local_rank)); // 机内序号
MPI_CALL(MPI_Comm_size(local_comm, &local_size)); // 本机进程数
MPI_CALL(MPI_Comm_free(&local_comm)); // 用完即释放
}
const int device_id = local_rank % num_devices; // 机内序号取模 → 绑卡
CUDA_RT_CALL(cudaSetDevice(device_id)); // 设定当前设备
CUDA_RT_CALL(cudaFree(0)); // 触发 CUDA context 惰性初始化
// ---- 前置条件 1：设备必须支持 FABRIC 句柄（MNNVL/NVLink fabric 环境） ----
int fabric_handle_supported = 0;
CUDA_CALL(cuDeviceGetAttribute(&fabric_handle_supported,
CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device_id));
if (!fabric_handle_supported) { // 不支持则报错退出
cudaDeviceProp prop;
CUDA_RT_CALL(cudaGetDeviceProperties(&prop, device_id));
fprintf(stderr, "ERROR: Creating fabric handles is not supported on device %d (%s)\n",
device_id, prop.name);
MPI_CALL(MPI_Finalize());
return -1;
}
// ---- 前置条件 2（仅 -use_mc_red）：必须支持 MC 组播对象（需 NVSwitch），
// 且不允许 MPS 超额订阅（多进程共用一卡时 MC 语义不成立） ----
if (use_mc_red) {
int multicast_supported = 0;
CUDA_CALL(cuDeviceGetAttribute(&multicast_supported,
CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, device_id));
if (!multicast_supported) { // 不支持 MC 对象 → 报错退出
cudaDeviceProp prop;
CUDA_RT_CALL(cudaGetDeviceProperties(&prop, device_id));
fprintf(stderr,
"ERROR: Creating Multicast Objects is not supported on device %d (%s)\n",
device_id, prop.name);
MPI_CALL(MPI_Finalize());
return -1;
}
if (1 < num_devices && num_devices < local_size) { // 进程数 > 卡数 → MPS 超订 → 拒绝
fprintf(stderr,
"ERROR: Creating Multicast Objects is not supported when oversubscribing a GPU "
"with MPS: %d ranks using %d (< %d) devices!\n",
local_size, num_devices, local_size);
MPI_CALL(MPI_Finalize());
return 1;
}
}
// ---- 每个进程先跑一遍单 GPU 参考实现：得参考解 a_ref_h 与基线耗时 runtime_serial ----
real* a_ref_h;
CUDA_RT_CALL(cudaMallocHost(&a_ref_h, nx * ny * sizeof(real))); // 锁页 host 缓冲：参考解
real* a_h;
CUDA_RT_CALL(cudaMallocHost(&a_h, nx * ny * sizeof(real))); // 锁页 host 缓冲：多 GPU 结果
double runtime_serial = single_gpu(nx, ny, iter_max, a_ref_h, nccheck, !csv && (0 == rank));
// ---- 行向域分解：全局内部行（去掉上下边界共 ny-2 行）尽量均分给 size 个 rank ----
// ny - 2 rows are distributed amongst `size` ranks in such a way
// that each rank gets either (ny - 2) / size or (ny - 2) / size + 1 rows.
// This optimizes load balancing when (ny - 2) % size != 0
int chunk_size; // 本 rank 实际分到的行数
int chunk_size_low = (ny - 2) / size; // 基础块大小（整除）
int chunk_size_high = chunk_size_low + 1; // 余数情况下的大块 = 基础块 + 1
// To calculate the number of ranks that need to compute an extra row,
// the following formula is derived from this equation:
// num_ranks_low * chunk_size_low + (size - num_ranks_low) * (chunk_size_low + 1) = ny - 2
// 【注】由上面的方程解出"拿小块的 rank 数"：num_ranks_low = size*(low+1) - (ny-2)
int num_ranks_low = size * chunk_size_low + size -
(ny - 2); // Number of ranks with chunk_size = chunk_size_low
if (rank < num_ranks_low)
chunk_size = chunk_size_low; // 前面的 rank 拿小块
else
chunk_size = chunk_size_high; // 后面的 rank 多拿 1 行
// Need to allocate with chunk_size_high on all ranks to ensure consistent sizes when mapping
// peer memory
// 【注】所有 rank 统一按"最大块 + 上下各 1 行 halo"分配：cuMemMap 映射对端时
// size 必须一致，而各 rank 行数可能差 1，统一按大块对齐。
fabric_mem a_fa; // 旧值缓冲（a）
allocate_fabric_mem(a_fa, (nx * (chunk_size_high + 2) * sizeof(real)), device_id);
real* a = reinterpret_cast<real*>(a_fa.ptr); // 转成普通指针供使用
fabric_mem a_new_fa; // 新值缓冲（a_new）
allocate_fabric_mem(a_new_fa, (nx * (chunk_size_high + 2) * sizeof(real)), device_id);
real* a_new = reinterpret_cast<real*>(a_new_fa.ptr);
CUDA_RT_CALL(cudaMemset(a, 0, nx * (chunk_size_high + 2) * sizeof(real))); // 清零 a
CUDA_RT_CALL(cudaMemset(a_new, 0, nx * (chunk_size_high + 2) * sizeof(real))); // 清零 a_new
// Calculate local domain boundaries
// 【注】计算本 rank 子区域在全局网格中的起始行号：小块 rank 段与
// 大块 rank 段的偏移公式不同（前半段每 rank low 行，后半段每 rank high 行）
int iy_start_global; // My start index in the global array
if (rank < num_ranks_low) {
iy_start_global = rank * chunk_size_low + 1; // 小块段（+1 跳过全局上边界）
} else {
iy_start_global =
num_ranks_low * chunk_size_low + (rank - num_ranks_low) * chunk_size_high + 1; // 大块段
}
int iy_end_global = iy_start_global + chunk_size - 1; // My last index in the global array
int iy_start = 1; // 本地缓冲首计算行（行 0 是顶部 halo）
int iy_end = iy_start + chunk_size; // 本地缓冲末计算行 +1（行 chunk_size+1 是底部 halo）
int iy_end_top; // 上邻居的 iy_end（我方写其底部 halo 的行号）
{
// Map memory of top and bottom peers to allow direct writes of halo data:
// 1. Exchange shareable fabric handles with MPI
// 2. Map fabric handles into local memory
// 3. Exchange top peers bottom boundary index (`iy_end_top`)
// 【注】halo 直写准备：环形确定邻居 → MPI 交换 a/a_new 的 fabric 句柄 →
// 交换上邻居的 iy_end（各 rank 块大小可能不同，不能假设）
const int top = rank > 0 ? rank - 1 : (size - 1); // 上邻居（rank0 绕到最后一名：周期边界）
const int bottom = (rank + 1) % size; // 下邻居（最后一名绕回 rank0）
CUmemFabricHandle fabric_handle_a_top; // 收：上邻居的 a 句柄
CUmemFabricHandle fabric_handle_a_bottom; // 收：下邻居的 a 句柄
// 发给上邻居、同时收下邻居发来的 a 句柄
MPI_CALL(MPI_Sendrecv(&a_fa.fabric_handle, sizeof(CUmemFabricHandle), MPI_BYTE, top, 0,
&fabric_handle_a_bottom, sizeof(CUmemFabricHandle), MPI_BYTE, bottom, 0,
MPI_COMM_WORLD, MPI_STATUSES_IGNORE));
// 发给下邻居、同时收上邻居发来的 a 句柄
MPI_CALL(MPI_Sendrecv(&a_fa.fabric_handle, sizeof(CUmemFabricHandle), MPI_BYTE, bottom, 0,
&fabric_handle_a_top, sizeof(CUmemFabricHandle), MPI_BYTE, top, 0,
MPI_COMM_WORLD, MPI_STATUSES_IGNORE));
// 把自己的 iy_end 发给下邻居，同时收上邻居的 iy_end（= iy_end_top）
MPI_CALL(MPI_Sendrecv(&iy_end, 1, MPI_INT, bottom, 0,
&iy_end_top, 1, MPI_INT, top, 0,
MPI_COMM_WORLD, MPI_STATUSES_IGNORE));
CUmemFabricHandle fabric_handle_a_new_top; // 收：上邻居的 a_new 句柄
CUmemFabricHandle fabric_handle_a_new_bottom; // 收：下邻居的 a_new 句柄
// a_new 的句柄交换（同上两步）
MPI_CALL(MPI_Sendrecv(&a_new_fa.fabric_handle, sizeof(CUmemFabricHandle), MPI_BYTE, top, 0,
&fabric_handle_a_new_bottom, sizeof(CUmemFabricHandle), MPI_BYTE, bottom, 0,
MPI_COMM_WORLD, MPI_STATUSES_IGNORE));
MPI_CALL(MPI_Sendrecv(&a_new_fa.fabric_handle, sizeof(CUmemFabricHandle), MPI_BYTE, bottom, 0,
&fabric_handle_a_new_top, sizeof(CUmemFabricHandle), MPI_BYTE, top, 0,
MPI_COMM_WORLD, MPI_STATUSES_IGNORE));
const CUmemAllocationHandleType handle_type = CU_MEM_HANDLE_TYPE_FABRIC; // 导入类型需一致
// 把收到的 4 个 fabric 句柄导入为本地可用的通用句柄
CUDA_CALL(cuMemImportFromShareableHandle(&a_fa.generic_handle_top, &fabric_handle_a_top,
handle_type));
CUDA_CALL(cuMemImportFromShareableHandle(&a_fa.generic_handle_bottom,
&fabric_handle_a_bottom, handle_type));
CUDA_CALL(cuMemImportFromShareableHandle(&a_new_fa.generic_handle_top,
&fabric_handle_a_new_top, handle_type));
CUDA_CALL(cuMemImportFromShareableHandle(&a_new_fa.generic_handle_bottom,
&fabric_handle_a_new_bottom, handle_type));
map_peers(a_fa); // 映射上/下邻居的 a（旧值缓冲）
map_peers(a_new_fa); // 映射上/下邻居的 a_new（新值缓冲）
}
// 4 个"对端视图"指针：kernel 内对它们 store 即写入邻居显存
real* a_top = reinterpret_cast<real*>(a_fa.ptr_top); // 上邻居的 a
real* a_bottom = reinterpret_cast<real*>(a_fa.ptr_bottom); // 下邻居的 a
real* a_new_top = reinterpret_cast<real*>(a_new_fa.ptr_top); // 上邻居的 a_new
real* a_new_bottom = reinterpret_cast<real*>(a_new_fa.ptr_bottom); // 下邻居的 a_new
// Set Dirichlet boundary conditions on left and right borders
// 【注】设左右竖边边界；offset = 本 rank 行 0（halo 行）对应的全局行号
launch_initialize_boundaries(a, a_new, PI, iy_start_global - 1, nx, (chunk_size + 2), ny);
CUDA_RT_CALL(cudaDeviceSynchronize()); // 等边界写完
cudaStream_t compute_stream;
CUDA_RT_CALL(cudaStreamCreate(&compute_stream)); // 计算流（所有 GPU 工作都排它）
real* l2_norm_d;
CUDA_RT_CALL(cudaMalloc(&l2_norm_d, sizeof(real))); // 设备端范数结果（标量）
real* l2_norm_h;
CUDA_RT_CALL(cudaMallocHost(&l2_norm_h, sizeof(real))); // 锁页 host 端范数（异步 D2H 用）
/* ------------------------------------------------------------------------
* MC 组播对象建立（-use_mc_red 时）------cuMulticast 系列 API 标准流程：
* 查粒度 → rank0 Create → 导出句柄 → MPI_Bcast → 其余导入 → AddDevice →
* Barrier 配齐 → 各 rank 分配 UC 物理内存并 BindMem → 映射 MC/UC 双视图 → 初始化
* ----------------------------------------------------------------------*/
uc_mc_pair partial_l2_handles; // MC/UC 账本
real_int_pair* partial_l2_norm = nullptr; // UC 视图指针（普通读写）
real_int_pair* partial_l2_norm_mc = nullptr; // MC 视图指针（multimem 指令用）
if (use_mc_red) {
const CUmemAllocationHandleType handle_type = CU_MEM_HANDLE_TYPE_FABRIC;
// Get the minimum/recommended granularity for the multicast object
// 【注】MC 对象属性：成员数=进程数、逻辑大小=一个 real_int_pair、FABRIC 句柄
CUmulticastObjectProp mc_prop = {};
mc_prop.numDevices = size; // 组成员数（每进程一卡）
mc_prop.size = sizeof(real_int_pair); // 逻辑空间大小（随后会对齐放大）
mc_prop.handleTypes = handle_type; // 可跨进程共享
size_t min_granularity; // 最小粒度（UC 物理分配对齐用）
size_t granularity; // 推荐粒度（MC 逻辑大小/映射对齐用）
CUDA_CALL(cuMulticastGetGranularity(&min_granularity, &mc_prop,
CU_MULTICAST_GRANULARITY_MINIMUM));
CUDA_CALL(cuMulticastGetGranularity(&granularity, &mc_prop,
CU_MULTICAST_GRANULARITY_RECOMMENDED));
mc_prop.size = round_up(mc_prop.size, granularity); // MC 逻辑大小对齐到推荐粒度
partial_l2_handles.mc_size = mc_prop.size; // 记录 MC 大小
partial_l2_handles.uc_size = round_up(sizeof(real_int_pair), min_granularity); // UC 大小对齐
CUmemFabricHandle fh; // 待广播的 MC fabric 句柄
if (0 == rank) {
// Allocate the multicast object
// 【注】rank0 创建 MC 对象（向 NVSwitch 登记组播组，不占物理显存）
CUDA_CALL(cuMulticastCreate(&partial_l2_handles.mc_handle, &mc_prop));
// 导出为 fabric 句柄供广播
CUDA_CALL(cuMemExportToShareableHandle(&fh, partial_l2_handles.mc_handle, handle_type, 0));
}
MPI_CALL(MPI_Bcast(&fh, sizeof(CUmemFabricHandle), MPI_BYTE, 0, MPI_COMM_WORLD)); // 广播句柄
if (0 != rank) {
// 其余 rank 导入同一个 MC 对象
CUDA_CALL(cuMemImportFromShareableHandle(&partial_l2_handles.mc_handle, &fh, handle_type));
}
// 每个 rank 把本卡注册为组成员（交换机侧登记成员端口）
CUDA_CALL(cuMulticastAddDevice(partial_l2_handles.mc_handle, device_id));
// Ensure all devices in this process are added BEFORE binding mem on any device
// 【注】全员 AddDevice 完成前 Bind/Map 会阻塞------先栅栏确保组配齐
MPI_CALL(MPI_Barrier(MPI_COMM_WORLD));
// 各 rank 分配本地物理显存（UC），作为 MC 对象偏移 0 处的本地副本
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; // 物理页固定
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; // 显存
prop.location.id = device_id; // 本卡
prop.requestedHandleTypes = handle_type; // 可共享（MC 绑定要求）
CUDA_CALL(cuMemCreate(&partial_l2_handles.uc_handle, partial_l2_handles.uc_size, &prop,
0 /*flags*/)); // 分配 UC 物理内存
// 绑定：把本卡 UC 内存登记为 MC 偏移 [0, uc_size) 处的本地副本
CUDA_CALL(cuMulticastBindMem(partial_l2_handles.mc_handle, 0, partial_l2_handles.uc_handle,
0, partial_l2_handles.uc_size, 0));
// MC Mapping
// 【注】映射 MC 组播视图：对该 VA 用 multimem.* 访问 → 交换机对所有副本归约/组播
CUDA_CALL(cuMemAddressReserve(&partial_l2_handles.mc_ptr, mc_prop.size, granularity,
0 /*baseVA*/, 0 /*flags*/)); // 保留 VA
CUDA_CALL(cuMemMap(partial_l2_handles.mc_ptr, mc_prop.size, 0 /*offset*/,
partial_l2_handles.mc_handle, 0 /*flags*/)); // 映射 MC 对象
CUmemAccessDesc desc = {};
desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
desc.location.id = device_id;
desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
CUDA_CALL(cuMemSetAccess(partial_l2_handles.mc_ptr, mc_prop.size, &desc, 1 /*count*/)); // 授权
// UC Mapping
// 【注】映射 UC 本地视图：同一块物理内存的普通视图，读写只看本 GPU 副本
CUDA_CALL(cuMemAddressReserve(&partial_l2_handles.uc_ptr, partial_l2_handles.uc_size,
granularity, 0 /*baseVA*/, 0 /*flags*/)); // 保留 VA
CUDA_CALL(cuMemMap(partial_l2_handles.uc_ptr, partial_l2_handles.uc_size, 0 /*offset*/,
partial_l2_handles.uc_handle, 0 /*flags*/)); // 映射 UC 句柄
CUDA_CALL(cuMemSetAccess(partial_l2_handles.uc_ptr, partial_l2_handles.uc_size, &desc, 1 /*count*/)); // 授权
partial_l2_norm = reinterpret_cast<real_int_pair*>(partial_l2_handles.uc_ptr); // UC 视图指针
partial_l2_norm_mc = reinterpret_cast<real_int_pair*>(partial_l2_handles.mc_ptr); // MC 视图指针
// 初始化本地副本：value=0、计数器=0（计数器跨迭代累计，永不复位）
real_int_pair partial_l2_norm_init;
partial_l2_norm_init.value = 0.0;
partial_l2_norm_init.arrival_counter = 0;
CUDA_RT_CALL(cudaMemcpy(partial_l2_norm, &partial_l2_norm_init, sizeof(real_int_pair), cudaMemcpyHostToDevice));
}
/* ============================================================================
* main ------ 主流程（第二段：主迭代循环、校验、输出、资源销毁）
* ==========================================================================*/
if (!csv && 0 == rank) {
printf(
"Jacobi relaxation: %d iterations on %d x %d mesh with norm check "
"every %d iterations\n",
iter_max, ny, nx, nccheck); // 运行参数横幅（仅 rank0、非 csv 模式）
}
int iter = 0; // 迭代计数
real l2_norm = 1.0; // 范数初值（> tol 保证进入循环）
bool calculate_norm = true; // boolean to store whether l2 norm will be calculated in
// an iteration or not
// 【注】calculate_norm：本轮是否在 kernel 里融合计算范数
CUDA_RT_CALL(cudaDeviceSynchronize()); // 等所有初始化工作完成
MPI_CALL(MPI_Barrier(MPI_COMM_WORLD)); // 全体对齐后开始计时
double start = MPI_Wtime(); // 计时起点
PUSH_RANGE("Jacobi solve", 0) // NVTX 区间开始（nsys 可见）
/* ========================================================================
* 主迭代循环：每轮 = 清零范数累加器 → jacobi_p2p_kernel（计算+直写邻居 halo）
* → 若查范数: MC 在网栅栏+求和 或 D2H+MPI_Allreduce → 判收敛 → swap 双缓冲
* ======================================================================*/
while (l2_norm > tol && iter < iter_max) {
// 清零范数累加器：MC 模式清 UC 本地副本的 value 字段；传统模式清 l2_norm_d
CUDA_RT_CALL(cudaMemsetAsync((use_mc_red ? &(partial_l2_norm->value) : l2_norm_d), 0,
sizeof(real), compute_stream));
// 本轮是否算范数：按 nccheck 周期；非 csv 时每 100 轮也强制查一次（便于打印进度）
calculate_norm = (iter % nccheck) == 0 || (!csv && (iter % 100) == 0);
// 核心 kernel：五点差分计算 + 首/末行顺手直写进上/下邻居的 halo 行
launch_jacobi_p2p_kernel(a_new, a, (use_mc_red ? &(partial_l2_norm->value) : l2_norm_d),
iy_start, iy_end, nx, a_new_top, iy_end_top, a_new_bottom, 0,
calculate_norm, compute_stream);
if (calculate_norm) {
if (use_mc_red) {
// MC 路径：单线程 kernel 完成 全员计数器+1(栅栏) → 自旋等全员 →
// NVSwitch 对所有副本 value 求和 → 开方写回 l2_norm_d
launch_all_reduce_norm_barrier_kernel(l2_norm_d, partial_l2_norm, partial_l2_norm_mc, size, iter, compute_stream);
}
// 把范数结果异步拷回 host（锁页内存 + 同流，kernel 完成后执行）
CUDA_RT_CALL(cudaMemcpyAsync(l2_norm_h, l2_norm_d, sizeof(real), cudaMemcpyDeviceToHost,
compute_stream));
}
if (calculate_norm) {
CUDA_RT_CALL(cudaStreamSynchronize(compute_stream)); // 等本轮 GPU 工作全部完成
if (!use_mc_red) {
// 传统路径：host 侧 MPI_Allreduce 求和（SUM）后开方
MPI_CALL(MPI_Allreduce(l2_norm_h, &l2_norm, 1, MPI_REAL_TYPE, MPI_SUM, MPI_COMM_WORLD));
l2_norm = std::sqrt(l2_norm);
} else {
// Need to ensure that partial_l2_norm is not reset by any GPU before all GPUs are
// done reading it
// 【注】MC 路径的保护栅栏：必须等所有 GPU 都读完本轮归约结果，
// 才允许任何 GPU 在下一轮清零 value------否则快 GPU 的清零会覆盖
// 慢 GPU 尚未读的数据（归约只有加法语义，没有"读完成"时序保证）。
MPI_CALL(MPI_Barrier(MPI_COMM_WORLD));
l2_norm = *l2_norm_h; // 读取已在网求和并开方后的范数
}
if (!csv && 0 == rank && (iter % 100) == 0) {
printf("%5d, %0.6f\n", iter, l2_norm); // 每 100 轮打印一次进度
}
} else {
// 【注】不查范数的轮次也要栅栏：halo 是 push 直写模型、无接收确认，
// 必须防止快 rank 领先过多轮、把邻居仍在读的 a_new halo 行提前覆盖。
MPI_CALL(MPI_Barrier(MPI_COMM_WORLD));
}
// 双缓冲交换：本地指针与对端视图指针同步交换（邻居写我的 a_new ↔ 我写邻居的 a_new）
std::swap(a_new, a);
std::swap(a_new_top, a_top);
std::swap(a_new_bottom, a_bottom);
iter++;
}
double stop = MPI_Wtime(); // 计时终点
POP_RANGE // NVTX 区间结束
// ---- 把本 rank 的结果行拷回 host（a 的第 1 行起，即跳过顶部 halo） ----
// 【注】拷贝长度取 min：最后一个 rank 的块可能延伸到全局下边界，防止越界
CUDA_RT_CALL(cudaMemcpy(a_h + iy_start_global * nx, a + nx,
std::min((ny - iy_start_global) * nx, chunk_size * nx) * sizeof(real),
cudaMemcpyDeviceToHost));
// ---- 与单 GPU 参考解逐点比对（只比内点，跳过左右边界列） ----
int result_correct = 1;
for (int iy = iy_start_global; result_correct && (iy < iy_end_global); ++iy) {
for (int ix = 1; result_correct && (ix < (nx - 1)); ++ix) {
if (std::fabs(a_ref_h[iy * nx + ix] - a_h[iy * nx + ix]) > tol) {
fprintf(stderr,
"ERROR on rank %d: a[%d * %d + %d] = %f does not match %f "
"(reference)\n",
rank, iy, nx, ix, a_h[iy * nx + ix], a_ref_h[iy * nx + ix]);
result_correct = 0; // 发现不一致即标记并跳出
}
}
}
// 全体 rank 校验结果取 MIN：任何一个 rank 错则整体判错
int global_result_correct = 1;
MPI_CALL(MPI_Allreduce(&result_correct, &global_result_correct, 1, MPI_INT, MPI_MIN,
MPI_COMM_WORLD));
result_correct = global_result_correct;
// ---- rank0 输出性能结果 ----
if (rank == 0 && result_correct) {
if (csv) {
// csv 格式：实现名, nx, ny, 迭代上限, nccheck, GPU数, 1(占位), 多GPU耗时, 单GPU耗时
printf("multi_node_p2p, %d, %d, %d, %d, %d, 1, %f, %f\n", nx, ny, iter_max, nccheck,
size, (stop - start), runtime_serial);
} else {
printf("Num GPUs: %d.\n", size);
// 人类可读格式：单/多 GPU 耗时、加速比、并行效率(%)
printf(
"%dx%d: 1 GPU: %8.4f s, %d GPUs: %8.4f s, speedup: %8.2f, "
"efficiency: %8.2f \n",
ny, nx, runtime_serial, size, (stop - start), runtime_serial / (stop - start),
runtime_serial / (size * (stop - start)) * 100);
}
}
CUDA_RT_CALL(cudaStreamDestroy(compute_stream)); // 销毁计算流
CUDA_RT_CALL(cudaFreeHost(l2_norm_h)); // 释放 host 锁页范数缓冲
CUDA_RT_CALL(cudaFree(l2_norm_d)); // 释放设备端范数缓冲
// ---- MC 对象销毁：先栅栏确保所有进程都不再使用，再解绑/解映射/释放 ----
if (use_mc_red) {
// Need to ensure that all processes are done using the MC object before unbinding it
MPI_CALL(MPI_Barrier(MPI_COMM_WORLD));
// 解绑：mcOffset=0、size 必须与当初 Bind 时一致（否则行为未定义）
CUDA_CALL(cuMulticastUnbind(partial_l2_handles.mc_handle, device_id, 0,
partial_l2_handles.uc_size));
CUDA_CALL(cuMemUnmap(partial_l2_handles.mc_ptr, partial_l2_handles.mc_size)); // 解 MC 视图
CUDA_CALL(cuMemUnmap(partial_l2_handles.uc_ptr, partial_l2_handles.uc_size)); // 解 UC 视图
CUDA_CALL(cuMemRelease(partial_l2_handles.uc_handle)); // 释放 UC 物理内存
CUDA_CALL(cuMemRelease(partial_l2_handles.mc_handle)); // 销毁 MC 对象
CUDA_CALL(cuMemAddressFree(partial_l2_handles.mc_ptr, partial_l2_handles.mc_size)); // 释放 MC VA
CUDA_CALL(cuMemAddressFree(partial_l2_handles.uc_ptr, partial_l2_handles.uc_size)); // 释放 UC VA
}
unmap_peers(a_new_fa); // 解除 a_new 的上/下邻居映射
unmap_peers(a_fa); // 解除 a 的上/下邻居映射
free_fabric_mem(a_new_fa); // 释放 a_new 本地 fabric 显存
free_fabric_mem(a_fa); // 释放 a 本地 fabric 显存
CUDA_RT_CALL(cudaFreeHost(a_h)); // 释放 host 结果缓冲
CUDA_RT_CALL(cudaFreeHost(a_ref_h)); // 释放 host 参考解缓冲
MPI_CALL(MPI_Finalize()); // MPI 收尾
return (result_correct == 1) ? 0 : 1; // 校验通过返回 0，否则返回 1
}
/* ============================================================================
* single_gpu ------ 单 GPU 参考实现（第三段）
* 作用：产生参考解（供多 GPU 校验）+ 串行基线时间（供加速比计算）。
* 用 3 条 stream + 3 个 event 实现"计算与周期边界拷贝 overlap"。
* ==========================================================================*/
double single_gpu(const int nx, const int ny, const int iter_max, real* const a_ref_h,
const int nccheck, const bool print) {
real* a; // 旧值缓冲
real* a_new; // 新值缓冲
cudaStream_t compute_stream; // 主流：Jacobi 计算
cudaStream_t push_top_stream; // 拷贝流1：末行 → 顶部 halo（周期边界上半）
cudaStream_t push_bottom_stream; // 拷贝流2：首行 → 底部 halo（周期边界下半）
cudaEvent_t compute_done; // 事件：计算完成
cudaEvent_t push_top_done; // 事件：上半拷贝完成
cudaEvent_t push_bottom_done; // 事件：下半拷贝完成
real* l2_norm_d; // 设备端范数
real* l2_norm_h; // host 端范数
int iy_start = 1; // 首计算行（行 0 是顶部 halo）
int iy_end = (ny - 1); // 末计算行 +1（行 ny-1 是底部 halo）
CUDA_RT_CALL(cudaMalloc(&a, nx * ny * sizeof(real))); // 分配旧值缓冲
CUDA_RT_CALL(cudaMalloc(&a_new, nx * ny * sizeof(real))); // 分配新值缓冲
CUDA_RT_CALL(cudaMemset(a, 0, nx * ny * sizeof(real))); // 清零
CUDA_RT_CALL(cudaMemset(a_new, 0, nx * ny * sizeof(real)));
// Set diriclet boundary conditions on left and right boarder
// 【注】设左右竖边边界（offset=0：单 GPU 覆盖全局网格）
launch_initialize_boundaries(a, a_new, PI, 0, nx, ny, ny);
CUDA_RT_CALL(cudaDeviceSynchronize()); // 等边界写完
CUDA_RT_CALL(cudaStreamCreate(&compute_stream)); // 建主流
CUDA_RT_CALL(cudaStreamCreate(&push_top_stream)); // 建拷贝流1
CUDA_RT_CALL(cudaStreamCreate(&push_bottom_stream)); // 建拷贝流2
// 建 3 个事件（DisableTiming：只用于同步，不测时，开销更小）
CUDA_RT_CALL(cudaEventCreateWithFlags(&compute_done, cudaEventDisableTiming));
CUDA_RT_CALL(cudaEventCreateWithFlags(&push_top_done, cudaEventDisableTiming));
CUDA_RT_CALL(cudaEventCreateWithFlags(&push_bottom_done, cudaEventDisableTiming));
CUDA_RT_CALL(cudaMalloc(&l2_norm_d, sizeof(real))); // 设备范数
CUDA_RT_CALL(cudaMallocHost(&l2_norm_h, sizeof(real))); // host 锁页范数
CUDA_RT_CALL(cudaDeviceSynchronize()); // 开计时前确保设备空闲
if (print)
printf(
"Single GPU jacobi relaxation: %d iterations on %d x %d mesh with "
"norm "
"check every %d iterations\n",
iter_max, ny, nx, nccheck); // 运行参数横幅
int iter = 0;
real l2_norm = 1.0;
bool calculate_norm = true;
double start = MPI_Wtime(); // 计时起点
PUSH_RANGE("Jacobi solve", 0)
// ---- 单 GPU 主循环 ----
while (l2_norm > tol && iter < iter_max) {
CUDA_RT_CALL(cudaMemsetAsync(l2_norm_d, 0, sizeof(real), compute_stream)); // 清零范数
// 主流等上一轮的两条 halo 拷贝完成（保护将被覆写的 halo 行）
CUDA_RT_CALL(cudaStreamWaitEvent(compute_stream, push_top_done, 0));
CUDA_RT_CALL(cudaStreamWaitEvent(compute_stream, push_bottom_done, 0));
calculate_norm = (iter % nccheck) == 0 || (iter % 100) == 0; // 本轮是否算范数
// 单 GPU 版 Jacobi kernel（无邻居直写）
launch_jacobi_kernel(a_new, a, l2_norm_d, iy_start, iy_end, nx, calculate_norm,
compute_stream);
CUDA_RT_CALL(cudaEventRecord(compute_done, compute_stream)); // 记录"计算完成"事件
if (calculate_norm) {
// 范数异步拷回 host（同流，排在 kernel 之后）
CUDA_RT_CALL(cudaMemcpyAsync(l2_norm_h, l2_norm_d, sizeof(real), cudaMemcpyDeviceToHost,
compute_stream));
}
// Apply periodic boundary conditions
// 【注】周期边界：末计算行(iy_end-1) → 顶部 halo(行0)；首计算行(iy_start) → 底部 halo(iy_end)。
// 两条拷贝各在自己的流上，与后续计算 overlap。
CUDA_RT_CALL(cudaStreamWaitEvent(push_top_stream, compute_done, 0)); // 拷贝流1等计算完
CUDA_RT_CALL(cudaMemcpyAsync(a_new, a_new + (iy_end - 1) * nx, nx * sizeof(real),
cudaMemcpyDeviceToDevice, push_top_stream)); // 末行→顶部
CUDA_RT_CALL(cudaEventRecord(push_top_done, push_top_stream)); // 记事件
CUDA_RT_CALL(cudaStreamWaitEvent(push_bottom_stream, compute_done, 0)); // 拷贝流2等计算完
CUDA_RT_CALL(cudaMemcpyAsync(a_new + iy_end * nx, a_new + iy_start * nx, nx * sizeof(real),
cudaMemcpyDeviceToDevice, compute_stream)); // 首行→底部（排主流上）
CUDA_RT_CALL(cudaEventRecord(push_bottom_done, push_bottom_stream)); // 记事件
if (calculate_norm) {
CUDA_RT_CALL(cudaStreamSynchronize(compute_stream)); // 等本轮完成
l2_norm = *l2_norm_h; // 读范数
l2_norm = std::sqrt(l2_norm); // 开方得 L2 范数
if (print && (iter % 100) == 0) printf("%5d, %0.6f\n", iter, l2_norm); // 打印进度
}
std::swap(a_new, a); // 双缓冲交换
iter++;
}
POP_RANGE
double stop = MPI_Wtime(); // 计时终点
CUDA_RT_CALL(cudaMemcpy(a_ref_h, a, nx * ny * sizeof(real), cudaMemcpyDeviceToHost)); // 拷出参考解
// ---- 资源清理 ----
CUDA_RT_CALL(cudaEventDestroy(push_bottom_done)); // 销毁事件
CUDA_RT_CALL(cudaEventDestroy(push_top_done));
CUDA_RT_CALL(cudaEventDestroy(compute_done));
CUDA_RT_CALL(cudaStreamDestroy(push_bottom_stream)); // 销毁流
CUDA_RT_CALL(cudaStreamDestroy(push_top_stream));
CUDA_RT_CALL(cudaStreamDestroy(compute_stream));
CUDA_RT_CALL(cudaFreeHost(l2_norm_h)); // 释放 host 范数
CUDA_RT_CALL(cudaFree(l2_norm_d)); // 释放设备范数
CUDA_RT_CALL(cudaFree(a_new)); // 释放新值缓冲
CUDA_RT_CALL(cudaFree(a)); // 释放旧值缓冲
return (stop - start); // 返回单 GPU 总耗时（基线）
}7.2 cu 原文件注释
cpp
/* ============================================================================
* jacobi_kernels.cu ------ 逐行中文注释版（原代码逐行保留，未做任何改动）
*
* 本文件包含 4 个 kernel 及其 host 侧 launch 封装：
* 1. initialize_boundaries ------ 写左右两条边的 Dirichlet 边界（正弦波）
* 2. jacobi_kernel ------ 单 GPU 五点差分迭代（无 halo 交换）
* 3. jacobi_p2p_kernel ------ 多 GPU 版：计算时把边界行直接 store 进
* 上/下邻居 GPU 的显存（NVLink fabric 直写）
* 4. all_reduce_norm_barrier_kernel ------ 基于 NVSwitch 组播(MC)内存的
* "全员到达栅栏 + L2 范数在网求和"
* ==========================================================================*/
/* Copyright (c) 2017-2018, 2024, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions
* are met:
* * Redistributions of source code must retain the above copyright
* notice, this list of conditions and the following disclaimer.
* * Redistributions in binary form must reproduce the above copyright
* notice, this list of conditions and the following disclaimer in the
* documentation and/or other materials provided with the distribution.
* * Neither the name of NVIDIA CORPORATION nor the names of its
* contributors may be used to endorse or promote products derived
* from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
* EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
* PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
* CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
* EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
* PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
* PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
* OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/
#include <cassert> // assert()：校验 barrier kernel 只以单线程启动
#include <cmath> // std::sqrt：对归约出的全局平方和开方得到 L2 范数
#include <cstdio> // fprintf：CUDA_RT_CALL 宏里的错误输出
#include <cstdlib> // exit：CUDA_RT_CALL 宏里的出错退出
#include <cuda/atomic> // libcu++ 的 cuda::atomic_ref，用于系统域自旋等待计数器
#ifdef HAVE_CUB // 编译期开关：有 CUB 就用块内归约（默认开启）
#include <cub/block/block_reduce.cuh> // cub::BlockReduce：块内 warp 级求和，减少全局原子加次数
#endif // HAVE_CUB
// ---- CUDA Runtime 错误检查宏：任何 runtime 调用失败即打印位置与错误串并退出 ----
// 【注】宏体是多行续行结构，注释只能加在宏外；逐行含义：
// 第1行 #define：宏名 CUDA_RT_CALL(call)，参数为被调用的表达式
// 花括号块：把宏体包成复合语句，保证 if/else 中使用不出语法问题
// cudaError_t cudaStatus = call：执行调用并保存返回码
// if (cudaSuccess != cudaStatus)：判断是否失败
// fprintf(...)：打印出错的调用文本(#call 字符串化)、行号、文件名、错误串与错误码
// exit(cudaStatus)：以错误码直接终止进程（示例代码不做恢复）
#define CUDA_RT_CALL(call) \
{ \
cudaError_t cudaStatus = call; \
if (cudaSuccess != cudaStatus) { \
fprintf(stderr, \
"ERROR: CUDA RT call \"%s\" in line %d of file %s failed " \
"with " \
"%s (%d).\n", \
#call, __LINE__, __FILE__, cudaGetErrorString(cudaStatus), cudaStatus); \
exit(cudaStatus); \
} \
}
// ---- 单/双精度编译期选择 ----
#ifdef USE_DOUBLE // 编译加 -DUSE_DOUBLE 则用 double
typedef double real; // 统一浮点类型别名：双精度
#define MPI_REAL_TYPE MPI_DOUBLE // 配套的 MPI 数据类型（.cu 里实际未用到，仅为与 jacobi.cpp 对齐）
#else // 默认单精度
typedef float real;
#define MPI_REAL_TYPE MPI_FLOAT
#endif
// ---- MC/UC 组播内存中的共享结构 ----
// 每个 GPU 的本地物理副本(UC)上各存一份；MC 地址是所有副本的"组播集合视图"。
struct real_int_pair {
real value; // 本 rank 的 L2 范数部分平方和（待跨 GPU 求和）
unsigned int arrival_counter; // 栅栏计数器：本轮已到达的 GPU 数（跨迭代累计，永不复位）
};
/* ----------------------------------------------------------------------------
* kernel 1: initialize_boundaries
* 设置左(ix=0)、右(ix=nx-1)两条竖边的 Dirichlet 边界值 sin(2*pi*y_global/(ny-1))。
* 参数：
* a_new, a : 双缓冲，两份都要写（迭代中会 swap）
* pi : 圆周率
* offset : 本 rank 第 0 行（halo 行）在全局网格中的行号
* nx : 行宽（列数）
* my_ny : 本地网格总行数（含上下各 1 行 halo）
* ny : 全局网格总行数
* --------------------------------------------------------------------------*/
__global__ void initialize_boundaries(real* __restrict__ const a_new, real* __restrict__ const a,
const real pi, const int offset, const int nx,
const int my_ny, const int ny) {
// grid-stride 循环：每个线程按全局线程号 stride 处理多行（行数 > 线程数时）
for (int iy = blockIdx.x * blockDim.x + threadIdx.x; iy < my_ny; iy += blockDim.x * gridDim.x) {
// 该行对应的全局纵坐标为 offset+iy，算出正弦边界值
const real y0 = sin(2.0 * pi * (offset + iy) / (ny - 1));
a[iy * nx + 0] = y0; // 旧缓冲的左边界
a[iy * nx + (nx - 1)] = y0; // 旧缓冲的右边界
a_new[iy * nx + 0] = y0; // 新缓冲的左边界（swap 后仍是边界）
a_new[iy * nx + (nx - 1)] = y0; // 新缓冲的右边界
}
}
// host 侧封装：按行数估算 block 数（每 block 128 线程），一维网格
void launch_initialize_boundaries(real* __restrict__ const a_new, real* __restrict__ const a,
const real pi, const int offset, const int nx, const int my_ny,
const int ny) {
initialize_boundaries<<<my_ny / 128 + 1, 128>>>(a_new, a, pi, offset, nx, my_ny, ny);
CUDA_RT_CALL(cudaGetLastError()); // 捕获 kernel 启动配置错误（异步错误不在这里）
}
/* ----------------------------------------------------------------------------
* kernel 2: jacobi_kernel ------ 单 GPU 参考实现用
* 标准五点差分：new = 0.25*(左+右+下+上)；可选融合计算局部 L2 平方和。
* 模板参数 BLOCK_DIM_X/Y：编译期块尺寸，供 CUB 临时存储定型。
* --------------------------------------------------------------------------*/
template <int BLOCK_DIM_X, int BLOCK_DIM_Y>
__global__ void jacobi_kernel(real* __restrict__ const a_new, const real* __restrict__ const a,
real* __restrict__ const l2_norm, const int iy_start,
const int iy_end, const int nx, const bool calculate_norm) {
#ifdef HAVE_CUB
// 定义块内归约类型：BLOCK_REDUCE_WARP_REDUCTIONS = 以 warp 洗牌为主的归约算法
typedef cub::BlockReduce<real, BLOCK_DIM_X, cub::BLOCK_REDUCE_WARP_REDUCTIONS, BLOCK_DIM_Y>
BlockReduce;
__shared__ typename BlockReduce::TempStorage temp_storage; // CUB 所需的共享内存暂存区
#endif // HAVE_CUB
int iy = blockIdx.y * blockDim.y + threadIdx.y + iy_start; // 本线程负责的行（从 iy_start 起）
int ix = blockIdx.x * blockDim.x + threadIdx.x + 1; // 本线程负责的列（+1 跳过左边界）
real local_l2_norm = 0.0; // 本线程的局部残差平方和
if (iy < iy_end && ix < (nx - 1)) { // 只算内点（末列 nx-1 是边界）
// 五点差分：上下左右四邻居取平均
const real new_val = 0.25 * (a[iy * nx + ix + 1] + a[iy * nx + ix - 1] +
a[(iy + 1) * nx + ix] + a[(iy - 1) * nx + ix]);
a_new[iy * nx + ix] = new_val; // 写回新值
if (calculate_norm) {
real residue = new_val - a[iy * nx + ix]; // 新旧值之差（残差）
local_l2_norm += residue * residue; // 累加残差平方
}
}
if (calculate_norm) {
#ifdef HAVE_CUB
// 块内归约：全块线程的 local_l2_norm 求和，结果在 0 号线程有效
real block_l2_norm = BlockReduce(temp_storage).Sum(local_l2_norm);
// 由块内 (0,0) 线程把块和原子加到全局累加器（每块仅 1 次原子加）
if (0 == threadIdx.y && 0 == threadIdx.x) atomicAdd(l2_norm, block_l2_norm);
#else
atomicAdd(l2_norm, local_l2_norm); // 无 CUB：每线程直接原子加
#endif // HAVE_CUB
}
}
// host 侧封装：32x32 线程块，二维网格覆盖整个本地子区域
void launch_jacobi_kernel(real* __restrict__ const a_new, const real* __restrict__ const a,
real* __restrict__ const l2_norm, const int iy_start, const int iy_end,
const int nx, const bool calculate_norm, cudaStream_t stream) {
constexpr int dim_block_x = 32; // 块宽（x 方向线程数）
constexpr int dim_block_y = 32; // 块高（y 方向线程数）
dim3 dim_grid((nx + dim_block_x - 1) / dim_block_x, // x 向上取整
((iy_end - iy_start) + dim_block_y - 1) / dim_block_y, 1); // y 向上取整
jacobi_kernel<dim_block_x, dim_block_y><<<dim_grid, {dim_block_x, dim_block_y, 1}, 0, stream>>>(
a_new, a, l2_norm, iy_start, iy_end, nx, calculate_norm);
CUDA_RT_CALL(cudaGetLastError());
}
/* ----------------------------------------------------------------------------
* kernel 3: jacobi_p2p_kernel ------ 多 GPU 核心 kernel：计算 + halo 直写融合
*
* 在 jacobi_kernel 基础上多了 4 个参数：
* a_new_top : 上方邻居(rank-1) a_new 缓冲映射进本进程的地址
* top_iy : 写入点在邻居缓冲中的行号（= 邻居的 iy_end，即其底部 halo 行）
* a_new_bottom : 下方邻居(rank+1) a_new 缓冲的映射地址
* bottom_iy : 写入点行号（实际恒为 0，即邻居顶部 halo 行）
* 对 a_new_top/bottom 的 store 是普通全局写，硬件经 NVLink fabric 直达对端显存。
* --------------------------------------------------------------------------*/
template <int BLOCK_DIM_X, int BLOCK_DIM_Y>
__global__ void jacobi_p2p_kernel(real* __restrict__ const a_new, const real* __restrict__ const a,
real* __restrict__ const l2_norm, const int iy_start,
const int iy_end, const int nx,
real* __restrict__ const a_new_top, const int top_iy,
real* __restrict__ const a_new_bottom, const int bottom_iy,
const bool calculate_norm) {
#ifdef HAVE_CUB
typedef cub::BlockReduce<real, BLOCK_DIM_X, cub::BLOCK_REDUCE_WARP_REDUCTIONS, BLOCK_DIM_Y>
BlockReduce;
__shared__ typename BlockReduce::TempStorage temp_storage; // 同单 GPU 版：CUB 共享暂存
#endif // HAVE_CUB
int iy = blockIdx.y * blockDim.y + threadIdx.y + iy_start; // 本线程的行
int ix = blockIdx.x * blockDim.x + threadIdx.x + 1; // 本线程的列（跳过左边界）
real local_l2_norm = 0.0;
if (iy < iy_end && ix < (nx - 1)) {
// 五点差分（与单 GPU 版相同）
const real new_val = 0.25 * (a[iy * nx + ix + 1] + a[iy * nx + ix - 1] +
a[(iy + 1) * nx + ix] + a[(iy - 1) * nx + ix]);
a_new[iy * nx + ix] = new_val; // 本地写回新值
// 若本线程算的是"本 rank 第一计算行"：同时写进上方邻居的底部 halo 行
// ------ 跨 NVLink 的远端 store，邻居下一轮即可直接读到
if (iy_start == iy) {
a_new_top[top_iy * nx + ix] = new_val;
}
// 若本线程算的是"本 rank 最后计算行"：写进下方邻居的顶部 halo 行（行号 0）
if ((iy_end - 1) == iy) {
a_new_bottom[bottom_iy * nx + ix] = new_val;
}
if (calculate_norm) {
real residue = new_val - a[iy * nx + ix]; // 残差
local_l2_norm += residue * residue; // 局部平方和
}
}
if (calculate_norm) {
#ifdef HAVE_CUB
real block_l2_norm = BlockReduce(temp_storage).Sum(local_l2_norm); // 块内求和
if (0 == threadIdx.y && 0 == threadIdx.x) atomicAdd(l2_norm, block_l2_norm);
#else
atomicAdd(l2_norm, local_l2_norm);
#endif // HAVE_CUB
}
}
// host 侧封装：网格/块配置与单 GPU 版一致，多传的 4 个参数是对端视图
void launch_jacobi_p2p_kernel(real* __restrict__ const a_new, const real* __restrict__ const a,
real* __restrict__ const l2_norm, const int iy_start,
const int iy_end, const int nx, real* __restrict__ const a_new_top,
const int top_iy, real* __restrict__ const a_new_bottom,
const int bottom_iy, const bool calculate_norm, cudaStream_t stream) {
constexpr int dim_block_x = 32;
constexpr int dim_block_y = 32;
dim3 dim_grid((nx + dim_block_x - 1) / dim_block_x,
((iy_end - iy_start) + dim_block_y - 1) / dim_block_y, 1);
jacobi_p2p_kernel<dim_block_x, dim_block_y>
<<<dim_grid, {dim_block_x, dim_block_y, 1}, 0, stream>>>(
a_new, a, l2_norm, iy_start, iy_end, nx, a_new_top, top_iy, a_new_bottom, bottom_iy,calculate_norm);
CUDA_RT_CALL(cudaGetLastError());
}
/* ----------------------------------------------------------------------------
* kernel 4: all_reduce_norm_barrier_kernel ------ NVSwitch 在网计算的精华
* 仅 1 线程执行（<<<1,1>>>，assert 校验），一次完成：
* 栅栏（multimem.red 全员计数器+1）→ 别名 fence → 自旋等全员到达
* → 在网求和（multimem.ld_reduce 对所有副本 value 求和）→ 开方写回。
* 替代传统路径：D2H 拷贝 + MPI_Allreduce + MPI_Barrier。
* --------------------------------------------------------------------------*/
__global__ void all_reduce_norm_barrier_kernel(real* const l2_norm,
real_int_pair* partial_l2_norm_uc,
real_int_pair* partial_l2_norm_mc,
const unsigned int expected_count) {
// 防御性校验：本 kernel 的语义要求全 GPU 只有 1 个线程在执行它
assert(1 == blockDim.x * blockDim.y * blockDim.z * gridDim.x * gridDim.y * gridDim.z);
real l2_norm_sum = 0.0; // 接收在网归约结果（全局 L2 平方和）
#if __CUDA_ARCH__ >= 900 // multimem 指令需 sm_90(Hopper)+；低档架构下 kernel 为空操作
// atomic reduction to all replicas
// this can be conceptually thought of as __threadfence_system(); atomicAdd_system(arrival_counter_mc, 1);
// See https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-multimem-ld-reduce-multimem-st-multimem-red
// for multimem PTX doc
// 【行注】multimem.red.release.sys.global.add.u32：
// red = 归约写（无返回值）；.release = 保证本 GPU 此前的写（本 rank 的 L2 部分和）
// 对系统域可见后计数器才 +1；.sys = 系统作用域（覆盖所有 GPU 与 host）；
// 该请求发到 NVSwitch，由交换机把 +1 原子地施加到【所有成员 GPU 的物理副本】。
// "l" 约束传 64 位地址，"n" 约束传立即数 1，"memory" 告知编译器内存被改写。
asm volatile ("multimem.red.release.sys.global.add.u32 [%0], %1;" ::"l"(&(partial_l2_norm_mc->arrival_counter)), "n"(1) : "memory");
// Need a fence between MC and UC access to the same memory:
// - fence.proxy instructions establish an ordering between memory accesses that may happen through different proxies
// - Value .alias of the .proxykind qualifier refers to memory accesses performed using virtually aliased addresses to the same memory location.
// from https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar
// 【行注】fence.proxy.alias：MC 地址与 UC 地址是同一块物理内存的两个虚拟别名，
// 经由不同 proxy（代理）访问；此 fence 建立"上面的 MC 写"与"下面的 UC 读"之间的顺序。
asm volatile ("fence.proxy.alias;" ::: "memory");
// spin wait with acquire ordering on UC mapping till all peers have arrived in this iteration
// Note: all ranks reach an MPI_Barrier after this kernel, such that it is not possible for the barrier to be unblocked by an
// arrival of a rank for the next iteration if some other rank is slow.
// 【行注】在本地 UC 副本上对计数器做系统域原子引用，acquire 自旋等待：
// 直到计数器达到 expected_count（= num_gpus*(iter+1)，跨迭代累计），即全员已到达。
// acquire 保证：看到计数器达标后，再读 value 时能看到所有 GPU 已写入的部分和。
cuda::atomic_ref<unsigned int, cuda::thread_scope_system> ac(partial_l2_norm_uc->arrival_counter);
while (expected_count > ac.load(cuda::memory_order_acquire));
// Atomic load reduction from all replicas. It does not provide ordering so it can be relaxed.
// 【行注】multimem.ld_reduce.relaxed.sys.global.add.f32/f64：
// 由 NVSwitch 读取【所有成员副本】同一偏移的 value 并求和(sum)后返回------在网归约；
// .relaxed 因为顺序已由上面的栅栏保证，这里不需要额外排序语义。
// "=d"/"=f" 约束把结果写入 l2_norm_sum（double 用 d 寄存器，float 用 f 寄存器）。
#ifdef USE_DOUBLE
asm volatile ("multimem.ld_reduce.relaxed.sys.global.add.f64 %0, [%1];" : "=d"(l2_norm_sum) : "l"(&(partial_l2_norm_mc->value)) : "memory");
#else
asm volatile ("multimem.ld_reduce.relaxed.sys.global.add.f32 %0, [%1];" : "=f"(l2_norm_sum) : "l"(&(partial_l2_norm_mc->value)) : "memory");
#endif
#endif
// 全局平方和开方得到 L2 范数，写入 l2_norm（随后被 D2H 拷回 host 做收敛判断）
*l2_norm = std::sqrt(l2_norm_sum);
}
// host 侧封装
void launch_all_reduce_norm_barrier_kernel(real* __restrict__ const l2_norm,
real_int_pair* __restrict__ partial_l2_norm_uc,
real_int_pair* __restrict__ partial_l2_norm_mc,
const int num_gpus, const int iter,
cudaStream_t stream) {
// calculating expected count as unsigned for well defined overflow to correctly handle large
// iteration counts with many GPUs
// 【行注】用 unsigned 计算阈值：溢出回绕行为有定义，大迭代数×多 GPU 时比较仍正确
unsigned int expected_count = num_gpus;
// iter starts at 0 so need to scale with iter+1
// 【行注】计数器从 0 起每轮每 GPU +1：第 iter 轮（0 基）达标值 = num_gpus*(iter+1)
expected_count *= (iter + 1);
// 单 block 单线程启动（与 kernel 内 assert 对应）
all_reduce_norm_barrier_kernel<<<1, 1, 0, stream>>>(l2_norm, partial_l2_norm_uc,
partial_l2_norm_mc, expected_count);
CUDA_RT_CALL(cudaGetLastError());
}7.3. Makefile
bash
# =============================================================================
# Makefile ------ multi_node_p2p 构建脚本（逐行中文注释版，原内容未改动）
#
# 产物：可执行文件 jacobi
# jacobi_kernels.cu → nvcc 编译（设备端 kernel）
# jacobi.cpp → mpicxx 编译（host 逻辑 + MPI + CUDA Driver API）
# =============================================================================
# Copyright (c) 2017-2018, NVIDIA CORPORATION. All rights reserved.
NP ?= 1 # mpirun 进程数，?= 表示可被环境变量/命令行覆盖（如 make run NP=8）
NVCC=nvcc # CUDA 编译器
MPICXX=mpicxx # MPI 的 C++ 编译器包装器
MPIRUN ?= mpirun # MPI 启动器（可覆盖为 srun/jsrun 等）
CUDA_HOME ?= /usr/local/cuda # CUDA 安装路径
# ---- 历代架构的 gencode 选项（系列仓库的历史保留） ----
GENCODE_SM30 := -gencode arch=compute_30,code=sm_30 # Kepler
GENCODE_SM35 := -gencode arch=compute_35,code=sm_35 # Kepler
GENCODE_SM37 := -gencode arch=compute_37,code=sm_37 # Kepler
GENCODE_SM50 := -gencode arch=compute_50,code=sm_50 # Maxwell
GENCODE_SM52 := -gencode arch=compute_52,code=sm_52 # Maxwell
GENCODE_SM60 := -gencode arch=compute_60,code=sm_60 # Pascal
GENCODE_SM70 := -gencode arch=compute_70,code=sm_70 # Volta
GENCODE_SM80 := -gencode arch=compute_80,code=sm_80 # Ampere
# Hopper：同时产出 SASS(sm_90) 与 PTX(compute_90 前向兼容)；
# 本示例的 multimem/fabric 能力最低需要 sm_90 + NVSwitch
GENCODE_SM90 := -gencode arch=compute_90,code=sm_90 -gencode arch=compute_90,code=compute_90
GENCODE_FLAGS := $(GENCODE_SM70) $(GENCODE_SM80) $(GENCODE_SM90) # 默认编三档：70/80/90
ifdef BUILD_SM_ARCH
# 只编指定架构，如 make BUILD_SM_ARCH=90 → 只保留对应一档
GENCODE_FLAGS := $(GENCODE_SM$(BUILD_SM_ARCH))
endif
ifdef DISABLE_CUB
# 不用 CUB 块归约：逐线程 atomicAdd，并让 ptxas 优化浮点原子加以减少序列化
NVCC_FLAGS = -Xptxas --optimize-float-atomics
else
# 默认：定义 HAVE_CUB，kernel 内使用 cub::BlockReduce
NVCC_FLAGS = -DHAVE_CUB
endif
ifdef SKIP_CUDA_AWARENESS_CHECK
# 跳过 MPI 的 CUDA-aware 检查（本文件实际未用此宏，为系列统一保留）
MPICXX_FLAGS = -DSKIP_CUDA_AWARENESS_CHECK
endif
# -lineinfo：保留源码行号供 nsys/ncu 分析；-std=c++14：libcu++ 的 cuda::atomic_ref 所需
NVCC_FLAGS += -lineinfo $(GENCODE_FLAGS) -std=c++14
MPICXX_FLAGS += -DUSE_NVTX -I$(CUDA_HOME)/include -std=c++14 # 打开 NVTX 时间线标注
# 链接：CUDA runtime(-lcudart) + driver(-lcuda，提供 cuMem*/cuMulticast*) + dl
LD_FLAGS = -L$(CUDA_HOME)/lib64 -lcudart -lcuda -ldl
# 最终链接：用 mpicxx 把 host 对象与 kernel 对象链成 jacobi
jacobi: Makefile jacobi.cpp jacobi_kernels.o
$(MPICXX) $(MPICXX_FLAGS) jacobi.cpp jacobi_kernels.o $(LD_FLAGS) -o jacobi
# kernel 编译规则：nvcc -c 生成目标文件
jacobi_kernels.o: Makefile jacobi_kernels.cu
$(NVCC) $(NVCC_FLAGS) jacobi_kernels.cu -c
.PHONY.: clean # 声明伪目标（注意原写法带点号，实为书写习惯）
clean: # 清理：可执行、目标文件、nsys 报告、sanitizer 日志
rm -f jacobi jacobi_kernels.o *.nsys-rep jacobi.*.compute-sanitizer.log
# 显存正确性检查：compute-sanitizer 随 mpirun 每进程跑一份，%q{} 按 rank 展开日志名
sanitize: jacobi
$(MPIRUN) -np $(NP) compute-sanitizer --log-file jacobi.%q{OMPI_COMM_WORLD_RANK}.compute-sanitizer.log ./jacobi -niter 10
run: jacobi # 直接运行（默认参数：16384x16384 网格、1000 迭代）
$(MPIRUN) -np $(NP) ./jacobi
# 性能剖析：Nsight Systems 跟踪 MPI/CUDA/NVTX，每进程输出一份报告
profile: jacobi
$(MPIRUN) -np $(NP) nsys profile --trace=mpi,cuda,nvtx -o jacobi.%q{OMPI_COMM_WORLD_RANK} ./jacobi -niter 10