feat(rccl): stage symmetric kernel tiles with gfx1250 TDM instead of TMA

projects/rccl/CHANGELOG.md

@@ -21,6 +21,7 @@ Full documentation for RCCL is available at [https://rccl.readthedocs.io](https:
 * Multiple GIN proxy progress threads via `NCCL_GIN_PROXY_NTHREADS`.
 * Enabled hierarchical Copy Engine `ncclAllGather` and `ncclAlltoAll` on multi-node communicators: the intra-node phase runs over SDMA/CE across the LSA team and the inter-node phase over one-sided RMA through the CPU (GIN) proxy. The path is off by default and takes two opt-ins: `NCCL_CTA_POLICY=2` (equivalently `config.CTAPolicy = NCCL_CTA_POLICY_ZERO`), and symmetric registration of both the send and the receive buffer through `ncclCommWindowRegister`. On architectures other than gfx1250 (MI450) a third is needed, `NCCL_CUMEM_ENABLE=1`, because symmetric memory depends on cuMem and cuMem auto-detection enables itself only on gfx1250. `NCCL_GIN_ENABLE`, `NCCL_NUM_RMA_CTX`, `NCCL_WIN_ENABLE` and `NCCL_DMABUF_ENABLE` need no change, since their defaults already qualify; the inter-node rail draws from the internal RMA context pool (`NCCL_NUM_RMA_INT_CTX`, default 4) rather than the user-addressable contexts, so `numRmaCtx` only has to stay above zero. Selection also requires an available GIN backend and the same number of ranks on every node; communicators that do not qualify keep using the kernel path. Validated for correctness and path selection from 1 to 32 nodes.
 * Added `ncclGinFenceLevel` barrier semantics on the Anvil SDMA GIN backend (`NCCL_GIN_TYPE=7`). `gin.get`, `flushAsync` and `wait` are implemented instead of trapping, and a barrier's signal is now ordered behind the payload on the peer's SDMA queue, so `Put`, `Get` and the default `Put | Get` fences drain what they promise. `Put` covers a rank's puts to itself, and `ncclGinAllContexts` fences every GIN context. Single-node validated on gfx950.
+* Added gfx1250 (MI450) Tensor Data Mover staging to the symmetric AllGather, AllReduce and ReduceScatter kernels, where NCCL uses NVIDIA's TMA. Their deep loops move each tile through the warp's shared-memory scratch with the DMA engine instead of per-lane vector loads; on ROCm that path previously compiled out entirely, so the `Tma` kernels were never generated. Launch tuning was also corrected so these kernels can use up to a quarter of a GPU's WGPs, where only 14 were reached before, which is where most of the large-message gain comes from. Opt-in via `NCCL_SYM_TMA_ENABLE=1`, and reached only when the application opts into symmetric memory. `NCCL_SYM_TMA_ENABLE=2` additionally forces these kernels wherever the collective has one, skipping both the tuner's cost comparison and the deep-loop size bar, which is intended for A/B measurement rather than production. The new kernels are generated only when gfx1250 is in `GPU_TARGETS`, and every other architecture keeps the vector path unchanged.
 * Removed GIN rocSHMEM GDA dependency on rocSHMEM GDA bitcode: GIN rocSHMEM GDA is now header-only. The CMake `roc::rccl` target provides the required include directories in its `INSTALL_INTERFACE` property.
 
 ### Changed
@@ -52,7 +53,7 @@ Full documentation for RCCL is available at [https://rccl.readthedocs.io](https:
 * Fixed hierarchical Copy Engine collectives dereferencing a NULL RMA proxy context when `NCCL_RMA_DISABLE=1` was combined with the hierarchical opt-ins. `NCCL_RMA_DISABLE` suppresses the window registration that creates those contexts, but the hierarchical availability check did not consult it, so a multi-node communicator with `NCCL_CTA_POLICY=2` and symmetric send and receive windows admitted the path and then crashed at launch. The internal RMA context predicate now defers to `ncclRmaProxyEnabled`, so such a communicator declines the hierarchical path and uses the kernel path instead.
 
 ### Known issues
-* The following NCCL 2.31 features are NVIDIA-specific and are not available in RCCL: Compute Fabric Transport, the GDAKI and EFA GDA GIN backends, PAT combined with NVLS, NVLink-multicast AllGather, TMA-based symmetric kernels, and the CuTeDSL and `nccl4rust` bindings.
+* The following NCCL 2.31 features are NVIDIA-specific and are not available in RCCL: Compute Fabric Transport, the GDAKI and EFA GDA GIN backends, PAT combined with NVLS, NVLink-multicast AllGather, and the CuTeDSL and `nccl4rust` bindings.
 
 ## RCCL 2.30.7 for ROCm 10.1.0 (Unreleased)
 

projects/rccl/src/CMakeLists.txt

@@ -855,9 +855,11 @@ if (gen_py_result)
 endif()
 
 if (GENERATE_SYM_KERNELS)
-  # Execute the python script to generate required symmetric memory kernels
+  # Execute the python script to generate required symmetric memory kernels.
+  # GPU_TARGETS selects the arch-specific algos (the Tma* kernels need gfx1250's
+  # Tensor Data Mover).
   execute_process(
-    COMMAND ${Python3_EXECUTABLE} ${RCCL_SOURCE_DIR}/src/device/symmetric/generate.py ${GEN_SYM_DIR}
+    COMMAND ${Python3_EXECUTABLE} ${RCCL_SOURCE_DIR}/src/device/symmetric/generate.py ${GEN_SYM_DIR} "${GPU_TARGETS}"
     WORKING_DIRECTORY ${RCCL_SOURCE_DIR}
     RESULT_VARIABLE gen_sym_py_result
     ERROR_VARIABLE gen_sym_py_error

projects/rccl/src/device/symmetric/all_gather.cuh

@@ -14,7 +14,9 @@
 #include "primitives.cuh"
 #endif
 
-template <int BytePerPack, int UnrollPacks, int UnrollPeers, bool EnableTma>
+// TileAligned: the caller's tier guarantees every tile address is ncclSymkTileLine-aligned, which
+// lets the staging wrappers drop TDM's head peel.
+template <int BytePerPack, int UnrollPacks, int UnrollPeers, bool EnableTma, bool TileAligned = false>
 static __device__ void bcastDeep(ncclSymkArgsHandler const& handler, int tn, int t, bool waitNeeded,
                                  ncclLsaBarrierSession<ncclCoopCta>& bar, ncclSymPtr<char> input,
                                  ncclSymPtr<char> output, bool inPlace, int nIters) {
@@ -27,52 +29,44 @@ static __device__ void bcastDeep(ncclSymkArgsHandler const& handler, int tn, int
 
   Pack* inpPacks = (Pack*)input.localPtr() + intptr_t(w) * UnrollPacks * WARP_SIZE +
                    (
-#if __CUDA_ARCH__ >= 1000
+#if NCCL_SYMK_ASYNC_TILE
                      EnableTma ? 0 :
 #endif
                                  lane);
 
   ncclSymPtr<Pack> outPacks = (ncclSymPtr<Pack>)output + intptr_t(w) * UnrollPacks * WARP_SIZE +
                               (
-#if __CUDA_ARCH__ >= 1000
+#if NCCL_SYMK_ASYNC_TILE
                                 EnableTma ? 0 :
 #endif
                                             lane);
 
   Pack tmp[UnrollPacks];
 
-#if __CUDA_ARCH__ >= 1000
+#if NCCL_SYMK_ASYNC_TILE
   int lw = threadIdx.x / WARP_SIZE;
-  extern __shared__ char smemScratch[];
   using tmaSmemStruct_t = tmaSmemStruct<Pack, UnrollPacks>;
-  constexpr int smemSizePerWarp = ncclTmaShmemScratchWarpSize();
-  tmaSmemStruct_t* tmaSmem = reinterpret_cast<tmaSmemStruct_t*>(smemScratch + lw * smemSizePerWarp);
+  tmaSmemStruct_t* tmaSmem = ncclSymkTileSmem<tmaSmemStruct_t>(lw);
   constexpr size_t tileSize = UnrollPacks * WARP_SIZE * BytePerPack;
 #endif
   bool skip = false; // all lanes issue loads/stores
 
-#if __CUDA_ARCH__ >= 1000
+#if NCCL_SYMK_ASYNC_TILE
+  size_t pending = 0;
   if NCCL_IF_CONSTEXPR (EnableTma) {
-    if (lane == 0) {
-      // lane0 issues async.cp.bulk commands
-      init(&tmaSmem->bar, 1);
-    } else {
-      // other lanes can skip the loop
-      skip = true;
-    }
+    ncclSymkTileBarInit(&tmaSmem->bar, /*arrivers=*/1, lane);
+    // TMA issues from lane 0, so the rest of the warp has nothing left to do.
+    // TDM needs the whole warp to build the descriptor, so every lane stays in.
+    skip = NCCL_SYMK_TILE_TMA && lane != 0;
   }
 #endif
 
   nIters -= w;
   if (0 < nIters) {
-#if __CUDA_ARCH__ >= 1000
+#if NCCL_SYMK_ASYNC_TILE
     if NCCL_IF_CONSTEXPR (EnableTma) {
-      if (lane == 0) {
-        cuda::device::memcpy_async_tx(tmaSmem->buff[0], inpPacks, cuda::aligned_size_t<16>(tileSize), tmaSmem->bar);
-        cuda::barrier<cuda::thread_scope_block>::arrival_token token =
-          cuda::device::barrier_arrive_tx(tmaSmem->bar, 1, tileSize);
-        tmaSmem->bar.wait(std::move(token));
-      }
+      ncclSymkTileLoad<TileAligned>(tmaSmem->buff[0], inpPacks, tileSize, tmaSmem->bar, pending, lane);
+      ncclSymkTileLoadWait</*Arrivers=*/1>(tmaSmem->bar, pending, lane);
     } else
 #endif
     {
@@ -97,9 +91,9 @@ static __device__ void bcastDeep(ncclSymkArgsHandler const& handler, int tn, int
           NVCC_PRAGMA_UNROLL_AUTO
           for (int ur = 0; ur < UnrollPeers - partial; ur++) {
             if (partial && dr == nRanks) break;
-#if __CUDA_ARCH__ >= 1000
+#if NCCL_SYMK_ASYNC_TILE
             if NCCL_IF_CONSTEXPR (EnableTma) {
-              ptx::cp_async_bulk(ptx::space_global, ptx::space_shared, outPacks.lsaPtr(r), tmaSmem->buff[0], tileSize);
+              ncclSymkTileStore<TileAligned>(outPacks.lsaPtr(r), tmaSmem->buff[0], tileSize, lane);
             } else
 #endif
             {
@@ -110,12 +104,9 @@ static __device__ void bcastDeep(ncclSymkArgsHandler const& handler, int tn, int
             }
             if (++r == nRanks) r = 0;
           }
-#if __CUDA_ARCH__ >= 1000
+#if NCCL_SYMK_ASYNC_TILE
           if NCCL_IF_CONSTEXPR (EnableTma) {
-            if (lane == 0) {
-              ptx::cp_async_bulk_commit_group();
-              ptx::cp_async_bulk_wait_group_read(ptx::n32_t<0>());
-            }
+            ncclSymkTileStoreWait(lane);
           }
 #endif
         }
@@ -124,14 +115,10 @@ static __device__ void bcastDeep(ncclSymkArgsHandler const& handler, int tn, int
       outPacks += intptr_t(wn) * UnrollPacks * WARP_SIZE;
       nIters -= wn;
       if (nIters <= 0) break;
-#if __CUDA_ARCH__ >= 1000
+#if NCCL_SYMK_ASYNC_TILE
       if NCCL_IF_CONSTEXPR (EnableTma) {
-        if (lane == 0) {
-          cuda::device::memcpy_async_tx(tmaSmem->buff[0], inpPacks, cuda::aligned_size_t<16>(tileSize), tmaSmem->bar);
-          cuda::barrier<cuda::thread_scope_block>::arrival_token token =
-            cuda::device::barrier_arrive_tx(tmaSmem->bar, 1, tileSize);
-          tmaSmem->bar.wait(std::move(token));
-        }
+        ncclSymkTileLoad<TileAligned>(tmaSmem->buff[0], inpPacks, tileSize, tmaSmem->bar, pending, lane);
+        ncclSymkTileLoadWait</*Arrivers=*/1>(tmaSmem->bar, pending, lane);
       } else
 #endif
       {
@@ -185,15 +172,15 @@ static __device__ void bcast(ncclSymkArgsHandler const& handler, int tn, int t,
 
   uint32_t alignment = uint32_t(input.offset - output.offset);
   uint32_t nPreBytes =
-#if __CUDA_ARCH__ >= 1000
+#if NCCL_SYMK_ASYNC_TILE
     (EnableTma && alignment % 256 == 0) ? (256 - input.offset) % 256 :
 #endif
                                           (16 - input.offset) % 16;
 
   nPreBytes = min((size_t)nPreBytes, nBytes);
   uintptr_t cursor = nPreBytes;
 
-#if __CUDA_ARCH__ >= 1000
+#if NCCL_SYMK_ASYNC_TILE
   if NCCL_IF_CONSTEXPR (EnableTma) {
     if (alignment % 256 == 0) {
       constexpr int BytePerPack = ncclSymkBytePerPack, UnrollPacks = ncclSymkAlign256BDeepUnrollPacks, UnrollPeers = 2;
@@ -202,10 +189,12 @@ static __device__ void bcast(ncclSymkArgsHandler const& handler, int tn, int t,
       chunks -= imodFast32(chunks, nBlocks, nBlocks_rcp32);
       if (chunks != 0) {
         uintptr_t cursorAfter = cursor + uintptr_t(chunks) * BytePerChunk;
-        bcastDeep<BytePerPack, UnrollPacks, UnrollPeers, EnableTma>(handler, tn, t, waitNeeded, bar,
-                                                                    (ncclSymPtr<char>)input + cursor,
-                                                                    (ncclSymPtr<char>)output + cursor, inPlace,
-                                                                    chunks * ncclSymkMinWarpsPerBlock);
+        // Both sides start on a 256 B boundary here (nPreBytes peels input to one, alignment % 256
+        // carries output with it, window bases are page-aligned) and every tile advances by a
+        // multiple of 256, so the tiles are ncclSymkTileLine-aligned and the peel can go.
+        bcastDeep<BytePerPack, UnrollPacks, UnrollPeers, EnableTma, /*TileAligned=*/true>(
+          handler, tn, t, waitNeeded, bar, (ncclSymPtr<char>)input + cursor,
+          (ncclSymPtr<char>)output + cursor, inPlace, chunks * ncclSymkMinWarpsPerBlock);
         cursor = cursorAfter;
         waitNeeded = false;
       }