From 245c2de909e42d24dd0ba56e8954a8f60d80a81f Mon Sep 17 00:00:00 2001 From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com> Date: Sun, 16 Mar 2025 15:10:05 -0700 Subject: [PATCH] Enable LL128 on gfx942 (#1549) --- src/device/generate.py | 8 +- src/device/op128.h | 4 +- src/device/prims_ll128.h | 156 ++++++++++++++++++--------------------- src/enqueue.cc | 5 +- src/graph/tuning.cc | 17 +++-- 5 files changed, 92 insertions(+), 98 deletions(-) diff --git a/src/device/generate.py b/src/device/generate.py index 9f540eed31..2d076a7b31 100755 --- a/src/device/generate.py +++ b/src/device/generate.py @@ -308,7 +308,7 @@ with open(os.path.join(gensrc, "device_table.h"), "w") as f: for fn in primary_funcs: sym = paste("_", "ncclDevFunc", *fn) if fn[2] == "LL128": - out("#if defined(__gfx90a__) && defined(ENABLE_LL128)\n") + out("#if (defined(__gfx90a__) || defined(__gfx942__)) && defined(ENABLE_LL128)\n") out("%s %s();\n#else\n" % (func_declaration, sym)) fn_ll = fn[:2] + ("LL",) + fn[3:] sym_ll = paste("_", "ncclDevFunc", *fn_ll) @@ -325,7 +325,7 @@ with open(os.path.join(gensrc, "device_table.h"), "w") as f: if unroll != "2": continue sym = paste("_", "ncclDevFunc", *fn) if fn[2] == "LL128": - out("#if defined(__gfx90a__) && defined(ENABLE_LL128)\n") + out("#if (defined(__gfx90a__) || defined(__gfx942__)) && defined(ENABLE_LL128)\n") out("/*%4d*/ %s,\n#else\n" % (index, sym)) fn_ll = fn[:2] + ("LL",) + fn[3:] sym_ll = paste("_", "ncclDevFunc", *fn_ll) @@ -342,7 +342,7 @@ with open(os.path.join(gensrc, "device_table.h"), "w") as f: if unroll != "4": continue sym = paste("_", "ncclDevFunc", *fn) if fn[2] == "LL128": - out("#if defined(__gfx90a__) && defined(ENABLE_LL128)\n") + out("#if (defined(__gfx90a__) || defined(__gfx942__)) && defined(ENABLE_LL128)\n") out("/*%4d*/ %s,\n#else\n" % (index4, sym)) fn_ll = fn[:2] + ("LL",) + fn[3:] sym_ll = paste("_", "ncclDevFunc", *fn_ll) @@ -494,7 +494,7 @@ for name in name_to_funcs.keys(): (coll, algo, proto, redop, ty, unroll) = fn sym = paste("_", coll, algo, proto, redop, ty, unroll) if proto == "LL128": - out("#if defined(__gfx90a__) && defined(ENABLE_LL128)\n") + out("#if (defined(__gfx90a__) || defined(__gfx942__)) && defined(ENABLE_LL128)\n") out( "DEFINE_ncclDevFunc({sym}, ncclFunc{coll}, {redop_cxx}, {ty_cxx}, NCCL_ALGO_{algo}, NCCL_PROTO_{proto}, {unroll})\n" .format(sym=sym, coll=coll, redop_cxx=redop_to_cxx[redop], ty_cxx=ty_to_cxx[ty], diff --git a/src/device/op128.h b/src/device/op128.h index 7fc1bcdc33..99e4fa268e 100644 --- a/src/device/op128.h +++ b/src/device/op128.h @@ -54,12 +54,12 @@ inline __device__ void loadShmemMisaligned128(T *ptr, uint64_t &v0, uint64_t &v1 else if(sizeof(T) == 4) { #pragma unroll for(int e=0; e < 4; e++) - tmp4[e] = __builtin_nontemporal_load(ptr+e); + tmp4[e] = __builtin_nontemporal_load(reinterpret_cast(ptr)+e); } else /*sizeof(T)==8*/ { #pragma unroll for(int e=0; e < 2; e++) - tmp8[e] = __builtin_nontemporal_load(ptr+e); + tmp8[e] = __builtin_nontemporal_load(reinterpret_cast(ptr)+e); } v0 = tmp8[0]; v1 = tmp8[1]; diff --git a/src/device/prims_ll128.h b/src/device/prims_ll128.h index 771ee82a25..2114516422 100644 --- a/src/device/prims_ll128.h +++ b/src/device/prims_ll128.h @@ -113,53 +113,60 @@ private: if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += 1); } inline __device__ void postSend() { - if (sendConnTailPtr) { STORE((unsigned long long *)sendConnTailPtr, sendConnTail += 1); } + if (sendConnTailPtr) { +#if __CUDA_ARCH__ >= 900 + __threadfence_system(); +#else + __threadfence(); +#endif + STORE((unsigned long long *)sendConnTailPtr, sendConnTail += 1); + } } template __device__ __forceinline__ void loadRegsBegin(uint64_t(®s)[WordPerThread], T const *src, int eltN) { constexpr int EltPer16B = 16/sizeof(T); - /* We are aligned to 16 bytes, so load directly to registers no shmem. - * Flag threads load half as much data which gets shuffled to the even - * registers during Finish. The point of splitting into two phases is to - * defer that shuffle, which incurs a dependency stall, until after other - * memops are launched by the caller. - */ - #pragma unroll - for(int g=0; g < WordPerThread/2; g++) { - int ix = g*WARP_SIZE - 16*(g/2) + wid - (g%2)*(wid/4); - if(!flagThread || g%2==0) { - if(ix*EltPer16B < eltN) { - if(reinterpret_cast(src)%4 == 0) { - regs[2*g+0] = __builtin_nontemporal_load((uint64_t*)(src + ix*EltPer16B)); - regs[2*g+1] = __builtin_nontemporal_load((uint64_t*)(src + ix*EltPer16B)+1); - } else { - union { - uint64_t regs64[WordPerThread]; - uint32_t regs32[WordPerThread*2]; - uint16_t regs16[WordPerThread*4]; - uint8_t regs8[WordPerThread*8]; - }; - if (sizeof(T) == 8) { - uint64_t *src64 = (uint64_t*)(src+ix*EltPer16B); - for (int i=0; i < 2; i++) - regs64[2*g+i] = __builtin_nontemporal_load(src64+i); - } else if (sizeof(T) == 4) { - uint32_t *src32 = (uint32_t*)(src+ix*EltPer16B); - for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T); i++) - regs32[2*g+i] = __builtin_nontemporal_load(src32+i); - } else if (sizeof(T) == 2) { - uint16_t *src16 = (uint16_t*)(src+ix*EltPer16B); - for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T); i++) - regs16[2*g+i] = __builtin_nontemporal_load(src16+i); - } else if (sizeof(T) == 1) { - uint8_t *src8 = (uint8_t*)(src+ix*EltPer16B); - for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T); i++) - regs8[2*g+i] = __builtin_nontemporal_load(src8+i); - } - regs[2*g+0] = regs64[2*g+0]; - regs[2*g+1] = regs64[2*g+1]; - } + if(reinterpret_cast(src)%16 == 0) { + /* We are aligned to 16 bytes, so load directly to registers no shmem. + * Flag threads load half as much data which gets shuffled to the even + * registers during Finish. The point of splitting into two phases is to + * defer that shuffle, which incurs a dependency stall, until after other + * memops are launched by the caller. + */ + #pragma unroll + for(int g=0; g < WordPerThread/2; g++) { + int ix = g*WARP_SIZE - 16*(g/2) + wid - (g%2)*(wid/4); + if(!flagThread || g%2==0) { + if(ix*EltPer16B < eltN) + load128((uint64_t*)(src + ix*EltPer16B), regs[2*g+0], regs[2*g+1]); + } + } + } + else { + // Not aligned. Stage the smallest 16 byte aligned region subsuming the + // buffer into shmem. + int misalignment = reinterpret_cast(src) % 16; + uint64_t *src8 = reinterpret_cast(reinterpret_cast(src) & -uintptr_t(16)); + uint64_t *shm8 = shmemCvtPtr((uint64_t*)ncclScratchForWarp(warpInBlock)); + #pragma unroll + for(int g=0; g < WordPerThread/2; g++) + if((g*WARP_SIZE + wid)*16 < misalignment + eltN*sizeof(T)) + load128(src8 + 2*(g*WARP_SIZE + wid), regs[2*g+0], regs[2*g+1]); + #pragma unroll + for(int g=0; g < WordPerThread/2; g++) + storeShmem128(shm8 + 2*(g*WARP_SIZE + wid), regs[2*g+0], regs[2*g+1]); + + __syncwarp(); + + // Now load from shmem stage to regs. Preserve the same pre-shuffled layout + // as the aligned case since Finish() will be applied regardless. + T *shm = (T*)shm8 + misalignment/sizeof(T); + #pragma unroll + for(int g=0; g < WordPerThread/2; g++) { + int ix = g*WARP_SIZE - 16*(g/2) + wid - (g%2)*(wid/4); + if(!flagThread || g%2==0) { + if(ix*EltPer16B < eltN) + loadShmemMisaligned128(shm + ix*EltPer16B, regs[2*g+0], regs[2*g+1]); } } } @@ -183,44 +190,25 @@ private: if (flagThread) regs[2*g-1] = regs[2*g]; } // Write to dst if 4-byte aligned, shmem otherwise. - int misalignment = reinterpret_cast(dst)%4; + int misalignment = reinterpret_cast(dst)%16; + uint64_t *shm8 = shmemCvtPtr((uint64_t*)ncclScratchForWarp(warpInBlock)); #pragma unroll for(int g=0; g < WordPerThread/2; g++) { int ix = g*WARP_SIZE - 16*(g/2) + wid - (g%2)*(wid/4); if (!flagThread || g%2==0) { - if(misalignment == 0 && (ix+1)*EltPer16B <= eltN) { - __builtin_nontemporal_store(regs[2*g+0], (uint64_t*)(dst + ix*EltPer16B)); - __builtin_nontemporal_store(regs[2*g+1], (uint64_t*)(dst + ix*EltPer16B)+1); - } else { - union { - uint64_t regs64[WordPerThread]; - uint32_t regs32[WordPerThread*2]; - uint16_t regs16[WordPerThread*4]; - uint8_t regs8[WordPerThread*8]; - }; - regs64[2*g+0] = regs[2*g+0]; - regs64[2*g+1] = regs[2*g+1]; - int remaining = eltN - ix*EltPer16B; - if (sizeof(T) == 8) { - uint64_t *dst64 = (uint64_t*)(dst+ix*EltPer16B); - for (int i=0; i < 2 && i < remaining; i++) - __builtin_nontemporal_store(regs64[2*g+i], dst64+i); - } else if (sizeof(T) == 4) { - uint32_t *dst32 = (uint32_t*)(dst+ix*EltPer16B); - for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T) && i < remaining; i++) - __builtin_nontemporal_store(regs32[2*g+i], dst32+i); - } else if (sizeof(T) == 2) { - uint16_t *dst16 = (uint16_t*)(dst+ix*EltPer16B); - for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T) && i < remaining; i++) - __builtin_nontemporal_store(regs16[2*g+i], dst16+i); - } else if (sizeof(T) == 1) { - uint8_t *dst8 = (uint8_t*)(dst+ix*EltPer16B); - for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T) && i < remaining; i++) - __builtin_nontemporal_store(regs8[2*g+i], dst8+i); - } - } + if(misalignment == 0 && (ix+1)*EltPer16B <= eltN) + store128((uint64_t*)(dst + ix*EltPer16B), regs[2*g+0], regs[2*g+1]); + else + storeShmem128(shm8+2*ix, regs[2*g+0], regs[2*g+1]); } } + __syncwarp(); + // Write rest from shmem to dst. No need to coalesce stores to 16-bytes, + // the hardware keeps up fine. + T *shm = (T*)ncclScratchForWarp(warpInBlock); + int skip = misalignment == 0 ? eltN & -EltPer16B : 0; + for(int i=skip+wid; i < eltN; i += WARP_SIZE) + dst[i] = shm[i]; } #define WARP_MASK 0xffffffff @@ -241,12 +229,14 @@ private: needReload = false; #pragma unroll for (int u=0; ufunc == ncclFuncAllGather || info->func == ncclFuncReduceScatter) && comm->nNodes > 1) continue; for (int p=0; ptopo->type != RCCL_TOPO_XGMI_ALL) continue; + if (p == NCCL_PROTO_LL128 && !(comm->topo->type & RCCL_TOPO_XGMI_ALL)) { + table[a][p] = NCCL_ALGO_PROTO_IGNORE; + continue; + } bool backup; float time; NCCLCHECK(ncclTopoGetAlgoTime(comm, info->func, a, p, nBytes, numPipeOps, &time, &backup)); diff --git a/src/graph/tuning.cc b/src/graph/tuning.cc index 570bc375d2..3c90705fa5 100644 --- a/src/graph/tuning.cc +++ b/src/graph/tuning.cc @@ -224,29 +224,29 @@ static struct tuningModel tuning_model_4 { static struct tuningModel tuning_model_5 { .hwLat = { /* NVLINK */ - { /* Tree (LL/LL128/Simple)*/ { 0.9, 0.0, 2.3 }, /* Ring (LL/LL128/Simple)*/ { 0.8, 0.0, 2.1 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 0.9 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 0.0 }, /* NVLS */ { 0, 0, 0 }, /* NVLS Tree */ { 0, 0, 0 } }, + { /* Tree (LL/LL128/Simple)*/ { 0.9, 0.9, 2.3 }, /* Ring (LL/LL128/Simple)*/ { 0.8, 0.8, 2.1 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 0.9 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 0.0 }, /* NVLS */ { 0, 0, 0 }, /* NVLS Tree */ { 0, 0, 0 } }, /* PCI */ { /* Tree (LL/LL128/Simple)*/ { 2.2, 2.2, 5.7 }, /* Ring (LL/LL128/Simple)*/ { 2.2, 2.2, 5.7 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 5.7 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 5.7 }, /* NVLS */ { 0, 0, 0 }, /* NVLS Tree */ { 0, 0, 0 } }, /* NET */ - { /* Tree (LL/LL128/Simple)*/ { 10.5, 0.0, 25.0 }, /* Ring (LL/LL128/Simple)*/ { 9.5, 0.0, 320.0 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 10.5 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 0.0 }, /* NVLS */ { 0, 0, 0 }, /* NVLS Tree */ { 0, 0, 0 } }, + { /* Tree (LL/LL128/Simple)*/ { 10.5, 10.5, 25.0 }, /* Ring (LL/LL128/Simple)*/ { 9.5, 9.5, 320.0 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 10.5 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 0.0 }, /* NVLS */ { 0, 0, 0 }, /* NVLS Tree */ { 0, 0, 0 } }, }, .bwRatio = { /* 2 nodes */ - { /* Tree (LL/LL128/Simple)*/ { 0.06, 0.00, 0.11 }, /* Ring (LL/LL128/Simple)*/ { 0.08, 0.00, 1.00 }, /* CollNetDirect (Simple)*/ { 0.00, 0.00, 1.00 }, /* CollNetChain (Simple)*/ { 0.00, 0.00, 1.00 }, /* NVLS */ { 0, 0, 0 }, /* NVLS Tree */ { 0, 0, 0 } }, + { /* Tree (LL/LL128/Simple)*/ { 0.06, 0.06, 0.11 }, /* Ring (LL/LL128/Simple)*/ { 0.08, 0.08, 1.00 }, /* CollNetDirect (Simple)*/ { 0.00, 0.00, 1.00 }, /* CollNetChain (Simple)*/ { 0.00, 0.00, 1.00 }, /* NVLS */ { 0, 0, 0 }, /* NVLS Tree */ { 0, 0, 0 } }, /* more than 2 nodes */ - { /* Tree (LL/LL128/Simple)*/ { 0.06, 0.00, 0.59 }, /* Ring (LL/LL128/Simple)*/ { 0.08, 0.00, 1.00 }, /* CollNetDirect (Simple)*/ { 0.00, 0.00, 1.00 }, /* CollNetChain (Simple)*/ { 0.00, 0.00, 1.00 }, /* NVLS */ { 0, 0, 0 }, /* NVLS Tree */ { 0, 0, 0 } }, + { /* Tree (LL/LL128/Simple)*/ { 0.06, 0.06, 0.59 }, /* Ring (LL/LL128/Simple)*/ { 0.08, 0.08, 1.00 }, /* CollNetDirect (Simple)*/ { 0.00, 0.00, 1.00 }, /* CollNetChain (Simple)*/ { 0.00, 0.00, 1.00 }, /* NVLS */ { 0, 0, 0 }, /* NVLS Tree */ { 0, 0, 0 } }, }, .treeCorrectionFactor = { { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 1.0, 1.0, 1.0, 1.0, 1.0, 0.6, 1.0, 0.9, 1.0, 1.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, }, - { 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, }, + { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 1.0, 1.0, 1.0, 1.0, 1.0, 0.6, 1.0, 0.9, 1.0, 1.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, }, { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.7, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.8, }, }, .ringCorrectionFactor = { { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 1.0, 0.2, 1.0, 0.4, 0.4, 0.1, 0.2, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, }, - { 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, }, + { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 1.0, 0.2, 1.0, 0.4, 0.4, 0.1, 0.2, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, }, { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.8, 1.0, 1.0, 1.0, }, }, }; @@ -504,12 +504,13 @@ ncclResult_t ncclTopoTuneModel(struct ncclComm* comm, int minCompCap, int maxCom for (int c=0; ctopo->nodes[GPU].nodes[0].gpu.gcn, "gfx12")) ? 0 : protoEnable[p]; - if (pEnable == 2 && p == NCCL_PROTO_LL128) { + if (p == NCCL_PROTO_LL128) { #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) #if defined(ENABLE_LL128) // Enable LL128 by default only on gfx90a with available tuning table pEnable = (graphs[a]->typeInter <= PATH_PXB) && graphs[a]->typeIntra <= PATH_NVL && - (IsArchMatch(comm->topo->nodes[GPU].nodes[0].gpu.gcn, "gfx90a") && comm->topo->ll128Enabled) ? 1 : 0; + ((IsArchMatch(comm->topo->nodes[GPU].nodes[0].gpu.gcn, "gfx90a") || + IsArchMatch(comm->topo->nodes[GPU].nodes[0].gpu.gcn, "gfx942")) && comm->topo->ll128Enabled) ? 1 : 0; #else pEnable = 0; #endif