[RCCL] gfx1250 AllToAll in place: never return wrong data

projects/rccl/CHANGELOG.md

@@ -9,6 +9,7 @@ Full documentation for RCCL is available at [https://rccl.readthedocs.io](https:
 
 ### Changed
 * On gfx1250, cuMem auto-detection (`NCCL_CUMEM_ENABLE=-2`, the default) no longer requires Linux kernel 6.8 or later; it enables cuMem whenever the ROCm version, VMM support and runtime probe checks pass. The runtime probe now also exports a DMA-BUF, so hosts that cannot export one keep cuMem off. Kernel 6.8 or later is still recommended for performance, and RCCL logs a warning when cuMem runs on an older kernel. Set `NCCL_CUMEM_ENABLE=0` to keep cuMem off.
+* gfx1250 in-place `ncclAlltoAll` and `ncclAlltoAllConfig` (`sendbuff == recvbuff`) no longer fall back to per-peer Send/Recv, which returned wrong data in place. A call that the DDA fabric LL / LL128 paths do not take now runs CE-Scratch when it fits the DDA fabric scratch, and otherwise returns `ncclInvalidUsage` with a warning. This includes calls above the scratch, under graph capture, with window-registered buffers, with `NCCL_ENQUEUE_REARCH_ENABLE`, and every call on communicators without a DDA fabric scratch (children of `ncclCommSplit`, `ncclCommShrink` and `ncclCommGrow`, `NCCL_MNNVL_ENABLE=0`, multi-clique communicators, `RCCL_DDA_ENABLE=0`). `RCCL_FORCE_CE=0` now also turns off the in-place CE-Scratch route, so those calls return `ncclInvalidUsage`. On gfx1250 fabric communicators, a CE-Scratch AllToAll also zeroes the staged scratch and adds one Copy Engine barrier after the copy-back. Other architectures are unchanged.
 
 ### Removed
 

projects/rccl/src/ce_coll.cc

@@ -3889,6 +3889,17 @@ ncclResult_t ncclLaunchCeColl(struct ncclComm* comm, struct ncclKernelPlan* plan
       CUDACHECKGOTO(cudaMemcpyAsync(args->ddaUserRecvBuff, args->recvBuff /*scratch*/, fullBytes,
                                     cudaMemcpyDeviceToDevice, stream),
                     ret, fail);
+      // gfx1250 fabric comms: the DDA LL / LL128 kernels of every collective poll flag words in this
+      // same scratch with no entry barrier. Without the two steps below a peer's next LL kernel can
+      // write our scratch while this copy-back still reads it, or accept our leftover payload as a
+      // flag. Zero the staged range (0 is never a live LL epoch, ddaGetLLEpochInc skips it), then
+      // sync once more so every rank has finished its copy-back and its zeroing before any peer
+      // launches past this collective. CE-Scratch is never selected under graph capture
+      // (rcclSelectAlltoAll and taskAppend refuse it), so this path is eager only.
+      if (args->func == ncclFuncAlltoAll && comm->ddaFabricMemHandler != nullptr) {
+        CUDACHECKGOTO(cudaMemsetAsync(args->recvBuff /*scratch*/, 0, fullBytes, stream), ret, fail);
+        NCCLCHECKGOTO(ncclMemOpSync(comm, stream, args), ret, fail);
+      }
       break;
     }
   }

projects/rccl/src/enqueue/enqueue.cc

@@ -4566,6 +4566,13 @@ static ncclResult_t rawTaskAppend(struct ncclComm* comm, struct ncclInfo* info)
       }
     }
 
