pvallem · Oct 10, 2026 · Oct 9, 2026 · Oct 8, 2026 · Oct 10, 2026 · Oct 12, 2026
@@ -115,7 +115,8 @@ bool ncclAllGatherDdaFabricEligible(ncclComm* comm, const void* sendbuff, void*
     return false;
   }
 
-  size_t need = sendcount * ncclTypeSize(datatype);
+  // Gate on the gathered total: the launcher rejects totals above scratch.
+  size_t need = sendcount * comm->nRanks * ncclTypeSize(datatype);
   if (need > comm->ddaScratchBytes) {
     return false;
   }

@@ -94,7 +94,8 @@ bool ncclAllGatherDdaIpcEligible(ncclComm* comm, const void* sendbuff, void* rec
     return false;
   }
 
-  size_t need = sendcount * ncclTypeSize(datatype);
+  // Gate on the gathered total: the launcher rejects totals above scratch.
+  size_t need = sendcount * comm->nRanks * ncclTypeSize(datatype);
   if (need > comm->ddaScratchBytes) {
     return false;
   }

@@ -162,8 +162,9 @@ static inline __device__ void allGather(T* const* __restrict__ ipcbuffs, T* __re
   }
 }
 
-inline uint32_t divRoundUp(size_t a, size_t b) {
-  uint32_t y = static_cast<uint32_t>((a + b - 1) / b);
+// Returns size_t: at >= 64 GiB the thread and warp counts exceed uint32, so narrow only after clamping.
+inline size_t divRoundUp(size_t a, size_t b) {
+  size_t y = (a + b - 1) / b;
   if (y == 0) {
     y = 1;
   }
@@ -202,14 +203,13 @@ inline std::pair<dim3, dim3> getGridAndBlockDims(size_t count, int typeSize, siz
   dim3 threads(0, 1, 1);
   dim3 blocks(0, 1, 1);
   if (ddaAlltoAllSingleBlockGrid(count, typeSize)) {
-    threads.x = divRoundUp(count, elementsPerWarp) * kThreadsPerWarp;
+    threads.x = static_cast<uint32_t>(divRoundUp(count, elementsPerWarp) * kThreadsPerWarp);
     blocks.x = 1;
   } else {
     auto warpsRequired = divRoundUp(count, elementsPerWarp);
     blocks.x = calcBlockCount(divRoundUp(count, elementsPerThread), kThreadsPerBlock, maxBlocks);
     auto warpsPerBlock = divRoundUp(warpsRequired, blocks.x);
-    auto threadsPerBlock = std::min<uint32_t>(kThreadsPerBlock, warpsPerBlock * kThreadsPerWarp);
-    threads.x = threadsPerBlock;
+    threads.x = static_cast<uint32_t>(std::min<size_t>(kThreadsPerBlock, warpsPerBlock * kThreadsPerWarp));
   }
 
   return std::make_pair(blocks, threads);

@@ -1023,6 +1023,7 @@ if(BUILD_TESTS)
       DeviceApiMPITests.cpp
       CommNetMismatchMPITests.cpp
       test_sparse_send_recv.cpp
+      DdaAllGatherMPITests.cpp
       DdaAllReduceMPITests.cpp
       DdaFabricSimpleMPITests.cpp
       SymmetricAbortCheckModeMPITests.cpp

@@ -0,0 +1,322 @@
+/*************************************************************************
+ * Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+/**
+ * @file DdaAllGatherMPITests.cpp
+ * @brief DDA IPC AllGather (gfx942/gfx950, 8 ranks): slice placement, and
+ *        fallback when RCCL_DDA_THRESHOLD is above the IPC scratch.
+ *
+ * Run (example):
+ *   mpirun -np 8 ./rccl-UnitTestsMPI --gtest_filter=DdaMPI_AllGather.*
+ *   mpirun -np 8 -x RCCL_DDA_THRESHOLD=1073741824 ./rccl-UnitTestsMPI
+ *     --gtest_filter=DdaMPI_AllGatherThreshold.*
+ */
+
+#ifdef MPI_TESTS_ENABLED
+
+#include "DeviceBufferHelpers.hpp"
+#include "MPIHelpers.hpp"
+#include "MPITestBase.hpp"
+#include "ResourceGuards.hpp"
+#include "TestChecks.hpp"
+#include "algorithms/dda/dda_init_detail.h"
+#include "comm.h"
+#include "rccl_common.h"
+
+#include <gtest/gtest.h>
+#include <hip/hip_runtime.h>
+
+#include <cstddef>
+#include <string>
+#include <vector>
+
+using namespace RCCLTestGuards;
+using namespace RCCLTestHelpers;
+
+namespace
+{
+constexpr int kDdaIpcRanks = nccl_dda_detail::kDdaNranks;
+
+// Keeps f(rank, i) unique per slot and exact in fp32 (7 * 2^20 + 2^18 < 2^24).
+constexpr size_t kRankStride = size_t{1} << 20;
+
+// One block per rank, then 8 MiB total: the gfx942 cap and a full grid-stride.
+constexpr size_t kSmallCount = 1024;
+constexpr size_t kGridCount  = 256 * 1024;
+
+// No rank sends a negative value, so an unwritten slot cannot verify.
+constexpr float kUnwritten = -1.0f;
+
+float slicePattern(int rank, size_t i)
+{
+    return static_cast<float>(static_cast<size_t>(rank) * kRankStride + i);
+}
+
+// Unique per rank and exact in fp32 at any count (max 4095 * 8 + 7 < 2^24).
+float rankPattern(int rank, size_t i)
+{
+    return static_cast<float>((i % 4096) * 8 + static_cast<size_t>(rank));
+}
+
+bool isDdaIpcArch()
+{
+    hipDeviceProp_t props{};
+    if(hipGetDeviceProperties(&props, 0) != hipSuccess)
+    {
+        return false;
+    }
+    const std::string arch(props.gcnArchName);
+    return arch.find("gfx942") != std::string::npos
+           || arch.find("gfx950") != std::string::npos;
+}
+
+// ASSERT_MPI_* skips healthy ranks, so rank 0 can hide a peer's failure.
+int countFailingRanks(bool localOk)
+{
+    int failing = localOk ? 0 : 1;
+    MPI_Allreduce(MPI_IN_PLACE, &failing, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);
+    return failing;
+}
+} // namespace
+
+class DdaIpcAllGatherMPITest : public MPITestBase
+{
+protected:
+    // Empty when this job can reach DDA IPC; every rank gets the same answer. Call before comm init.
+    std::string preInitSkipReason()
+    {
+        if(!validateTestPrerequisites(kDdaIpcRanks, kDdaIpcRanks, false, 1, 1))
+        {
+            return "DDA IPC AllGather needs 8 ranks on one node";
+        }
+        const bool optedOut = rcclParamDdaEnable() == 0;
+        return mpiCoordinatedSkipReason(
+            !isDdaIpcArch() || optedOut,
+            optedOut ? "RCCL_DDA_ENABLE=0" : "DDA IPC needs gfx942 or gfx950");
+    }
+
+    // Implicit launch order disables DDA for the process lifetime.
+    std::string launchOrderSkipReason()
+    {
+        return mpiCoordinatedSkipReason(
+            getActiveCommunicator()->config.launchOrderImplicit == 1,
+            "NCCL_LAUNCH_ORDER_IMPLICIT=1");
+    }
+
+    int selectedAlgo(const void* sendBuf, void* recvBuf, size_t count)
+    {
+        int algo        = -1;
+        int protocol    = -1;
+        int maxChannels = -1;
+        EXPECT_EQ(ncclSuccess,
+                  rcclGetCollImplInfo(getActiveCommunicator(),
+                                      ncclFuncAllGather, count, ncclFloat32,
+                                      ncclSum, sendBuf, recvBuf,
+                                      /*graphCapturing=*/0, &algo, &protocol,
+                                      &maxChannels));
+        return algo;
+    }
+};
+
+class DdaMPI_AllGather : public DdaIpcAllGatherMPITest
+{
+protected:
+    int rank_{};
+    int nRanks_{};
+
+    void SetUp() override
+    {
+        MPITestBase::SetUp();
+        std::string skip = preInitSkipReason();
+        if(!skip.empty())
+        {
+            GTEST_SKIP() << skip;
+        }
+
+        const ncclResult_t res = createTestCommunicator();
+        ASSERT_EQ(0, countFailingRanks(res == ncclSuccess))
+            << "Rank " << getTestMpiRank() << ": " << ncclGetErrorString(res);
+        ncclComm_t comm = getActiveCommunicator();
+        ncclCommUserRank(comm, &rank_);
+        ncclCommCount(comm, &nRanks_);
+
+        skip = launchOrderSkipReason();
+        if(!skip.empty())
+        {
+            GTEST_SKIP() << skip;
+        }
+        ASSERT_EQ(0, countFailingRanks(comm->ddaIpcBarrierState != nullptr))
+            << "Rank " << rank_ << ": DDA IPC did not initialize";
+    }
+
+    void runPlacement(size_t count, bool inPlace)
+    {
+        const size_t totalCount = count * static_cast<size_t>(nRanks_);
+
+        void* recvBuf = nullptr;
+        ASSERT_EQ(0, countFailingRanks(
+                         hipMalloc(&recvBuf, totalCount * sizeof(float))
+                         == hipSuccess));
+        DeviceBufferAutoGuard recvGuard(recvBuf);
+        ASSERT_EQ(0, countFailingRanks(
+                         initializeBufferWithPattern<float>(
+                             recvBuf, totalCount,
+                             [](size_t) { return kUnwritten; })
+                         == hipSuccess));
+
+        void* sendBuf
+            = static_cast<float*>(recvBuf) + static_cast<size_t>(rank_) * count;
+        void* separateSendBuf = nullptr;
+        if(!inPlace)
+        {
+            ASSERT_EQ(0, countFailingRanks(
+                             hipMalloc(&separateSendBuf, count * sizeof(float))
+                             == hipSuccess));
+            sendBuf = separateSendBuf;
+        }
+        DeviceBufferAutoGuard sendGuard(separateSendBuf);
+        ASSERT_EQ(0, countFailingRanks(
+                         initializeBufferWithPattern<float>(
+                             sendBuf, count,
+                             [this](size_t i)
+                             {
+                                 return slicePattern(rank_, i);
+                             })
+                         == hipSuccess));
+
+        // A kernel-path fallback also gathers correctly, so pin the backend.
+        const int   algo     = selectedAlgo(sendBuf, recvBuf, count);
+        const char* algoName = nullptr;
+        rcclGetAlgoName(algo, &algoName);
+        ASSERT_EQ(0, countFailingRanks(algo == RCCL_DDA_IPC))
+            << "Rank " << rank_ << ": " << count << " floats selected "
+            << (algoName ? algoName : "?") << ", not DDA-IPC";
+
+        const ncclResult_t agRes = ncclAllGather(
+            sendBuf, recvBuf, count, ncclFloat32, getActiveCommunicator(),
+            getActiveStream());
+        const hipError_t syncErr = hipStreamSynchronize(getActiveStream());
+        ASSERT_EQ(0, countFailingRanks(agRes == ncclSuccess
+                                       && syncErr == hipSuccess))
+            << "Rank " << rank_ << ": " << ncclGetErrorString(agRes) << " / "
+            << hipGetErrorString(syncErr);
+
+        auto [copyErr, host] = downloadBuffer<float>(recvBuf, totalCount);
+        size_t mismatches = copyErr == hipSuccess ? 0 : totalCount;
+        size_t firstBad   = totalCount;
+        for(size_t j = 0; copyErr == hipSuccess && j < totalCount; ++j)
+        {
+            if(host[j] != slicePattern(static_cast<int>(j / count), j % count))
+            {
+                firstBad = mismatches == 0 ? j : firstBad;
+                ++mismatches;
+            }
+        }
+        ASSERT_EQ(0, countFailingRanks(mismatches == 0))
+            << "Rank " << rank_ << ": " << mismatches << " of " << totalCount
+            << " wrong, first at " << firstBad << " (count " << count << ")";
+    }
+};
+
+class DdaMPI_AllGatherThreshold : public DdaIpcAllGatherMPITest
+{
+};
+
+// recvbuff[r * count + i] must hold f(r, i) on every rank.
+TEST_F(DdaMPI_AllGather, SlicePlacementOutOfPlace)
+{
+    for(size_t count : {kSmallCount, kGridCount})
+    {
+        runPlacement(count, /*inPlace=*/false);
+        if(HasFailure())
+        {
+            return;
+        }
+    }
+}
+
+TEST_F(DdaMPI_AllGather, SlicePlacementInPlace)
+{
+    for(size_t count : {kSmallCount, kGridCount})
+    {
+        runPlacement(count, /*inPlace=*/true);
+        if(HasFailure())
+        {
+            return;
+        }
+    }
+}
+
+// A total above scratch but under the threshold must fall back, not fail.
+TEST_F(DdaMPI_AllGatherThreshold, AboveScratchFallsBack)
+{
+    const size_t threshold = 4 * static_cast<size_t>(DDA_IPC_BUFFER_SIZE);
+    // RCCL caches the param on first read, so set it before comm init.
+    MPIHelpers::MpiEnvGuard thresholdGuard(
+        "RCCL_DDA_THRESHOLD", std::to_string(threshold).c_str());
+    std::string skip = preInitSkipReason();
+    if(!skip.empty())
+    {
+        GTEST_SKIP() << skip;
+    }
+    skip = mpiCoordinatedSkipReason(
+        static_cast<size_t>(rcclParamDdaThreshold()) != threshold,
+        "RCCL_DDA_THRESHOLD was read earlier; run alone or set it at launch");
+    if(!skip.empty())
+    {
+        GTEST_SKIP() << skip;
+    }
+
+    ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
+    ncclComm* comm = getActiveCommunicator();
+    int rank = -1;
+    int nRanks = 0;
+    ASSERT_MPI_EQ(ncclSuccess, ncclCommUserRank(comm, &rank));
+    ASSERT_MPI_EQ(ncclSuccess, ncclCommCount(comm, &nRanks));
+    skip = launchOrderSkipReason();
+    if(!skip.empty())
+    {
+        GTEST_SKIP() << skip;
+    }
+
+    // 2x the per-rank share that fills scratch: 64 MiB each, 512 MiB total.
+    const size_t ranks      = static_cast<size_t>(nRanks);
+    const size_t fullCount  = comm->ddaScratchBytes / (ranks * sizeof(float));
+    const size_t count      = 2 * fullCount;
+    const size_t totalCount = count * ranks;
+    ASSERT_MPI_TRUE(count * sizeof(float) <= comm->ddaScratchBytes);
+    ASSERT_MPI_TRUE(totalCount * sizeof(float) > comm->ddaScratchBytes);
+    ASSERT_MPI_TRUE(totalCount * sizeof(float) <= threshold);
+
+    void* sendBuf = nullptr;
+    void* recvBuf = nullptr;
+    ASSERT_MPI_EQ(hipSuccess, hipMalloc(&sendBuf, count * sizeof(float)));
+    DeviceBufferAutoGuard sendGuard(sendBuf);
+    ASSERT_MPI_EQ(hipSuccess, hipMalloc(&recvBuf, totalCount * sizeof(float)));
+    DeviceBufferAutoGuard recvGuard(recvBuf);
+
+    // Control: a total that exactly fills scratch still takes DDA IPC.
+    ASSERT_MPI_EQ(RCCL_DDA_IPC, selectedAlgo(sendBuf, recvBuf, fullCount));
+    EXPECT_NE(RCCL_DDA_IPC, selectedAlgo(sendBuf, recvBuf, count))
+        << "Rank " << rank << ": DDA IPC picked for a total above scratch";
+
+    ASSERT_MPI_EQ(hipSuccess,
+                  initializeBufferWithPattern<float>(
+                      sendBuf, count,
+                      [rank](size_t i) { return rankPattern(rank, i); }));
+    ASSERT_MPI_EQ(hipSuccess,
+                  hipMemset(recvBuf, 0, totalCount * sizeof(float)));
+    ASSERT_MPI_EQ(ncclSuccess,
+                  ncclAllGather(sendBuf, recvBuf, count, ncclFloat32, comm,
+                                getActiveStream()));
+    ASSERT_MPI_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
+    ASSERT_MPI_TRUE(verifyBufferData<float>(
+        recvBuf, totalCount, [count](size_t j) {
+            return rankPattern(static_cast<int>(j / count), j % count);
+        }));
+}
+
+#endif // MPI_TESTS_ENABLED