fix(rccl/gin): resolve VA signals on the rocSHMEM GDA backend

projects/rccl/src/gin/gin_plugin_rocshmem_gda.cc

@@ -215,9 +215,10 @@ static ncclResult_t ginRocshmemGdaRegMrSym(void* collComm, void* data, size_t si
     return ncclSystemError;
   }
 
-  // Register buffer with QP set's PD
+  // Register buffer with QP set's PD. Remote atomics are required because a VA
+  // signal is an RDMA atomic into a cell of this window.
   uint32_t lkey, rkey;
-  if (rocshmem_gin_reg_mr_vmm(cctx->qpSet, data, size, /*atomic=*/0, &mh->mr, &lkey, &rkey) != 0) {
+  if (rocshmem_gin_reg_mr_vmm(cctx->qpSet, data, size, /*atomic=*/1, &mh->mr, &lkey, &rkey) != 0) {
     WARN("GIN rocshmem-gda: MR registration failed for buffer %p size %zu", data, size);
     (void)hipFree(mh->devHandle);
     free(mh);

projects/rccl/src/include/nccl_device/gin/rocshmem_gda/gin_rocshmem_gda.h

@@ -19,6 +19,25 @@ namespace detail {
 NCCL_DEVICE_INLINE ncclGinRocshmemGdaGPUContext* gdaGpuCtx(ncclGinCtx ctx) {
   return &reinterpret_cast<ncclGinRocshmemGdaGPUContext*>(ctx.handle)[ctx.contextId];
 }
+
+struct GdaSignalTarget {
+  uintptr_t raddr;
+  uint32_t rkey;
+};
+
+// Remote cell a signal increments on peer. vaSignal and indexedSignal share a
+// union, so the descriptor type must pick the field that is read.
+NCCL_DEVICE_INLINE GdaSignalTarget gdaSignalTarget(ncclGinRocshmemGdaGPUContext* rsCtx, int peer,
+                                                   const ncclGinSignalDescriptor& signal) {
+  using nccl::utility::loadConst;
+  if (signal.type == NCCL_GIN_SIGNAL_TYPE_VA) {
+    ncclGinRocshmemGdaMemHandle* mh = (ncclGinRocshmemGdaMemHandle*)signal.vaSignal.signalWindow;
+    return {loadConst(loadConst(&mh->remote_vas) + peer) + signal.vaSignal.signalOffset,
+            loadConst(loadConst(&mh->rkeys) + peer)};
+  }
+  return {loadConst(loadConst(&rsCtx->signal_raddrs) + peer) + sizeof(uint64_t) * signal.indexedSignal.signalId,
+          loadConst(loadConst(&rsCtx->signal_rkeys) + peer)};
+}
 }  // namespace detail
 }  // namespace rocshmem_gda
 }  // namespace gin
