当前位置：首页 > news >正文

从“自己搬砖”到“指挥别人搬砖”：DUSHMEM库如何让多卡通信更简单

news 2026/7/9 1:21:46

如果你写过GPU程序，一定遇到过这样的场景：你手上有8块DCU，每块卡上都跑着一个进程（PE，Processing Element），你想让0号卡上的数据能被1号卡直接读取——怎么做？

最直接的办法是：0号卡把数据从显存拷到CPU内存，1号卡再从CPU内存拷到自己的显存。数据兜了一大圈，CPU在中间当“搬运工”，延迟高、带宽低，CPU还被占着干不了别的活。

有没有更好的办法？有。让DCU之间直接通信，绕过CPU。

DUSHMEM就是干这个的。

一、DUSHMEM是什么？

DUSHMEM的全称是DCU Unified Shared Memory，是DTK中的一套DCU集合通信库。它提供了一种叫PGAS（Partitioned Global Address Space，分区全局地址空间）的并行编程模型。

听起来很唬人？拆开看就明白了。

“分区”：每个PE（进程）都有自己的内存，这些内存是“分区”的——各自管各自的。

“全局地址空间”：虽然内存是分区的，但从编程角度看，所有PE的内存构成了一个统一的“全局地址空间”。你可以用dushmem_ptr拿到远程PE上某个变量的地址【7†L19】，然后像访问本地内存一样读写它。

换句话说：DUSHMEM让每个PE都能直接读写其他PE的内存，不需要CPU插手。

这和之前讲的HIP Runtime API有什么不同？

HIP Runtime API解决的是“单块DCU上怎么写程序”——怎么分配显存、怎么启动Kernel、怎么管理流。
DUSHMEM解决的是“多块DCU之间怎么通信”——怎么让数据在卡之间高效流动、怎么让多卡协同干活。

一个管“单卡编程”，一个管“多卡通信”，两者是互补的关系。

二、核心概念：PE、对称内存、Team

要理解DUSHMEM，先记住三个概念：

PE（Processing Element）：就是DUSHMEM里的一个进程。每个PE有一个唯一的编号（从0开始），用dushmem_my_pe()查询当前PE的编号【7†L17】，用dushmem_n_pes()查询总共有多少个PE【7†L18】。

对称内存（Symmetric Memory）：这是一块特殊的内存区域，所有PE都能远程访问。用dushmem_malloc分配的内存就是对称内存【7†L24】——你在0号卡上分配一块内存，1号卡也能通过DUSHMEM API读写它。

Team：一组PE的集合。DUSHMEM的集体通信（如dushmem_barrier、dushmem_broadcast）都是基于Team进行的【7†L28】。

三、最简示例：8张卡，环形传数据

光说概念太抽象，看个完整例子【7†L9-L10】：

#include<dushmem.h>#include<dushmemx.h>__global__voidsimple_shift(int*destination){intmype=dushmem_my_pe();// 我是几号卡？intnpes=dushmem_n_pes();// 一共几张卡？intpeer=(mype+1)%npes;// 我要发给谁？（下一张卡）dushmem_int_p(destination,mype,peer);// 把我的ID写到peer的destination里}intmain(){dushmem_init();// 初始化DUSHMEM【7†L13】intmype_node=dushmem_team_my_pe(DUSHMEMX_TEAM_NODE);hipSetDevice(mype_node);// 每个PE绑定一块DCUhipStream_t stream;hipStreamCreate(&stream);int*destination=(int*)dushmem_malloc(sizeof(int));// 分配对称内存【7†L24】// 启动Kernel：每个PE把自己的ID写到下一张卡的destination里simple_shift<<<1,1,0,stream>>>(destination);dushmemx_barrier_all_on_stream(stream);// 等待所有通信完成【7†L27】intmsg;hipMemcpyAsync(&msg,destination,sizeof(int),hipMemcpyDeviceToHost,stream);hipStreamSynchronize(stream);// 验证结果：我收到的是上一张卡发来的IDintexpected=(dushmem_my_pe()-1+dushmem_n_pes())%dushmem_n_pes();printf("PE %d: received %d, expected %d\n",dushmem_my_pe(),msg,expected);dushmem_free(destination);// 释放对称内存【7†L24】dushmem_finalize();// 退出DUSHMEM【7†L12】return0;}

这段代码在做什么？

8张卡，每张卡上跑一个PE。每个PE把自己的编号写到下一张卡的destination里。最后，每张卡都收到了上一张卡发来的编号【7†L11】。

注意一个关键细节：destination是用dushmem_malloc分配的——这意味着它在所有PE上都是对称的。0号卡的destination在1号卡上通过dushmem_ptr拿到地址后可以直接访问【7†L19】。这就是PGAS模型的精髓：远程内存像本地内存一样访问。

四、API分类：四大金刚

DUSHMEM的API可以分成四类【7†L7-L59】：

1. 初始化和查询——告诉DUSHMEM“我准备好了”

dushmem_init()：初始化DUSHMEM库，所有PE必须调用【7†L13】
dushmem_finalize()：退出DUSHMEM，释放资源【7†L12】
dushmem_my_pe()/dushmem_n_pes()：查询PE编号和总数【7†L17-L18】
dushmem_ptr()：获取远程PE上某个对称地址的本地指针【7†L19】

