fix(rccl): fit AINIC scatter parts to the P2P channel pool

projects/rccl/CHANGELOG.md

@@ -56,7 +56,7 @@ Full documentation for RCCL is available at [https://rccl.readthedocs.io](https:
 * Fixed virtual address space exhaustion when symmetric windows backed by the same physical allocation are registered repeatedly.
 * Fixed Copy Engine `ncclAlltoAllv` hanging on intra-node communicators of more than 8 ranks with uneven send sizes. Each rank decided from its own send sizes whether to add intra-batch barriers, so ranks could issue different barrier counts. The decision now uses the full send-size matrix, which every rank shares.
 * Fixed `NCCL_P2P_DISABLE` and `NCCL_P2P_LEVEL` being ignored on Intel, ARM and Zhaoxin hosts, where the CPU-based default P2P level overwrote the user setting. On such hosts `NCCL_P2P_DISABLE=1` now falls back to SHM as documented, and `NCCL_P2P_LEVEL` is honored.
-* Fixed a 15 to 30% `ncclScatter` bandwidth loss from 8 MB upward on multi-node AINIC jobs of up to 8 nodes, caused by the move to two network channels per peer. Scatter now uses and connects one channel per peer when the peers cannot each get their own channels. Point-to-point operations and the other collectives keep two channels.
+* Fixed a 15 to 30% `ncclScatter` bandwidth loss from 16 MB upward on multi-node AINIC jobs, caused by the move to two network channels per peer (for example on 4 nodes of 8 GPUs). Scatter now uses and connects the most channels per peer whose flows still fit the P2P channel pool. Point-to-point operations and the other collectives keep their channel count.
 
 ### Known issues
 * The FP8 ReduceSum and ReduceCopy device APIs are not exported in the LLVM bitcode library.

projects/rccl/src/enqueue/enqueue.cc

@@ -1326,19 +1326,16 @@ static bool rcclP2pBatchEligible(struct ncclComm* comm, ssize_t sendBytes, ssize
                               rcclParamP2pBatchThreshold());
 }
 
-// Number of the p2pnChannelsPerPeer parts that a task of collAPI uses and connects.
-// On AINIC ncclTopoComputeP2pChannels raises p2pnChannelsPerPeer to 2 so one peer can saturate a NIC.
-// Scatter already drives all nRanks-1 peers at once from one root and loses 15-30% from 8 MB up with
-// two parts per peer, so it uses only part 0 once the peers cannot each own p2pnChannelsPerPeer
-// channels. It must not connect part 1 either: when scatter connects both parts itself, it stays at
-// or below the two-part speed even though it sends on part 0 only. Part 1 connected later by another
-// collective does not slow it down. The answer may depend only on values both ends of a pair share
-// (collAPI, comm-wide sizes), never on per-rank task counts, or sender and receiver would disagree on
-// the parts.
+// Number of the p2pnChannelsPerPeer parts a task of collAPI uses and connects. Multi-node AINIC scatter
+// uses and connects only the most parts whose flows from the root still fit the pool; connecting more
+// slows it down. A user-set nChannelsPerNetPeer is kept, as in ncclTopoComputeP2pChannels. Depend only
+// on values both ends of a pair share, or they disagree on the parts.
 int rcclP2pTaskParts(struct ncclComm* comm, ncclFunc_t collAPI) {
-  if (collAPI == ncclFuncScatter && comm->nNodes > 1 && comm->p2pnChannelsPerPeer > 1 && rcclUseAinic() &&
-      (comm->nRanks - 1) * comm->p2pnChannelsPerPeer > comm->p2pnChannels)
-    return 1;
+  if (collAPI == ncclFuncScatter && comm->nNodes > 1 && comm->config.nChannelsPerNetPeer == NCCL_CONFIG_UNDEF_INT &&
+      rcclUseAinic()) {
+    int partsThatFit = comm->p2pnChannels / (comm->nRanks - 1);
+    return std::min(comm->p2pnChannelsPerPeer, partsThatFit >= 1 ? pow2Down(partsThatFit) : 1);
+  }
   return comm->p2pnChannelsPerPeer;
 }
 
