fix(rccl): back host-memory shared P2P pools with huge pages [ROCM-33021]

projects/rccl/CHANGELOG.md

@@ -72,6 +72,7 @@ Full documentation for RCCL is available at [https://rccl.readthedocs.io](https:
 * Fixed LL-protocol latency on gfx950 when RCCL is built against a ROCm older than 7.0.2 without the 128-bit global load/store builtins (for example ROCm 7.0.1, or `-DDWORDX4_INTRINSICS=OFF`). Every store to the user output buffer was followed by a release fence and an L2 writeback, which made 4 KB to 64 KB `ncclBroadcast` about 30% slower on MI355X. Builds against ROCm 7.0.2 or later were not affected and are unchanged.
 * 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.
 * Fixed intermittent data corruption in multi-node collectives, seen as out-of-bounds values in 8-node AllToAll, on hosts that use GPU Direct RDMA through peermem. Since RCCL 2.30.4 the GDR flush scratchpad was registered only through dma-buf, so on peermem hosts the flush degraded to a single RDMA read of the receive buffer, which does not order earlier relaxed-ordering writes. The scratchpad is now also registered through peermem, restoring the strictly ordered RDMA write that precedes the flush read.
+* Fixed `ibv_reg_mr_iova2 failed with error Invalid argument` on AMD AINIC when several communicators open P2P network connections without GPU Direct RDMA, for example Direct AllGather with `NCCL_PXN_DISABLE=1`. The host-memory shared P2P pools used 4 KiB pages, so each 256 MiB pool took 65,536 of the NIC's per-LIF page-table entries. They now use 2 MiB pages; `RCCL_NET_SHARED_HOST_HUGEPAGES` (default `1`) set to `0` restores the previous allocation.
 
 ### Known issues
 * The FP8 ReduceSum and ReduceCopy device APIs are not exported in the LLVM bitcode library.

projects/rccl/docs/api-reference/env-variables.rst

@@ -606,6 +606,16 @@ in the following table.
       - | Integer value in bytes (default: ``128``)
         | ``N``: Split only when message size >= N bytes
 
+    * - | ``RCCL_NET_SHARED_HOST_HUGEPAGES``
+        | Backs the host-memory shared P2P network staging pools (used by
+          connections without GPU Direct RDMA) with 2 MiB pages, so each
+          pool takes fewer NIC memory-registration page-table entries.
+          Falls back to the regular pinned host allocation if the huge-page
+          allocation or registration fails. Hosts with transparent huge
+          pages set to ``never`` get 4 KiB pages either way.
+      - | ``1``: Use 2 MiB pages (default).
+        | ``0``: Use the regular pinned host allocation.
+
     * - | ``NCCL_NCHANNELS_PER_NET_PEER``
         | Sets the number of channels used per network (remote) peer.
         | This overrides the value of the ``nChannelsPerNetPeer`` field in

projects/rccl/src/include/proxy.h

@@ -271,6 +271,8 @@ struct ncclProxySharedP2p {
   int64_t size;
   char* cudaBuff;
   char* hostBuff;
+  char* hostBuffDev; // Device-side address of hostBuff
+  size_t hostBuffMapSize; // Non-zero when hostBuff is a huge-page mmap registered with cudaHostRegister
   // CUDA IPC
   ncclIpcDesc ipcDesc;
   int dmaBufFd; // DMA-BUF fd for cuMem allocations

projects/rccl/src/transport/net.cc

@@ -860,6 +860,125 @@ static ncclResult_t recvFree(struct ncclComm* comm, struct ncclConnector* recv)
 }
 
 #define NCCL_SHARED_STEPS 16
+
+#if defined(NCCL_OS_LINUX)
+// Host-memory shared P2P pools are registered on every NIC that serves a non-GDR P2P connection.
+// With 4 KiB pages a 256 MiB pool takes 65,536 NIC page-table entries; AINIC caps them per LIF
+// (rdma_npts_per_lif = 1,048,576), so one communicator's pools nearly fill a NIC and the next
+// ibv_reg_mr fails with EINVAL. Backing the pool with transparent huge pages lets the MR use
+// 2 MiB pages (128 entries per 256 MiB).
+RCCL_PARAM(NetSharedHostHugePages, "NET_SHARED_HOST_HUGEPAGES", 1);
+static constexpr size_t rcclHugePageSize = 2UL << 20;
+
+static size_t rcclAnonHugeKb(const void* addr) {
+  FILE* f = fopen("/proc/self/smaps", "r");
+  if (f == NULL) return 0;
+  char line[512];
+  bool inVma = false;
+  size_t kb = 0;
+  while (fgets(line, sizeof(line), f)) {
+    unsigned long start, end;
+    if (sscanf(line, "%lx-%lx ", &start, &end) == 2) {
+      if (inVma) break;
+      inVma = (uintptr_t)addr >= start && (uintptr_t)addr < end;
+    } else if (inVma && sscanf(line, "AnonHugePages: %zu kB", &kb) == 1) {
+      break;
+    }
+  }
+  fclose(f);
+  return kb;
+}
+
+// Backs state->hostBuff with a 2 MiB-aligned anonymous mapping that is registered with the GPU runtime. Returns false
+// when the huge-page path is disabled or fails; in that case nothing stays mapped or registered.
+static bool sharedHostBuffAllocHuge(struct ncclProxySharedP2p* state) {
+  if (!rcclParamNetSharedHostHugePages() || (size_t)state->size < rcclHugePageSize) return false;
+  size_t mapSize = ROUNDUP((size_t)state->size, rcclHugePageSize);
+  char* raw =
+    (char*)mmap(NULL, mapSize + rcclHugePageSize, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  if (raw == MAP_FAILED) {
+    ATTN("NET/Shared host pool: mmap of %zu bytes failed (%s), using hipHostMalloc", mapSize, strerror(errno));
+    return false;
+  }
+  char* ptr = (char*)ROUNDUP((uintptr_t)raw, rcclHugePageSize);
+  size_t head = ptr - raw;
+  if (head) munmap(raw, head);
+  if (rcclHugePageSize - head) munmap(ptr + mapSize, rcclHugePageSize - head);
+  if (madvise(ptr, mapSize, MADV_HUGEPAGE) != 0) {
+    INFO(NCCL_ALLOC | NCCL_NET, "NET/Shared host pool: madvise(MADV_HUGEPAGE) failed: %s", strerror(errno));
+  }
+  memset(ptr, 0, mapSize);
+
+  // No CUDACHECK between the mapping and the end of this block: an early return would leak the mapping.
+  void* devPtr = NULL;
+  bool registered = false;
+  cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
+  cudaError_t err = cudaThreadExchangeStreamCaptureMode(&mode);
+  if (err == cudaSuccess) {
+#if defined(HIP_HOST_UNCACHED_MEMORY)
+    err = cudaHostRegister(ptr, mapSize,
+                           cudaHostRegisterPortable | cudaHostRegisterMapped | hipExtHostRegisterUncached);
+#else
+    err = cudaHostRegister(ptr, mapSize, cudaHostRegisterPortable | cudaHostRegisterMapped);
+#endif
+    registered = (err == cudaSuccess);
+    if (registered) err = cudaHostGetDevicePointer(&devPtr, ptr, 0);
+    // Restore the capture mode whether or not the registration succeeded.
+    cudaError_t restoreErr = cudaThreadExchangeStreamCaptureMode(&mode);
+    if (err == cudaSuccess) err = restoreErr;
+  }
+  if (err != cudaSuccess) {
+    ATTN("NET/Shared host pool: huge-page registration failed (%s), using hipHostMalloc", cudaGetErrorString(err));
+    if (registered) (void)cudaHostUnregister(ptr);
+    (void)cudaGetLastError();
+    munmap(ptr, mapSize);
+    return false;
+  }
+  state->hostBuff = ptr;
+  state->hostBuffDev = (char*)devPtr;
+  state->hostBuffMapSize = mapSize;
+  INFO(NCCL_INIT | NCCL_NET, "NET/Shared host pool %p (dev %p) size %zu: %zu kB on huge pages", ptr, devPtr, mapSize,
+       rcclAnonHugeKb(ptr));
+  return true;
+}
+
+static ncclResult_t sharedHostBuffFreeHuge(struct ncclProxySharedP2p* state) {
+  // Same rule as ncclCudaHostFree: at process exit the GPU runtime may already be destroyed, so leave the mapping to
+  // the OS instead of calling into it.
+  if (rcclShutdownFlag().load(std::memory_order_acquire)) {
+    INFO(NCCL_ALLOC, "NET/Shared host pool: skipping free (process shutdown) pointer %p", state->hostBuff);
+    return ncclSuccess;
+  }
+  CUDACHECK(cudaHostUnregister(state->hostBuff));
+  munmap(state->hostBuff, state->hostBuffMapSize);
+  return ncclSuccess;
+}
+#else
+// Huge-page backing is Linux only; other platforms use hipHostMalloc.
+static bool sharedHostBuffAllocHuge(struct ncclProxySharedP2p*) { return false; }
+static ncclResult_t sharedHostBuffFreeHuge(struct ncclProxySharedP2p*) { return ncclSuccess; }
+#endif  // NCCL_OS_LINUX
+
+static ncclResult_t sharedHostBuffAlloc(struct ncclProxySharedP2p* state) {
+  state->hostBuffMapSize = 0;
+  if (sharedHostBuffAllocHuge(state)) return ncclSuccess;
+  NCCLCHECK(ncclCudaHostCalloc(&state->hostBuff, state->size));
+  state->hostBuffDev = state->hostBuff;
+  return ncclSuccess;
+}
+
+static ncclResult_t sharedHostBuffFree(struct ncclProxySharedP2p* state) {
+  if (state->hostBuffMapSize) {
+    NCCLCHECK(sharedHostBuffFreeHuge(state));
+  } else {
+    NCCLCHECK(ncclCudaHostFree(state->hostBuff));
+  }
+  state->hostBuff = NULL;
+  state->hostBuffDev = NULL;
+  state->hostBuffMapSize = 0;
+  return ncclSuccess;
+}
+
 static ncclResult_t sharedNetBuffersInit(struct ncclProxyState* proxyState, int cuda, int tpLocalRank, int type,
                                          int sameProcess, int nChannels, char** gpuPtr, char** cpuPtr, ssize_t* size,
                                          ncclIpcDesc* ipcDesc) {
@@ -910,10 +1029,10 @@ static ncclResult_t sharedNetBuffersInit(struct ncclProxyState* proxyState, int
     }
   }
   if (!cuda && state->hostBuff == NULL) {
-    NCCLCHECK(ncclCudaHostCalloc(&state->hostBuff, state->size));
+    NCCLCHECK(sharedHostBuffAlloc(state));
   }
   if (cpuPtr) *cpuPtr = cuda ? state->cudaBuff : state->hostBuff;
-  if (gpuPtr) *gpuPtr = (cpuPtr && sameProcess) ? *cpuPtr : NULL;
+  if (gpuPtr) *gpuPtr = (cpuPtr && sameProcess) ? (cuda ? state->cudaBuff : state->hostBuffDev) : NULL;
   if (ipcDesc) memcpy(ipcDesc, &state->ipcDesc, sizeof(state->ipcDesc));
   return ncclSuccess;
 }
@@ -945,7 +1064,7 @@ static ncclResult_t sharedNetBuffersDestroy(struct ncclProxyState* proxyState, i
       }
       NCCLCHECK(ncclCudaFree(state->cudaBuff, proxyState->memManager));
     }
-    if (state->hostBuff) NCCLCHECK(ncclCudaHostFree(state->hostBuff));
+    if (state->hostBuff) NCCLCHECK(sharedHostBuffFree(state));
   }
 
   if (peer->send.refcount || peer->recv.refcount) return ncclSuccess;