Add registered-source send to IBGDA transport (#3411)

Summary:

Add resumable and blocking registered-source send APIs that let IBGDA read directly from a registered caller buffer. Track posting separately from local NIC completion so callers can safely reuse the source only after an explicit drain.

Preserve the existing staged protocol cursor, receiver staging, signaling, credits, and default behavior.

Differential Revision: D114552007

Author: Santosh Narayankhedkar

comms/prims/tests/MultipeerIbgdaTransportTest.cc

@@ -9,6 +9,8 @@
 #include <chrono>
 #include <memory>
 #include <string>
+#include <thread>
+#include <tuple>
 #include <vector>
 
 #ifdef __HIP_PLATFORM_AMD__
@@ -1666,6 +1668,179 @@ TEST_P(
   }
 }
 
+TEST_F(
+    MultipeerIbgdaTransportTestFixture,
+    RegisteredSendTailsDrainAndStagingIsolation) {
+  if (numRanks != 2) {
+    GTEST_SKIP() << "Skipping test: requires exactly 2 ranks, got " << numRanks;
+  }
+  if (!test::supportsProgressSendRecv()) {
+    GTEST_SKIP() << "registered-source send is not supported for this build";
+  }
+
+  constexpr std::size_t perChannelSize = 64 * 1024;
+  constexpr int pipelineDepth = 2;
+  constexpr std::size_t pipelineChunk = perChannelSize / pipelineDepth;
+  constexpr std::size_t maxSignalBytes = 4 * 1024;
+  constexpr int numBlocks = 1;
+  constexpr int blockSize = 128;
+  constexpr uint8_t stagingPoison = 0xA5;
+  const int peerRank = globalRank == 0 ? 1 : 0;
+  const std::array<std::size_t, 8> sizes{
+      0,
+      1,
+      7,
+      15,
+      16,
+      17,
+      pipelineChunk - 1,
+      pipelineChunk + 1,
+  };
+
+  std::unique_ptr<MultipeerIbgdaTransport> transport;
+  try {
+    MultipeerIbgdaTransportConfig config{
+        .cudaDevice = localRank,
+        .perChannelSize = perChannelSize,
+        .max_num_channels = numBlocks,
+        .pipelineDepth = pipelineDepth,
+    };
+    auto bootstrap = std::make_shared<meta::comms::MpiBootstrap>();
+    transport = std::make_unique<MultipeerIbgdaTransport>(
+        globalRank, numRanks, bootstrap, config);
+    transport->exchange();
+  } catch (const std::exception& e) {
+    GTEST_SKIP() << "IBGDA transport not available: " << e.what();
+  }
+
+  auto* peerTransport = transport->getP2pTransportDevice(peerRank);
+  DeviceBuffer errorCountBuffer(sizeof(int));
+  DeviceBuffer observationBuffer(sizeof(test::RegisteredSendObservation));
+  auto* errorCount = static_cast<int*>(errorCountBuffer.get());
+  auto* observation =
+      static_cast<test::RegisteredSendObservation*>(observationBuffer.get());
+
+  test::testFillTransportStaging(
+      peerTransport,
+      globalRank == 0,
+      0,
+      perChannelSize,
+      stagingPoison,
+      numBlocks,
+      blockSize);
+  CUDACHECK_TEST(cudaDeviceSynchronize());
+
+  for (std::size_t index = 0; index < sizes.size(); ++index) {
+    const std::size_t nbytes = sizes[index];
+    const std::size_t allocationBytes = nbytes == 0 ? 1 : nbytes;
+    const uint8_t pattern = static_cast<uint8_t>(0x20 + index * 13);
+    DeviceBuffer sendBuffer(allocationBytes);
+    DeviceBuffer recvBuffer(allocationBytes);
+    IbgdaLocalBuffer registeredSource{};
+    if (globalRank == 0) {
+      if (nbytes > 0) {
+        registeredSource = transport->registerBuffer(sendBuffer.get(), nbytes);
+      }
+      test::fillBufferWithPattern(
+          sendBuffer.get(), nbytes, pattern, numBlocks, blockSize);
+      CUDACHECK_TEST(
+          cudaMemset(observation, 0, sizeof(test::RegisteredSendObservation)));
+    } else {
+      CUDACHECK_TEST(cudaMemset(recvBuffer.get(), 0, allocationBytes));
+    }
+    CUDACHECK_TEST(cudaDeviceSynchronize());
+    MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));
+
+    // Keep the receiver idle for the first non-empty transfer so the sender
+    // mutates its source immediately after Drained and before recv consumes it.
+    if (index == 1 && globalRank == 1) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(50));
+    }
+    test::testRegisteredSendRecv(
+        peerTransport,
+        registeredSource,
+        recvBuffer.get(),
+        nbytes,
+        maxSignalBytes,
+        globalRank == 0,
+        numBlocks,
+        blockSize,
+        globalRank == 0 ? observation : nullptr,
+        index == 4,
+        globalRank == 0 && nbytes > 0,
+        static_cast<uint8_t>(pattern ^ 0xFF));
+    CUDACHECK_TEST(cudaDeviceSynchronize());
+    MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));
+
+    if (globalRank == 0) {
+      test::RegisteredSendObservation hostObservation{};
+      CUDACHECK_TEST(cudaMemcpy(
+          &hostObservation,
+          observation,
+          sizeof(hostObservation),
+          cudaMemcpyDeviceToHost));
+      if (nbytes == 0 || index == 4) {
+        EXPECT_EQ(hostObservation.postedCount, 0);
+      } else {
+        EXPECT_EQ(hostObservation.postedCount, 1);
+      }
+      EXPECT_EQ(hostObservation.drainedCount, 1);
+    } else {
+      CUDACHECK_TEST(cudaMemset(errorCount, 0, sizeof(int)));
+      test::verifyBufferPattern(
+          recvBuffer.get(), nbytes, pattern, errorCount, numBlocks, blockSize);
+      CUDACHECK_TEST(cudaDeviceSynchronize());
+      int hostErrors = 0;
+      CUDACHECK_TEST(cudaMemcpy(
+          &hostErrors, errorCount, sizeof(hostErrors), cudaMemcpyDeviceToHost));
+      EXPECT_EQ(hostErrors, 0) << "registered send corrupted size " << nbytes;
+    }
+
+    if (index == 1 && globalRank == 1) {
+      CUDACHECK_TEST(cudaMemset(errorCount, 0, sizeof(int)));
+      test::testVerifyTransportStaging(
+          peerTransport,
+          false,
+          1,
+          protocol::Simple::kData - 1,
+          stagingPoison,
+          errorCount,
+          numBlocks,
+          blockSize);
+      CUDACHECK_TEST(cudaDeviceSynchronize());
+      int hostErrors = 0;
+      CUDACHECK_TEST(cudaMemcpy(
+          &hostErrors, errorCount, sizeof(hostErrors), cudaMemcpyDeviceToHost));
+      EXPECT_EQ(hostErrors, 0)
+          << "registered send wrote rounded tail bytes into recv staging";
+    }
+    if (globalRank == 0 && nbytes > 0) {
+      transport->deregisterBuffer(sendBuffer.get());
+    }
+    MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));
+  }
+
+  if (globalRank == 0) {
+    CUDACHECK_TEST(cudaMemset(errorCount, 0, sizeof(int)));
+    test::testVerifyTransportStaging(
+        peerTransport,
+        true,
+        0,
+        perChannelSize,
+        stagingPoison,
+        errorCount,
+        numBlocks,
+        blockSize);
+    CUDACHECK_TEST(cudaDeviceSynchronize());
+    int hostErrors = 0;
+    CUDACHECK_TEST(cudaMemcpy(
+        &hostErrors, errorCount, sizeof(hostErrors), cudaMemcpyDeviceToHost));
+    EXPECT_EQ(hostErrors, 0)
+        << "registered send modified transport send staging";
+  }
+  MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));
+}
+
 // =============================================================================
 // Sustained chunked send/recv - repro for the GB200 per-channel DATA_READY
 // deadlock (two NICs atomic-FA the same flag at maxGroups>=8). Streams a large