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.

Reviewed By: rmahidhar

Differential Revision: D114552007

Author: Santosh Narayankhedkar

comms/prims/P2pIbTransportBuildContractTest.py

@@ -7,10 +7,14 @@
 
 _PROGRESS_ONLY_FUNCTIONS = (
     "init_send_progress",
+    "init_registered_send_progress",
     "init_recv_progress",
     "progress_send_once",
+    "progress_registered_send_once",
+    "progress_registered_send_drain_once",
     "poll_recv_data_ready",
     "progress_recv_once",
+    "send_registered",
     "store_progress_state",
     "make_progress_geometry",
     "active_payload_offset",

comms/prims/tests/MultipeerIbgdaTransportTest.cc

@@ -6,9 +6,9 @@
 #include <folly/logging/xlog.h>
 #include <array>
 
-#include <chrono>
 #include <memory>
 #include <string>
+#include <tuple>
 #include <vector>
 
 #ifdef __HIP_PLATFORM_AMD__
@@ -1791,6 +1791,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 std::size_t simpleProtocolDataBytes = 16;
+  constexpr int numBlocks = 1;
+  constexpr int blockSize = 128;
+  constexpr uint8_t stagingPoison = 0xA5;
+  constexpr std::size_t zeroByteAfterPostedIndex = 2;
+  const int peerRank = globalRank == 0 ? 1 : 0;
+  const std::array<std::size_t, 8> sizes{
+      0,
+      1,
+      7,
+      simpleProtocolDataBytes - 1,
+      simpleProtocolDataBytes,
+      simpleProtocolDataBytes + 1,
+      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));
+
+    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),
+        index == zeroByteAfterPostedIndex);
+    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 (index == 4) {
+        EXPECT_EQ(hostObservation.postedCount, 0);
+      } else if (index == zeroByteAfterPostedIndex) {
+        EXPECT_EQ(hostObservation.postedCount, 2);
+      } 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,
+          simpleProtocolDataBytes - 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