dushmem_init是集体操作——所有PE必须一起调用，不能有的调有的不调。dushmem_finalize同理【7†L12】。

2. 内存管理——在“共享内存”里划地盘

dushmem_malloc(size)：从对称堆里分配一块内存【7†L24】
dushmem_calloc(count, size)：分配并清零【7†L26】
dushmem_free(ptr)：释放对称内存【7†L24】
dushmem_align(alignment, size)：按指定对齐分配【7†L24】

这些也是集体操作——所有PE必须用相同的参数调用【7†L25】。为啥？因为DUSHMEM需要在所有PE上维护一致的对称内存布局。

3. 点对点通信——直接“喊话”

点对点通信是DUSHMEM最核心的能力：一个PE直接读写另一个PE的内存。

Put（写）：把数据从本地写到远程【7†L42】：

dushmem_int_put(dest,source,nelems,pe);// 把source写到pe的dest里dushmem_int_p(dest,value,pe);// 单元素的极低延迟put【7†L44】dushmem_int_iput(dest,source,dst,sst,nelems,pe);// 带步长的put【7†L45】

Get（读）：从远程把数据读到本地【7†L46】：

dushmem_int_get(dest,source,nelems,pe);// 从pe的source读到本地的destdushmem_int_g(source,pe);// 单元素的极低延迟get，直接返回值【7†L48】dushmem_int_iget(dest,source,dst,sst,nelems,pe);// 带步长的get【7†L49】

原子操作：多PE同时读写同一块内存时保证数据一致性【7†L50-L58】：

dushmem_int_atomic_fetch_add(dest,value,pe);// 原子地读-加-写dushmem_int_atomic_compare_swap(dest,cond,value,pe);// 原子比较交换【7†L51】dushmem_int_atomic_fetch_and(dest,value,pe);// 原子位与【7†L55】dushmem_int_atomic_fetch_or(dest,value,pe);// 原子位或【7†L56】dushmem_int_atomic_fetch_xor(dest,value,pe);// 原子位异或【7†L58】

原子操作在分布式计数器、锁、累加器等场景中非常有用——多个PE同时更新同一个变量，不会出现数据竞争。

4. 集合通信——大家一起“开会”

当所有PE需要一起做事时，用集合通信【7†L27-L41】：

dushmem_barrier_all()：所有PE同步，等到大家都到了再继续【7†L27】
dushmem_barrier(team)：指定Team内的PE同步【7†L28】
dushmem_broadcast(dest, source, nelems, root)：从root PE把数据广播给所有人【7†L33】
dushmem_alltoall(dest, source, nelems)：所有人都给所有人发数据【7†L31】
dushmem_reduce(dest, source, nreduce, OP)：所有PE的数据做归约（求和、求最大等）【7†L37】

注意barrier和sync的区别【7†L29】：

dushmem_barrier：保证所有远程写操作完成，是“硬同步”
dushmem_sync：只保证本地写操作完成，不保证远程写操作完成，是“软同步”

五、在DCU上跑：Device端API

上面列的所有API都有Device端版本——也就是说，你可以在Kernel里直接调用DUSHMEM的通信函数【7†L17-L58】。

比如在Kernel里这样写：

__global__voidmy_kernel(int*data){intmype=dushmem_my_pe();// 在GPU上也能调用！intpeer=(mype+1)%dushmem_n_pes();dushmem_int_p(data,mype,peer);// 在GPU上直接发数据给另一张卡}

这意味着通信逻辑可以写在Kernel内部，而不需要回到CPU端发起通信。对于需要细粒度通信的算法（比如迭代式求解器、图计算），这能大幅减少CPU参与，降低延迟。

但有一个重要限制：如果Kernel里用了DUSHMEM的同步或集合通信API（如dushmem_barrier、dushmem_barrier_all），必须用dushmemx_collective_launch启动，不能用普通的hipLaunchKernel【7†L22-L23】。因为这类API需要所有PE协同，普通启动方式无法保证这种协同性。

六、与MPI的对比：为什么需要DUSHMEM？

如果你熟悉MPI（Message Passing Interface），可能会问：“DUSHMEM和MPI有什么区别？MPI不是也能做多卡通信吗？”

核心区别在于通信模式：

MPI是“双边”的：发送方调用MPI_Send，接收方调用MPI_Recv，两方必须配合。
DUSHMEM是“单边”的：发送方直接put到接收方的内存，接收方完全不需要参与。

单边通信的好处是解耦——发送方可以随时发数据，接收方不需要提前准备好接收操作。这在异步计算、动态负载均衡等场景中非常有优势。

更重要的是，DUSHMEM可以和MPI混合使用【7†L11-L12】。如果你想在节点内部用DUSHMEM做高速通信、节点之间用MPI，完全可以：

MPI_Init(&argc,&argv);MPI_Comm_rank(MPI_COMM_WORLD,&rank);dushmemx_init_attr_t attr;attr.mpi_comm=&MPI_COMM_WORLD;dushmemx_init_attr(DUSHMEMX_INIT_WITH_MPI_COMM,&attr);// 基于MPI初始化DUSHMEM【7†L14】// ... 混合使用MPI和DUSHMEM ...dushmem_finalize();MPI_Finalize();