@@ -69,10 +88,8 @@ struct ncclGinApi_Put<NCCL_NET_DEVICE_GIN_ROCSHMEM_GDA> {
 
       if (hasSignal) {
         if (signalOp == ncclGinSignalInc) signalOpArg = 1;
-        uintptr_t sigAddr =
-          loadConst(loadConst(&rsCtx->signal_raddrs) + peer) + sizeof(uint64_t) * signal.indexedSignal.signalId;
-        uint32_t sigRkey = loadConst(loadConst(&rsCtx->signal_rkeys) + peer);
-        qp->atomic_add(sigAddr, sigRkey, signalOpArg, wf_info, PostOpt{RingDB<true>});
+        auto target = nccl::gin::rocshmem_gda::detail::gdaSignalTarget(rsCtx, peer, signal);
+        qp->atomic_add(target.raddr, target.rkey, signalOpArg, wf_info, PostOpt{RingDB<true>});
       } else if (hasCounter) {
         qp->quiet(wf_info);
       }
@@ -127,10 +144,8 @@ struct ncclGinApi_PutValue<NCCL_NET_DEVICE_GIN_ROCSHMEM_GDA> {
 
       if (hasSignal) {
         if (signalOp == ncclGinSignalInc) signalOpArg = 1;
-        uintptr_t sigAddr =
-          loadConst(loadConst(&rsCtx->signal_raddrs) + peer) + sizeof(uint64_t) * signal.indexedSignal.signalId;
-        uint32_t sigRkey = loadConst(loadConst(&rsCtx->signal_rkeys) + peer);
-        qp->atomic_add(sigAddr, sigRkey, signalOpArg, wf_info, PostOpt{RingDB<true>});
+        auto target = nccl::gin::rocshmem_gda::detail::gdaSignalTarget(rsCtx, peer, signal);
+        qp->atomic_add(target.raddr, target.rkey, signalOpArg, wf_info, PostOpt{RingDB<true>});
       }
     }
     coop.sync();
@@ -164,9 +179,14 @@ struct ncclGinApi_GetSignalPtr<NCCL_NET_DEVICE_GIN_ROCSHMEM_GDA> {
 template <>
 struct ncclGinApi_ResetSignal<NCCL_NET_DEVICE_GIN_ROCSHMEM_GDA> {
   NCCL_DEVICE_INLINE static void call(ncclGinCtx ctx, ncclGinSignalDescriptor signal) {
+    using nccl::utility::loadConst;
     ncclGinRocshmemGdaGPUContext* rsCtx = nccl::gin::rocshmem_gda::detail::gdaGpuCtx(ctx);
-    if (signal.type == NCCL_GIN_SIGNAL_TYPE_INDEXED)
-      nccl::utility::loadConst(&rsCtx->signals)[signal.indexedSignal.signalId] = 0;
+    if (signal.type == NCCL_GIN_SIGNAL_TYPE_INDEXED) {
+      loadConst(&rsCtx->signals)[signal.indexedSignal.signalId] = 0;
+    } else if (signal.type == NCCL_GIN_SIGNAL_TYPE_VA) {
+      ncclGinRocshmemGdaMemHandle* mh = (ncclGinRocshmemGdaMemHandle*)signal.vaSignal.signalWindow;
+      *reinterpret_cast<uint64_t*>(loadConst(&mh->local_va) + signal.vaSignal.signalOffset) = 0;
+    }
   }
 };
 

projects/rccl/test/device/GinRocshmemGdaTemplate_test.cpp

@@ -567,7 +567,7 @@ TEST_F(GinRocshmemGdaTemplateTest, GetReset_SignalAndCounter) {
   EXPECT_EQ(ctr[0], 0ULL);   // set to 77 then reset
 }
 
-// G12: ResetSignal with a non-indexed descriptor is a no-op.
+// G12: ResetSignal with a NONE descriptor is a no-op.
 __global__ void kernelResetSignalNone(GdaHarness* h) {
   ncclGinCtx ginCtx{};
   ginCtx.handle = &h->ctx;
@@ -584,7 +584,7 @@ TEST_F(GinRocshmemGdaTemplateTest, ResetSignal_NoneIsNoOp) {
   kernelResetSignalNone<<<1, 1>>>(env.dHarness.ptr);
   syncAndCheck();
   auto sigs = env.signals.copyTo();
-  EXPECT_EQ(sigs[0], 42ULL);  // untouched by non-indexed reset
+  EXPECT_EQ(sigs[0], 42ULL);  // untouched by NONE reset
 }
 
 // G13: GDA device dispatch selects the GPU context indexed by ncclGinCtx::contextId
@@ -704,6 +704,84 @@ TEST_F(GinRocshmemGdaTemplateTest, PutValue_SelectLogicalContextSignalStripe) {
       << "contextId=1 PutValue must deliver the signal into the context-1 stripe";
 }
 
+// G16: VA signals (AICOMRCCL-3180) target a cell in the dst window, never the indexed signal buffer.
+constexpr size_t kVaSignalOff = 256;
+constexpr size_t kVaWindowBytes = kVaSignalOff + sizeof(uint64_t);
+
+__global__ void kernelPutVaSignal(GdaHarness* h, bool putValue, uint64_t val) {
+  ncclGinCtx ginCtx{};
+  ginCtx.handle = &h->ctx;
+  ginCtx.nRanks = 2;
+  ncclGinSignalDescriptor sig{};
+  sig.type = NCCL_GIN_SIGNAL_TYPE_VA;
+  sig.vaSignal.signalWindow = reinterpret_cast<ncclGinWindow_t>(&h->dstMh);
+  sig.vaSignal.signalOffset = kVaSignalOff;
+  if (putValue) {
+    ncclGinApi_PutValue<NCCL_NET_DEVICE_GIN_ROCSHMEM_GDA>::call(
+        ginCtx, ncclCoopThread{}, 1, reinterpret_cast<ncclGinWindow_t>(&h->dstMh), 0, val, sig, ncclGinSignalAdd, 9,
+        false, nullptr, cuda::thread_scope_system, cuda::thread_scope_system);
+  } else {
+    ncclGinApi_Put<NCCL_NET_DEVICE_GIN_ROCSHMEM_GDA>::call(
+        ginCtx, ncclCoopThread{}, 1, true, reinterpret_cast<ncclGinWindow_t>(&h->dstMh), 0,
+        reinterpret_cast<ncclGinWindow_t>(&h->srcMh), 0, sizeof(uint64_t), sig, ncclGinSignalAdd, 9, false, 0, false,
+        nullptr, cuda::thread_scope_system, cuda::thread_scope_system);
+  }
+}
+
+static uint64_t readU64At(const DeviceBuffer<uint8_t>& buf, size_t off) {
+  auto got = buf.copyTo();
+  uint64_t v = 0;
+  std::memcpy(&v, got.data() + off, sizeof(v));
+  return v;
+}
+
+TEST_F(GinRocshmemGdaTemplateTest, VaSignal_PutSignalsWindowCell) {
+  const uint64_t kVal = 0x1122334455667788ULL;
+  for (const bool putValue : {false, true}) {
+    SCOPED_TRACE(::testing::Message() << "putValue=" << putValue);
+    GdaEnv env(kVaWindowBytes);
+    std::vector<uint8_t> src(kVaWindowBytes, 0);
+    std::memcpy(src.data(), &kVal, sizeof(kVal));
+    env.src.copyFrom(src);
+    env.dst.zero();
+    env.build();
+    resetSignalCount();
+    kernelPutVaSignal<<<1, 1>>>(env.dHarness.ptr, putValue, kVal);
+    syncAndCheck();
+    EXPECT_EQ(readSignalCount(), 1ULL);
+    EXPECT_EQ(readU64At(env.dst, 0), kVal) << "payload must still land at the destination offset";
+    EXPECT_EQ(readU64At(env.dst, kVaSignalOff), 9ULL) << "the VA cell must take the signal";
+    EXPECT_EQ(env.signals.copyTo(), std::vector<uint64_t>(GdaEnv::kNSignals, 0))
+        << "a VA signal must not touch the indexed signal buffer";
+  }
+}
+
+// G17: ResetSignal clears a VA signal's cell through the window's local VA.
+__global__ void kernelResetVaSignal(GdaHarness* h) {
+  ncclGinCtx ginCtx{};
+  ginCtx.handle = &h->ctx;
+  ncclGinSignalDescriptor sig{};
+  sig.type = NCCL_GIN_SIGNAL_TYPE_VA;
+  sig.vaSignal.signalWindow = reinterpret_cast<ncclGinWindow_t>(&h->dstMh);
+  sig.vaSignal.signalOffset = kVaSignalOff;
+  ncclGinApi_ResetSignal<NCCL_NET_DEVICE_GIN_ROCSHMEM_GDA>::call(ginCtx, sig);
+}
+
+TEST_F(GinRocshmemGdaTemplateTest, ResetSignal_VaClearsWindowCell) {
+  GdaEnv env(kVaWindowBytes);
+  std::vector<uint8_t> dst(kVaWindowBytes, 0);
+  const uint64_t armed = 5;
+  std::memcpy(dst.data() + kVaSignalOff, &armed, sizeof(armed));
+  env.dst.copyFrom(dst);
+  env.build();
+  env.signals.copyFrom(std::vector<uint64_t>(GdaEnv::kNSignals, 7));
+  kernelResetVaSignal<<<1, 1>>>(env.dHarness.ptr);
+  syncAndCheck();
+  EXPECT_EQ(readU64At(env.dst, kVaSignalOff), 0ULL);
+  EXPECT_EQ(env.signals.copyTo(), std::vector<uint64_t>(GdaEnv::kNSignals, 7))
+      << "a VA reset must not clear indexed signals";
+}
+
 }  // namespace RcclUnitTesting
 
 #endif  // NCCL_GIN_ROCSHMEM_GDA_ENABLE

projects/rccl/test/gin/GinRocshmemGdaPlugin_test.cpp

@@ -243,4 +243,37 @@ TEST_F(GinRocshmemGdaPluginTest, CreateContext_SignalMrFailureReturnsErrorWithou
   EXPECT_EQ(Stubs::GetLiveHipAllocationCount(), hipAllocsBefore);
 }
 
+// A VA signal is an RDMA atomic into a user window, so the window MR must allow
+// remote atomics (AICOMRCCL-3180).
+TEST_F(GinRocshmemGdaPluginTest, RegMrSym_AllowsRemoteAtomicsForVaSignals) {
+  constexpr int kRanks = 2, kRank = 0;
+  constexpr size_t kBytes = 4096;
+  connect(kRanks, kRank);
+
+  void* buf = nullptr;
+  ASSERT_EQ(hipMalloc(&buf, kBytes), hipSuccess);
+  void* mhandle = nullptr;
+  void* ginHandle = nullptr;
+  ASSERT_EQ(plugin_.regMrSym(coll_, buf, kBytes, NCCL_PTR_CUDA, 0, &mhandle, &ginHandle), ncclSuccess);
+
+  ASSERT_EQ(Stubs::GetRegMrCalls().size(), 1u);
+  const Stubs::RegMrCall mr = Stubs::GetRegMrCalls()[0];
+  EXPECT_EQ(mr.addr, buf);
+  EXPECT_EQ(mr.size, kBytes);
+  EXPECT_EQ(mr.atomic, 1);
+
+  const auto devMh = CopyToHost(static_cast<const ncclGinRocshmemGdaMemHandle*>(ginHandle), 1)[0];
+  EXPECT_EQ(devMh.local_va, reinterpret_cast<uintptr_t>(buf));
+  const auto rkeys = CopyToHost(devMh.rkeys, kRanks);
+  const auto vas = CopyToHost(devMh.remote_vas, kRanks);
+  EXPECT_EQ(rkeys[kRank], mr.rkey);
+  EXPECT_EQ(rkeys[1], Stubs::PeerRkey(1));
+  EXPECT_EQ(vas[kRank], reinterpret_cast<uintptr_t>(buf));
+  EXPECT_EQ(vas[1], Stubs::PeerAddr(1));
+
+  EXPECT_EQ(plugin_.deregMrSym(coll_, mhandle), ncclSuccess);
+  EXPECT_EQ(Stubs::GetLiveMrCount(), 0);
+  EXPECT_EQ(hipFree(buf), hipSuccess);
+}
+
 }  // namespace RcclUnitTesting

projects/rccl/test/transport/GinDeviceMPITests.cpp

@@ -177,15 +177,6 @@ std::string proxyNthreadsSkipReason() {
   return "";
 }
 
-// rocSHMEM-GDA implements only INDEXED signals; VA signals are unsupported
-// (Put/PutValue address the signal via indexedSignal.signalId).
-std::string vaSignalTestSkipReason() {
-  if (requestedGinType() == NCCL_NET_DEVICE_GIN_ROCSHMEM_GDA)
-    return "VA signals not supported by rocSHMEM-GDA (NCCL_GIN_TYPE=" +
-           std::to_string(NCCL_NET_DEVICE_GIN_ROCSHMEM_GDA) + ")";
-  return "";
-}
-
 // BarrierFence_* exercises Anvil SDMA queue ordering and proxy per-context state.
 // rocSHMEM GDA (NCCL_GIN_TYPE=6) runs the rest of this file but skips these cases.
 std::string barrierFenceBackendSkipReason() {
@@ -6097,8 +6088,6 @@ __global__ void putVAMultiCtxConsumerKernel(
 void GinMPIDeviceTests::runVASignalPut(int nContexts) {
   if (auto reason = ginProxyTestSkipReason(); !reason.empty())
     GTEST_SKIP() << reason;
-  if (auto reason = vaSignalTestSkipReason(); !reason.empty())
-    GTEST_SKIP() << reason;
   if (!validateTestPrerequisites(/*min_processes=*/2, /*max_processes=*/2))
     GTEST_SKIP() << "Requires exactly 2 ranks";
 
@@ -6339,8 +6328,6 @@ __global__ void indexedSignalResetOnlyKernel(
 TEST_F(GinMPIDeviceTests, VASignal_NoPayload_IncAndAdd) {
   if (auto reason = ginProxyTestSkipReason(); !reason.empty())
     GTEST_SKIP() << reason;
-  if (auto reason = vaSignalTestSkipReason(); !reason.empty())
-    GTEST_SKIP() << reason;
   if (!validateTestPrerequisites(/*min_processes=*/2, /*max_processes=*/2))
     GTEST_SKIP() << "Requires exactly 2 ranks";
 
@@ -6454,8 +6441,6 @@ TEST_F(GinMPIDeviceTests, VASignal_NoPayload_IncAndAdd) {
 TEST_F(GinMPIDeviceTests, VASignal_ReadAndReset) {
   if (auto reason = ginProxyTestSkipReason(); !reason.empty())
     GTEST_SKIP() << reason;
-  if (auto reason = vaSignalTestSkipReason(); !reason.empty())
-    GTEST_SKIP() << reason;
   if (!validateTestPrerequisites(/*min_processes=*/2, /*max_processes=*/2))
     GTEST_SKIP() << "Requires exactly 2 ranks";
 
@@ -6810,8 +6795,6 @@ TEST_F(GinMPIDeviceTests, Signal_WaitRead_Low32Bits) {
 TEST_F(GinMPIDeviceTests, VASignal_StrictOrderingFlag) {
   if (auto reason = ginProxyTestSkipReason(); !reason.empty())
     GTEST_SKIP() << reason;
-  if (auto reason = vaSignalTestSkipReason(); !reason.empty())
-    GTEST_SKIP() << reason;
   if (!validateTestPrerequisites(/*min_processes=*/2, /*max_processes=*/2))
     GTEST_SKIP() << "Requires exactly 2 ranks";
 

projects/rccl/tools/ci/lib/gin-tests.json

@@ -127,6 +127,24 @@
       "env": ["NCCL_CUMEM_ENABLE=1", "RCCL_ENABLE_INTRANET=1", "NCCL_DMABUF_ENABLE=1", "HSA_NO_SCRATCH_RECLAIM=1", "NCCL_MSCCL_ENABLE=0", "NCCL_GIN_ENABLE=1", "NCCL_GIN_TYPE=7", "NCCL_DEBUG=WARN"],
       "args": "--gtest_filter=GinMPIDeviceTests.AnvilStrongSignal_*:GinMPIDeviceTests.VASignal_*"
     },
+    {
+      "_comment": "np 2: GDA VASignal_* coverage for AICOMRCCL-3180. MultiContext needs NCCL_GIN_NCONTEXTS>1; run that case separately.",
+      "name": "rccl-gin-mpi-gda-va-signal",
+      "kind": "mpi-fixtures",
+      "bin": "rccl-UnitTestsMPI",
+      "np": 2,
+      "env": ["NCCL_CUMEM_ENABLE=1", "RCCL_ENABLE_INTRANET=1", "NCCL_DMABUF_ENABLE=1", "HSA_NO_SCRATCH_RECLAIM=1", "NCCL_MSCCL_ENABLE=0", "NCCL_GIN_ENABLE=1", "NCCL_GIN_TYPE=6", "NCCL_DEBUG=WARN"],
+      "args": "--gtest_filter=GinMPIDeviceTests.VASignal_Put:GinMPIDeviceTests.VASignal_NoPayload_IncAndAdd:GinMPIDeviceTests.VASignal_ReadAndReset:GinMPIDeviceTests.VASignal_StrictOrderingFlag"
+    },
+    {
+      "_comment": "np 2: VASignal_Put_MultiContext self-skips unless NCCL_GIN_NCONTEXTS>1.",
+      "name": "rccl-gin-mpi-gda-va-signal-multictx",
+      "kind": "mpi-fixtures",
+      "bin": "rccl-UnitTestsMPI",
+      "np": 2,
+      "env": ["NCCL_CUMEM_ENABLE=1", "RCCL_ENABLE_INTRANET=1", "NCCL_DMABUF_ENABLE=1", "HSA_NO_SCRATCH_RECLAIM=1", "NCCL_MSCCL_ENABLE=0", "NCCL_GIN_ENABLE=1", "NCCL_GIN_TYPE=6", "NCCL_GIN_NCONTEXTS=2", "NCCL_DEBUG=WARN"],
+      "args": "--gtest_filter=GinMPIDeviceTests.VASignal_Put_MultiContext"
+    },
     {
       "name": "rccl-gin-rocshmem-sdma-internal",
       "kind": "rccl-tests",

projects/rccl/tools/scripts/test_runner/configs/gin_rocshmem_backends_mellanox.json

@@ -180,6 +180,13 @@
         { "name": "Properties_NLsaTeams",     "description": "GIN properties: number of LSA teams",                   "test_filter": "GinMPIDeviceTests.Properties_NLsaTeams" },
         { "name": "MultiContext_Exclusive",   "description": "GIN multi-context exclusive contexts",                  "test_filter": "GinMPIDeviceTests.MultiContext_Exclusive" },
         { "name": "VASignal_Put",             "description": "GIN VA-signal put",                                     "test_filter": "GinMPIDeviceTests.VASignal_Put" },
+        { "name": "VASignal_NoPayload_IncAndAdd", "description": "GIN VA signal-only Inc and Add",                   "test_filter": "GinMPIDeviceTests.VASignal_NoPayload_IncAndAdd" },
+        { "name": "VASignal_ReadAndReset",    "description": "GIN VA signal read and reset",                          "test_filter": "GinMPIDeviceTests.VASignal_ReadAndReset" },
+        { "name": "VASignal_StrictOrderingFlag", "description": "GIN VA signal with STRICT_ORDERING",               "test_filter": "GinMPIDeviceTests.VASignal_StrictOrderingFlag" },
+        {
+          "name": "VASignal_Put_MultiContext", "description": "GIN VA-signal put across multiple contexts", "test_filter": "GinMPIDeviceTests.VASignal_Put_MultiContext",
+          "env_variables": { "NCCL_GIN_NCONTEXTS": "2" }
+        },
         {
           "name": "RailConnection_Create", "description": "GIN rail connection create; rail-only mode (NCCL_CROSS_NIC=0)", "test_filter": "GinMPIDeviceTests.RailConnection_Create",
           "env_variables": { "NCCL_CROSS_NIC": "0" }