+    // AllToAll is always lowered to per-peer Send/Recv here (task_classify.cc), whatever
+    // rcclSelectAlltoAll chose, and that is wrong in place on gfx1250.
+    if (info->coll == ncclFuncAlltoAll) {
+      NCCLCHECK(rcclAlltoAllInPlaceGuard(comm, info->sendbuff, info->recvbuff,
+                                         (size_t)comm->nRanks * info->count * ncclTypeSize(info->datatype),
+                                         rcclA2aSiteRearch, /*capturing=*/false, /*registeredCe=*/false));
+    }
     NCCLCHECK(hostToDevRedOp(&opDev, info->op, info->datatype, comm));
     if (comm->nRanks == 1) {
       NCCLCHECK(ncclLaunchOneRank(info->recvbuff, info->sendbuff, info->count, opDev, info->datatype, info->stream,
@@ -4787,6 +4794,22 @@ static ncclResult_t taskAppend(struct ncclComm* comm, struct ncclInfo* info) {
          rcclAllGatherCeRegisteredWindow(comm, recvBytes, winRegType, ceCapturing));
       const bool allGatherDecided = (info->coll == ncclFuncAllGather && info->decisionValid);
       const bool alltoAllDecided = (info->coll == ncclFuncAlltoAll && info->decisionValid);
+      // An AllToAll without a decision (ncclAlltoAllConfig) gets the same in-place CE-Scratch exception
+      // as rcclSelectAlltoAll; the Send/Recv decomposition below refuses in place on gfx1250. So does a
+      // decided one whose decision is not CE: only a DDA fabric decision diverted here by NCCL_CHECK_MODE
+      // reaches this funnel in place (a live Direct decision in place was already refused by the selector).
+      const bool a2aCeScratchInPlace =
+        info->coll == ncclFuncAlltoAll && rcclAlltoAllCeScratchInPlace(comm, info->sendbuff, info->recvbuff);
+      // The undecided symmetric-window CE arm below. In place on gfx1250 an AllToAll must not take it
+      // (rcclAlltoAllRegisteredCeRefused); CE-Scratch refuses registered recv windows, so the call
+      // reaches the in-place guard before the Send/Recv decomposition.
+      const bool symWindowCe = !allGatherDecided && !alltoAllDecided &&
+                               (comm->config.CTAPolicy & NCCL_CTA_POLICY_ZERO) &&
+                               !(rcclNcclAlgoEnvIsSet() &&
+                                 (info->coll == ncclFuncAllGather || info->coll == ncclFuncReduceScatter)) &&
+                               (ceAvailable || hierCeAvailable) && !hasSysmemSegment;
+      const bool a2aRegCeRefused = symWindowCe && info->coll == ncclFuncAlltoAll &&
+                                   rcclAlltoAllRegisteredCeRefused(comm, info->sendbuff, info->recvbuff);
       if ((info->coll == ncclFuncAllReduce || info->coll == ncclFuncReduceScatter) && info->decisionValid) {
         // Backend was already chosen once by rcclSelectAllReduce() / rcclSelectReduceScatter();
         // honor it here instead of recomputing CE eligibility.
@@ -4819,17 +4842,15 @@ static ncclResult_t taskAppend(struct ncclComm* comm, struct ncclInfo* info) {
           NCCLCHECK(ceCollTaskAppend(comm, info, /*sendWin=*/nullptr, /*recvWin=*/nullptr,
                                      comm->ddaScratch, comm->ddaPeerPtrsHost, opDev));
         }
-      } else if ((!allGatherDecided && !alltoAllDecided &&
-                  (comm->config.CTAPolicy & NCCL_CTA_POLICY_ZERO) &&
-                  !(rcclNcclAlgoEnvIsSet() && (info->coll == ncclFuncAllGather || info->coll == ncclFuncReduceScatter)) &&
-                  (ceAvailable || hierCeAvailable) && !hasSysmemSegment) ||
-                 ceArSymRegistered || (!allGatherDecided && ceAgSymRegistered)) {
+      } else if ((symWindowCe && !a2aRegCeRefused) || ceArSymRegistered ||
+                 (!allGatherDecided && ceAgSymRegistered)) {
         INFO(NCCL_INIT, "Taking CE collective path with symmetric registered windows for user buffers");
         NCCLCHECK(ceCollTaskAppend(comm, info, sendWin, recvWin, /*ddaRecvBase=*/nullptr, /*ddaPeerBases=*/nullptr,
                                    opDev));
 
-      } else if (!rcclNcclAlgoEnvIsSet() && !allGatherDecided && !alltoAllDecided &&
-                 rcclForceCeScratchEnabled(comm) && CeScratchAvailable && !hasSysmemSegment &&
+      } else if (!allGatherDecided && (!alltoAllDecided || a2aCeScratchInPlace) &&
+                 ((!rcclNcclAlgoEnvIsSet() && rcclForceCeScratchEnabled(comm)) || a2aCeScratchInPlace) &&
+                 CeScratchAvailable && !hasSysmemSegment &&
                  comm->ddaScratch != nullptr && recvBytes <= comm->ddaScratchBytes && info->coll != ncclFuncAllReduce &&
                  info->coll != ncclFuncReduceScatter) {
         INFO(NCCL_TUNING, "Using DDA scratch for CE collective, count=%zu, recvBytes=%zu", info->count, recvBytes);
@@ -4853,6 +4874,9 @@ static ncclResult_t taskAppend(struct ncclComm* comm, struct ncclInfo* info) {
         NCCLCHECK(ncclCudaGetCapturingGraph(&graph, info->stream, comm->config.graphUsageMode));
         captured = ncclCudaGraphValid(graph);
         if (info->coll == ncclFuncAlltoAll) {
+          NCCLCHECK(rcclAlltoAllInPlaceGuard(comm, info->sendbuff, info->recvbuff, recvBytes,
+                                             alltoAllDecided ? rcclA2aSiteAlltoAllEnqueue : rcclA2aSiteAlltoAllConfig,
+                                             ceCapturing, a2aRegCeRefused));
           bool sendLocalValid = false;
           bool recvLocalValid = false;
           NCCLCHECK(ncclRegFind(comm, info->sendbuff, comm->nRanks * info->count * ncclTypeSize(info->datatype),

projects/rccl/src/include/rccl_common.h

@@ -290,6 +290,30 @@ bool rcclForceCeAllReduceEnabled(const struct ncclComm* comm);
 // RCCL_FORCE_CE=0 turns it off; otherwise NCCL_CTA_POLICY_ZERO forces it and
 // the arch table's forceCeScratch decides.
 bool rcclForceCeScratchEnabled(const struct ncclComm* comm);
+// gfx1250 in-place AllToAll. The per-peer Send/Recv decomposition is wrong in place there, so
+// every AllToAll funnel uses these two: rcclSelectAlltoAll, taskAppend for calls that carry no
+// decision (ncclAlltoAllConfig), and the NCCL_ENQUEUE_REARCH_ENABLE lowering.
+// - CeScratchInPlace: in place may take CE-Scratch (correct through ddaScratch) even though the
+//   arch table does not force it; RCCL_FORCE_CE=0 turns it off. Buffer-less queries are not in place.
+// - RegisteredCeRefused: in place must not take the registered-window CE AllToAll (single-node
+//   ncclCeAlltoAll or hierarchical CE). Both push each chunk into the peer's recv window while
+//   the peer may still read that region to send it, a race on user data. gfx1250 only (the same
+//   race exists upstream on every arch; other arches are left unchanged).
+// - InPlaceGuard: called where in place would otherwise reach Send/Recv; ncclInvalidUsage + WARN
+//   on gfx1250 (never wrong data), ncclSuccess otherwise. The WARN names the API (site) and why
+//   no correct path was left: capturing / registeredCe say what the caller saw; check mode, the
+//   RCCL_FORCE_CE setting and the CE scratch size are read from the comm.
+enum rcclA2aInPlaceSite {
+  rcclA2aSiteAlltoAll,         // rcclSelectAlltoAll (ncclAlltoAll)
+  rcclA2aSiteAlltoAllEnqueue,  // taskAppend, ncclAlltoAll carrying a non-CE decision
+  rcclA2aSiteAlltoAllConfig,   // taskAppend, no decision (ncclAlltoAllConfig)
+  rcclA2aSiteRearch,           // rawTaskAppend (NCCL_ENQUEUE_REARCH_ENABLE lowering)
+};
+bool rcclAlltoAllCeScratchInPlace(const struct ncclComm* comm, const void* sendbuff, const void* recvbuff);
+bool rcclAlltoAllRegisteredCeRefused(const struct ncclComm* comm, const void* sendbuff, const void* recvbuff);
+ncclResult_t rcclAlltoAllInPlaceGuard(const struct ncclComm* comm, const void* sendbuff, const void* recvbuff,
+                                      size_t totalBytes, enum rcclA2aInPlaceSite site, bool capturing,
+                                      bool registeredCe);
 // Sets ceAllReduceFastPath on each queued CE task from its receive window; no bootstrap.
 void rcclSetCeAllReduceFastPath(struct ncclComm* comm);
 RCCL_PARAM_DECLARE(CeArMaxMsgBytes);     // -1 = use ceNonRegMax[AR] (2-shot) from arch table