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

projects/rccl/CHANGELOG.md

@@ -36,6 +36,7 @@ Full documentation for RCCL is available at [https://rccl.readthedocs.io](https:
 * The CSV tuner configurations are no longer installed to `${ROCM_PATH}/share/rccl/tuner/` or copied into the build tree; they ship inside `librccl.so`. A CSV placed in that directory, or in `<librccl.so dir>/tuner/`, still overrides the built-in defaults. A copy left behind by an earlier build or by `cmake --install` over an existing prefix also still takes precedence, so delete it to pick up the embedded defaults; package upgrades remove it automatically.
 * Retuned the symmetric AllReduce, ReduceScatter and AllGather kernels for gfx950. The block width is now selected per collective and message size instead of a fixed 256 threads, and the work partitioning and unroll factors in the load-direct kernels and the AllGather store kernel were refitted to the 64-lane wavefront. AllGather also switches from LL to the store kernel from 4 MB instead of 16 MB. The multi-node GIN kernels and all other architectures are unchanged. To fit the wider LL blocks, the shared LL scratch buffer on gfx950 grows from 4 MiB to 8 MiB per 8-rank communicator.
 * gfx1250 no longer forces unregistered `ncclAllGather`, `ncclAlltoAll`, `ncclGather` and `ncclScatter` through the Copy Engine scratch path (CE-Scratch), because it was slower than the kernel paths it replaced. Whether CE-Scratch is forced by default now comes from the per-architecture table: gfx942 and gfx950 still force it. A communicator-level `NCCL_CTA_POLICY_ZERO` still forces it on every architecture; an explicit `RCCL_FORCE_CE=1` still forces it and `RCCL_FORCE_CE=0` still turns it off. gfx1250 `ncclAllGather` now uses DDA LL128 up to 128 MiB.
+* 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.
 * gfx1250 `ncclAllReduce` on single-node communicators keeps the LL128 protocol at larger sizes: Tree at every size, Ring up to nRanks x 32 MiB, and the DDA LL128 tier up to 256 MiB. The DDA LL128 two-shot ceiling now follows that tier ceiling (`RCCL_DDA_LL128_TWOSHOT_THRESHOLD` still overrides it); the DDA scratch slot lets two-shot reach 256 MiB from 32 ranks, about 240 MiB at 16 and 128 MiB at 2 to 8. On gfx1250 communicators of 2 to 8 ranks, the DDA fabric scratch buffer grows to about 546 MiB.
 * On gfx1250 single-node communicators of 8 or more ranks, the DDA tier, algorithm and protocol crossovers now scale with the rank count. Communicators below 8 ranks keep the previous thresholds.
   * `ncclAllReduce`: DDA LL two-shot up to nRanks x 1 MiB (`RCCL_DDA_LL_TWOSHOT_THRESHOLD` still overrides it); from 32 ranks, DDA VMM between 2 MiB and nRanks x 512 KiB for unregistered buffers unless `RCCL_DDA_THRESHOLD` or `RCCL_DDA_LL_THRESHOLD` is set; Tree instead of Ring up to nRanks x 1 MiB unless `NCCL_ALGO` or `NCCL_PROTO` is set, with LL up to nRanks x 32 KiB.

projects/rccl/src/ce_coll.cc

@@ -3610,6 +3610,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

@@ -4517,6 +4517,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,
@@ -4724,6 +4731,20 @@ 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.
+      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->decisionValid) {
         // AllReduce's backend was already chosen once by rcclSelectAllReduce();
         // honor it here instead of recomputing CE eligibility. rcclSelectAllReduce
@@ -4761,17 +4782,15 @@ static ncclResult_t taskAppend(struct ncclComm* comm, struct ncclInfo* info) {
         // ncclMakeSymmetricTaskList (collTaskAppend sets symkExtract from decision).
         INFO(NCCL_INIT, "Taking kernel-based collective path for ReduceScatter");
         NCCLCHECK(collTaskAppend(comm, info, 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 &&
+                 ((!rcclNcclAlgoEnvIsSet() && rcclForceCeScratchEnabled(comm)) || a2aCeScratchInPlace) &&
+                 CeScratchAvailable && !hasSysmemSegment &&
                  comm->ddaScratch != nullptr && recvBytes <= comm->ddaScratchBytes && info->coll != ncclFuncAllReduce) {
         INFO(NCCL_TUNING, "Using DDA scratch for CE collective, count=%zu, recvBytes=%zu", info->count, recvBytes);
         NCCLCHECK(ceCollTaskAppend(comm, info, /*sendWin=*/nullptr, /*recvWin=*/nullptr, comm->ddaScratch,
@@ -4794,6 +4813,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

@@ -266,6 +266,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);
 RCCL_PARAM_DECLARE(CeArMaxMsgBytes);     // -1 = use ceNonRegMax[AR] (2-shot) from arch table
 RCCL_PARAM_DECLARE(CeArRegMaxMsgBytes);  // -1 = use ceRegMax[AR] (registered) from arch table
 RCCL_PARAM_DECLARE(CeArStagingBytes);    // -1 = use NCCL_CE_AR_STAGING_BYTES; sizes ceARTmpBuf, not the cap