@@ -1348,16 +1345,17 @@ static int rcclP2pPolicyChannels(struct ncclComm* comm, struct ncclTaskP2p* task
     ncclFunc_t collAPI;
     int nChannels;
   };
-  // This workaround targets direct P2P/IPC traffic only. When P2P is disabled,
-  // preserve the normal channel selection for the fallback transport.
   if (task == nullptr) return -1;
+  // The Scatter part limit is about network flows, so it applies with P2P disabled too.
   int taskParts = rcclP2pTaskParts(comm, task->collAPI);
   if (taskParts < comm->p2pnChannelsPerPeer) return taskParts;
+  // The policies below target direct P2P/IPC traffic only. When P2P is disabled,
+  // preserve the normal channel selection for the fallback transport.
   if (ncclParamP2pDisable()) return -1;
   constexpr Policy policies[] = {
     // Two-channel full-mesh traffic has high variance on gfx110x. Keep the
-    // workaround specific to AllToAll so Gather, Scatter, and SendRecv retain
-    // their higher-throughput multi-channel paths.
+    // workaround specific to AllToAll so Gather, SendRecv and Scatter (outside the
+    // AINIC part limit above) retain their higher-throughput multi-channel paths.
     {"gfx110", ncclFuncAlltoAll, 1},
   };
   for (const auto& policy : policies) {

projects/rccl/src/enqueue/task_prep/task_posttuning.cc

@@ -245,8 +245,7 @@ static ncclResult_t postTuneP2pRegisterBuffer(struct ncclComm* comm, struct nccl
   if (network) {
     bool pxnUsed = !ncclPxnDisable(comm) && comm->isAllNvlink && comm->maxLocalRanks > 1;
     if (proxySameProcess && !pxnUsed && (conn->conn.flags & NCCL_DIRECT_NIC)) {
-      int nParts = rcclP2pTaskParts(comm, task->collAPI);
-      for (int part = 0; part < nParts; part++) {
+      for (int part = 0; part < comm->p2pnChannelsPerPeer; part++) {
         int partChannelId = ncclP2pChannelForPart(comm->p2pnChannels, base, part, comm->p2pnChannelsPerPeer,
                                                   comm->nNodes, comm->p2pChannelShiftSize);
         struct ncclConnector* partConn = isSendNotRecv ? &comm->channels[partChannelId].peers[peer]->send[connIndex] :

projects/rccl/test/host/enqueue-test.cc

@@ -339,6 +339,68 @@ TEST_F(EnqueueMicrotest, ShmemDynamicSize_DeviceLinker_IsAlwaysZero) {
 }
 #endif
 
+// ===========================================================================
+// rcclP2pTaskParts (enqueue.cc:1333) and rcclP2pPolicyChannels (enqueue.cc:1342)
+//
+// Multi-node AINIC Scatter takes the most power-of-two parts whose nRanks - 1
+// flows fit p2pnChannels, at most p2pnChannelsPerPeer; everything else takes
+// p2pnChannelsPerPeer. Both ends of a pair must get the same answer, so it may
+// depend only on comm-wide fields.
+// ===========================================================================
+
+namespace {
+std::unique_ptr<ncclComm> P2pPartsComm(int nNodes, int nRanks, int p2pnChannels, int perPeer) {
+  std::unique_ptr<ncclComm> comm{new ncclComm{}};
+  comm->nNodes = nNodes;
+  comm->nRanks = nRanks;
+  comm->p2pnChannels = p2pnChannels;
+  comm->p2pnChannelsPerPeer = perPeer;
+  comm->config.nChannelsPerNetPeer = NCCL_CONFIG_UNDEF_INT;
+  return comm;
+}
+}  // namespace
+
+TEST_F(EnqueueMicrotest, P2pTaskParts_AinicScatterTakesThePartsWhoseFlowsFitThePool) {
+  g_rcclUseAinicValue = true;
+  struct Case {
+    int nNodes, nRanks, p2pnChannels, perPeer, scatterParts;
+  };
+  // 8 GPUs per node with the MI355X pools, the pool boundary (16 peers x 2 parts fill 32 channels,
+  // 17 peers do not), power-of-two rounding (64 / 19 = 3) and the one-part floor (64 / 127 = 0).
+  for (Case k : {Case{2, 16, 32, 2, 2}, Case{3, 24, 64, 4, 2}, Case{4, 32, 32, 2, 1}, Case{8, 64, 64, 2, 1},
+                 Case{2, 17, 32, 2, 2}, Case{2, 18, 32, 2, 1}, Case{3, 20, 64, 8, 2}, Case{16, 128, 64, 2, 1}}) {
+    auto comm = P2pPartsComm(k.nNodes, k.nRanks, k.p2pnChannels, k.perPeer);
+    EXPECT_EQ(k.scatterParts, rcclP2pTaskParts(comm.get(), ncclFuncScatter))
+      << k.nRanks << " ranks, " << k.p2pnChannels << " channels";
+    for (ncclFunc_t f : {ncclFuncSendRecv, ncclFuncSend, ncclFuncRecv, ncclFuncAlltoAll, ncclFuncGather,
+                         ncclFuncAlltoAllv}) {
+      EXPECT_EQ(k.perPeer, rcclP2pTaskParts(comm.get(), f)) << "collAPI=" << int(f);
+    }
+  }
+}
+
+TEST_F(EnqueueMicrotest, P2pTaskParts_ScatterKeepsPerPeerOutsideMultiNodeAinicOrWithAUserCount) {
+  auto comm = P2pPartsComm(/*nNodes=*/4, /*nRanks=*/32, /*p2pnChannels=*/32, /*perPeer=*/2);
+  g_rcclUseAinicValue = false;
+  EXPECT_EQ(2, rcclP2pTaskParts(comm.get(), ncclFuncScatter)) << "not AINIC";
+  g_rcclUseAinicValue = true;
+  comm->config.nChannelsPerNetPeer = 2;
+  EXPECT_EQ(2, rcclP2pTaskParts(comm.get(), ncclFuncScatter)) << "NCCL_NCHANNELS_PER_NET_PEER set";
+  comm->config.nChannelsPerNetPeer = NCCL_CONFIG_UNDEF_INT;
+  comm->nNodes = 1;
+  EXPECT_EQ(2, rcclP2pTaskParts(comm.get(), ncclFuncScatter)) << "single node";
+}
+
+TEST_F(EnqueueMicrotest, P2pPolicyChannels_ScatterBelowPerPeerReturnsItsParts) {
+  // Returns before the arch policy table, so comm->topo is never read.
+  g_rcclUseAinicValue = true;
+  auto comm = P2pPartsComm(/*nNodes=*/4, /*nRanks=*/32, /*p2pnChannels=*/32, /*perPeer=*/2);
+  ncclTaskP2p task{};
+  task.collAPI = ncclFuncScatter;
+  EXPECT_EQ(1, rcclP2pPolicyChannels(comm.get(), &task));
+  EXPECT_EQ(-1, rcclP2pPolicyChannels(comm.get(), nullptr));
+}
+
 // ===========================================================================
 // calcP2pChannelCount (enqueue.cc:1550)
 //

projects/rccl/test/host/fakes/enqueue_fakes.cc

@@ -17,7 +17,7 @@
 
 #include "signature-drift.h"
 
-#include "channel.h"            // rcclEffectiveP2pBatchEnable
+#include "channel.h"            // rcclEffectiveP2pBatchEnable, rcclP2pTaskParts
 #include "enqueue_fakes.h"
 #include "nccl_fakes.h"         // g_loadParam, for the NCCL_PARAM body this file stands in for
 #include "sym_kernels_fakes.h"  // g_symkAvailable and the sym_kernels.cc seams below's canonical home
@@ -31,6 +31,7 @@ ASSERT_HOOK_MATCHES_PROD(g_addProxyOpIfNeeded, ncclAddProxyOpIfNeeded);
 ASSERT_HOOK_MATCHES_PROD(g_getCollNetSupport, ncclGetCollNetSupport);
 ASSERT_HOOK_MATCHES_PROD(g_getRegBuff, ncclGetRegBuff);
 ASSERT_HOOK_MATCHES_PROD(g_rcclEffectiveP2pBatchEnable, rcclEffectiveP2pBatchEnable);
+ASSERT_HOOK_MATCHES_PROD(g_rcclP2pTaskParts, rcclP2pTaskParts);
 #undef ASSERT_HOOK_MATCHES_PROD
 
 // Generous default: a deny-everything default would make even a single small task look over budget.
@@ -75,6 +76,10 @@ std::function<ncclResult_t(struct ncclComm*, struct ncclKernelPlan*, struct nccl
 static int DefaultEffectiveP2pBatchEnable(struct ncclComm*) { return 0; }
 std::function<int(struct ncclComm*)> g_rcclEffectiveP2pBatchEnable = DefaultEffectiveP2pBatchEnable;
 
+// enqueue.cc's answer for every task outside multi-node AINIC Scatter.
+static int DefaultP2pTaskParts(struct ncclComm* comm, ncclFunc_t) { return comm->p2pnChannelsPerPeer; }
+std::function<int(struct ncclComm*, ncclFunc_t)> g_rcclP2pTaskParts = DefaultP2pTaskParts;
+
 // enqueue.cc:437-441 splits the default on ROCm version; mirrored rather than pinned to one arm.
 #if ROCM_VERSION >= 71200
 int64_t ncclParamGraphRegister() { return g_loadParam("GRAPH_REGISTER", 1); }
@@ -105,11 +110,15 @@ void ResetEnqueueFakes() {
   g_getCollNetSupport = DefaultGetCollNetSupport;
   g_getRegBuff = DefaultGetRegBuff;
   g_rcclEffectiveP2pBatchEnable = DefaultEffectiveP2pBatchEnable;
+  g_rcclP2pTaskParts = DefaultP2pTaskParts;
 }
 
 int rcclEffectiveP2pBatchEnable(struct ncclComm* comm) {
   return g_rcclEffectiveP2pBatchEnable(comm);
 }
+int rcclP2pTaskParts(struct ncclComm* comm, ncclFunc_t collAPI) {
+  return g_rcclP2pTaskParts(comm, collAPI);
+}
 
 // src/enqueue/enqueue.cc
 ncclResult_t ncclGetAlgoInfo(struct ncclComm* comm, struct ncclTaskColl* task, int collNetSupport, int nvlsSupport,

projects/rccl/test/host/fakes/enqueue_fakes.h

@@ -52,6 +52,9 @@ extern std::function<ncclResult_t(struct ncclComm*, struct ncclTaskColl*, int*)>
 // enqueue.cc's rcclEffectiveP2pBatchEnable: default 0 is batching off, the single-node answer.
 extern std::function<int(struct ncclComm*)> g_rcclEffectiveP2pBatchEnable;
 
+// enqueue.cc's rcclP2pTaskParts: default p2pnChannelsPerPeer; install the multi-node AINIC Scatter limit per test.
+extern std::function<int(struct ncclComm*, ncclFunc_t)> g_rcclP2pTaskParts;
+
 void ResetEnqueueFakes();
 
 #endif  // RCCL_TEST_HOST_FAKES_ENQUEUE_FAKES_H_