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

如果你写过GPU程序一定遇到过这样的场景你手上有8块DCU每块卡上都跑着一个进程PEProcessing Element你想让0号卡上的数据能被1号卡直接读取——怎么做最直接的办法是0号卡把数据从显存拷到CPU内存1号卡再从CPU内存拷到自己的显存。数据兜了一大圈CPU在中间当“搬运工”延迟高、带宽低CPU还被占着干不了别的活。有没有更好的办法有。让DCU之间直接通信绕过CPU。DUSHMEM就是干这个的。一、DUSHMEM是什么DUSHMEM的全称是DCU Unified Shared Memory是DTK中的一套DCU集合通信库。它提供了一种叫PGASPartitioned 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先记住三个概念PEProcessing 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】#includedushmem.h#includedushmemx.h__global__voidsimple_shift(int*destination){intmypedushmem_my_pe();// 我是几号卡intnpesdushmem_n_pes();// 一共几张卡intpeer(mype1)%npes;// 我要发给谁下一张卡dushmem_int_p(destination,mype,peer);// 把我的ID写到peer的destination里}intmain(){dushmem_init();// 初始化DUSHMEM【7†L13】intmype_nodedushmem_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_shift1,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()-1dushmem_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){intmypedushmem_my_pe();// 在GPU上也能调用intpeer(mype1)%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如果你熟悉MPIMessage 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_commMPI_COMM_WORLD;dushmemx_init_attr(DUSHMEMX_INIT_WITH_MPI_COMM,attr);// 基于MPI初始化DUSHMEM【7†L14】// ... 混合使用MPI和DUSHMEM ...dushmem_finalize();MPI_Finalize();这种灵活性让现有MPI应用可以逐步迁移到DUSHMEM而不需要推倒重来。七、环境变量调优的“旋钮”DUSHMEM提供了丰富的环境变量来控制运行时行为【7†L59-L64】DUSHMEM_SYMMETRIC_SIZE每个PE的对称堆大小默认1GB【7†L59】DUSHMEM_REMOTE_TRANSPORT选择远程传输协议默认ibrcInfiniBand RC【7†L62】DUSHMEM_HCA_LIST指定使用哪些InfiniBand网卡【7†L63】DUSHMEM_IB_ENABLE_IBGDA启用IBGDAInfiniBand GPU Direct Async功能【7†L64】DUSHMEM_DEBUG开启调试信息【7†L60】这些变量让你可以在不重新编译的情况下调整DUSHMEM的行为——比如在调试时打开日志、在生产环境调大对称堆大小。八、总结回到开头的问题多卡通信为什么难因为数据在卡之间搬动需要有人协调——谁来发、谁来收、什么时候发、什么时候收。传统方式让CPU当“中间人”慢且占资源。DUSHMEM的思路是让DCU自己管自己。用对称内存让所有PE共享一个逻辑上的全局地址空间用put/get实现单边通信接收方不需要参与用原子操作保证多PE同时访问时的数据一致性用集合通信实现PE之间的协同用Device端API让通信逻辑直接在Kernel里完成如果说HIP Runtime API解决的是“怎么让一块DCU干活”那DUSHMEM解决的就是“怎么让一堆DCU一起干活”——而且是高效地、直接地、不需要CPU插手地一起干活。对于多卡编程的开发者来说DUSHMEM提供了一套顺手且强大的工具。从单卡到多卡你不需要重新学习一套完全不同的编程模型——put、get、atomic、barrier这些概念在DUSHMEM里都能找到直观的对应。