Add registered-source comparison to IBGDA benchmark (meta-pytorch#3413)

Summary:

Add an opt-in `--ibgda_sendrecv_enable_registered` mode to the IBGDA send/receive benchmark so registered-source sends can be measured against the existing staged path. Keep staged sends as the default, use the same one-block benchmark geometry for both paths, and explicitly drain registered sends before their source storage can be reused.

Reviewed By: rmahidhar

Differential Revision: D114552005

Author: Santosh Narayankhedkar

comms/prims/benchmarks/IbgdaSendRecv.cu

@@ -409,7 +409,8 @@ __global__ void __launch_bounds__(512, 1) ibgda_progress_send_kernel(
     std::size_t totalBytes,
     int numBlocks,
     std::size_t maxSignalBytes,
-    Timeout timeout) {
+    Timeout timeout,
+    bool waitForSlotFree) {
   auto group = make_block_group();
 
   const std::size_t sectionBytes = section_bytes(transport, totalBytes);
@@ -423,6 +424,50 @@ __global__ void __launch_bounds__(512, 1) ibgda_progress_send_kernel(
            IbgdaSendRecvProgressStatus::Done) {
     }
   }
+
+  if (waitForSlotFree) {
+    auto& channel = transport->local_channel(group.group_id);
+    transport->wait_signal(
+        group,
+        channel.slotFree,
+        static_cast<uint64_t>(channel.sendProgress.nextStep),
+        timeout);
+  }
+}
+
+__global__ void __launch_bounds__(512, 1) ibgda_registered_progress_send_kernel(
+    P2pIbgdaTransportDevice* transport,
+    IbgdaLocalBuffer src,
+    std::size_t totalBytes,
+    int numBlocks,
+    std::size_t maxSignalBytes,
+    Timeout timeout) {
+  auto group = make_block_group();
+
+  auto status = IbgdaRegisteredSendProgressStatus::Waiting;
+  const std::size_t sectionBytes = section_bytes(transport, totalBytes);
+  const std::size_t totalSections = totalBytes / sectionBytes;
+  for (std::size_t s = 0; s < totalSections; ++s) {
+    const auto section = src.subBuffer(s * sectionBytes);
+    transport->init_registered_send_progress(
+        group, sectionBytes, maxSignalBytes);
+    status = IbgdaRegisteredSendProgressStatus::Waiting;
+    while (status != IbgdaRegisteredSendProgressStatus::Posted &&
+           status != IbgdaRegisteredSendProgressStatus::Drained) {
+      status = transport->progress_registered_send_once(
+          group, section, sectionBytes, maxSignalBytes, timeout);
+    }
+  }
+  while (status != IbgdaRegisteredSendProgressStatus::Drained) {
+    status = transport->progress_registered_send_drain_once(group, timeout);
+  }
+
+  auto& channel = transport->local_channel(group.group_id);
+  transport->wait_signal(
+      group,
+      channel.slotFree,
+      static_cast<uint64_t>(channel.sendProgress.nextStep),
+      timeout);
 }
 
 __global__ void __launch_bounds__(512, 1) ibgda_progress_recv_kernel(
@@ -467,7 +512,36 @@ void launch_ibgda_progress_send(
   printf("[PIPES] progress send benchmark is NVIDIA-only\n");
 #else
   ibgda_progress_send_kernel<<<numBlocks, 512, 0, stream>>>(
-      transport, src, nbytes, numBlocks, maxSignalBytes, timeout);
+      transport, src, nbytes, numBlocks, maxSignalBytes, timeout, false);
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess) {
+    printf(
+        "[PIPES] progress send kernel launch failed: %s\n",
+        cudaGetErrorString(err));
+  }
+#endif
+}
+
+void launch_ibgda_progress_send_complete(
+    P2pIbgdaTransportDevice* transport,
+    char* src,
+    std::size_t nbytes,
+    int numBlocks,
+    cudaStream_t stream,
+    std::size_t maxSignalBytes,
+    Timeout timeout) {
+#ifdef __HIP_PLATFORM_AMD__
+  (void)transport;
+  (void)src;
+  (void)nbytes;
+  (void)numBlocks;
+  (void)stream;
+  (void)maxSignalBytes;
+  (void)timeout;
+  printf("[PIPES] progress send benchmark is NVIDIA-only\n");
+#else
+  ibgda_progress_send_kernel<<<numBlocks, 512, 0, stream>>>(
+      transport, src, nbytes, numBlocks, maxSignalBytes, timeout, true);
   cudaError_t err = cudaGetLastError();
   if (err != cudaSuccess) {
     printf(
@@ -477,6 +551,35 @@ void launch_ibgda_progress_send(
 #endif
 }
 
+void launch_ibgda_registered_progress_send(
+    P2pIbgdaTransportDevice* transport,
+    const IbgdaLocalBuffer& src,
+    std::size_t nbytes,
+    int numBlocks,
+    cudaStream_t stream,
+    std::size_t maxSignalBytes,
+    Timeout timeout) {
+#ifdef __HIP_PLATFORM_AMD__
+  (void)transport;
+  (void)src;
+  (void)nbytes;
+  (void)numBlocks;
+  (void)stream;
+  (void)maxSignalBytes;
+  (void)timeout;
+  printf("[PIPES] registered progress send benchmark is NVIDIA-only\n");
+#else
+  ibgda_registered_progress_send_kernel<<<numBlocks, 512, 0, stream>>>(
+      transport, src, nbytes, numBlocks, maxSignalBytes, timeout);
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess) {
+    printf(
+        "[PIPES] registered progress send kernel launch failed: %s\n",
+        cudaGetErrorString(err));
+  }
+#endif
+}
+
 void launch_ibgda_progress_recv(
     P2pIbgdaTransportDevice* transport,
     char* dst,

comms/prims/benchmarks/IbgdaSendRecv.cuh

@@ -86,6 +86,19 @@ __global__ void ibgda_progress_send_kernel(
     std::size_t totalBytes,
     int numBlocks,
     std::size_t maxSignalBytes,
+    Timeout timeout,
+    bool waitForSlotFree);
+
+/**
+ * Unidirectional registered-source progress send kernel.
+ * Grid: numBlocks. Block: 512 threads.
+ */
+__global__ void ibgda_registered_progress_send_kernel(
+    P2pIbgdaTransportDevice* transport,
+    IbgdaLocalBuffer src,
+    std::size_t totalBytes,
+    int numBlocks,
+    std::size_t maxSignalBytes,
     Timeout timeout);
 
 /**

comms/prims/benchmarks/IbgdaSendRecv.h

@@ -7,6 +7,7 @@
 #include <cstdint>
 
 #include "comms/prims/core/Timeout.cuh"
+#include "comms/prims/transport/ibgda/IbgdaBuffer.h"
 
 namespace comms::prims {
 class P2pIbgdaTransportDevice;
@@ -128,6 +129,33 @@ void launch_ibgda_progress_send(
     std::size_t maxSignalBytes = 0,
     Timeout timeout = Timeout());
 
+/**
+ * Launch the staged progress sender and wait for the receiver's final credit.
+ */
+void launch_ibgda_progress_send_complete(
+    P2pIbgdaTransportDevice* transport,
+    char* src,
+    std::size_t nbytes,
+    int numBlocks,
+    cudaStream_t stream,
+    std::size_t maxSignalBytes = 0,
+    Timeout timeout = Timeout());
+
+/**
+ * Launch a unidirectional registered-source progress send kernel.
+ *
+ * The source is read directly by the NIC. The kernel drains local NIC reads
+ * and waits for the receiver's final slot-free credit before returning.
+ */
+void launch_ibgda_registered_progress_send(
+    P2pIbgdaTransportDevice* transport,
+    const IbgdaLocalBuffer& src,
+    std::size_t nbytes,
+    int numBlocks,
+    cudaStream_t stream,
+    std::size_t maxSignalBytes = 0,
+    Timeout timeout = Timeout());
+
 /**
  * Launch unidirectional progress recv kernel. All blocks receive.
  */