refactor(ipc): add GPU-side verification and improve IPC benchmark

Author: crazyguitar

experiments/ipc/main.cu

@@ -21,90 +21,53 @@
  */
 #include <affinity/affinity.h>
 #include <bench/arguments.h>
-#include <io/progress.h>
 
 #include <bench/modules/ipc.cuh>
 #include <bench/mpi/fabric.cuh>
 
 /**
- * @brief IPC benchmark with configurable parallelism
- * @tparam BufType Buffer type (SymmetricDMAMemory)
+ * @brief IPC benchmark test with configurable grid/block dimensions
  * @tparam NumBlocks Number of thread blocks (grid dimension)
  * @tparam NumThreads Number of threads per block
  */
-template <typename BufType, unsigned int NumBlocks, unsigned int NumThreads>
-struct TestConfig {
-  static BenchResult Run(size_t size, const Options& opts, FabricBench& peer, std::string_view name = "IPCWrite") {
+template <unsigned int NumBlocks, unsigned int NumThreads>
+struct Test {
+  static BenchResult Run(size_t size, const Options& opts, FabricBench& peer, std::string_view name) {
     int rank = peer.mpi.GetWorldRank();
-    int world = peer.mpi.GetWorldSize();
     int local_size = peer.mpi.GetLocalSize();
-    int local_rank = peer.mpi.GetLocalRank();
 
-    int next = (rank + 1) % world;
-    auto mem = std::make_unique<BufType>(peer.channels[next], size, world, peer.device);
-    auto local_world_ranks = peer.Handshake(mem);
-
-    auto ctx = mem->GetContext();
-    size_t len = size / sizeof(int);
+    // Allocate IPC buffer and exchange handles with local ranks
+    auto bufs = peer.AllocIPC<SymmetricDMAMemory>(size);
+    auto local_world_ranks = peer.Handshake(bufs, std::true_type{});
 
     auto& affinity = GPUloc::Get().GetGPUAffinity()[peer.device];
-    size_t ipc_bw_bps = affinity.mem_support.nvlink_bw * 8;
-
-    cudaLaunchConfig_t cfg{.gridDim = {NumBlocks, 1, 1}, .blockDim = {NumThreads, 1, 1}, .stream = peer.stream};
+    size_t ipc_bw = affinity.mem_support.nvlink_bw * 8;
 
+    // Benchmark rank 0 writing to each local peer via IPC
     double sum_bw = 0, sum_time = 0;
     for (int t = 1; t < local_size; ++t) {
-      int target_rank = local_world_ranks[t];
-      Progress progress(opts.repeat, ipc_bw_bps, name);
-      MPI_Barrier(MPI_COMM_WORLD);
-      auto start = std::chrono::high_resolution_clock::now();
-
-      if (rank == 0) {
-        for (int i = 0; i < opts.repeat; ++i) {
-          LAUNCH_KERNEL(&cfg, IPCWriteKernel, ctx.ipc_ptrs, target_rank, len, 1);
-          if ((i + 1) % Progress::kPrintFreq == 0) {
-            CUDA_CHECK(cudaStreamSynchronize(peer.stream));
-            progress.Print(std::chrono::high_resolution_clock::now(), size, i + 1);
-          }
-        }
-        CUDA_CHECK(cudaStreamSynchronize(peer.stream));
-      }
-
-      MPI_Barrier(MPI_COMM_WORLD);
-      auto end = std::chrono::high_resolution_clock::now();
-
-      // GPU-side verification: target rank checks received data
-      if (rank == target_rank) {
-        int* result;
-        CUDA_CHECK(cudaMallocManaged(&result, sizeof(int)));
-        *result = 1;
-        LAUNCH_KERNEL(&cfg, IPCVerifyKernel, static_cast<int*>(mem->Data()), target_rank, len, result);
-        CUDA_CHECK(cudaStreamSynchronize(peer.stream));
-        bool ok = (*result == 1);
-        CUDA_CHECK(cudaFree(result));
-        if (!ok) throw std::runtime_error("IPC verification failed");
-      }
-
-      double elapsed_us = std::chrono::duration<double, std::micro>(end - start).count();
-      double avg_us = elapsed_us / opts.repeat;
-      double bw_gbps = (size * 8.0) / (avg_us * 1000.0);
-      sum_bw += bw_gbps;
-      sum_time += avg_us;
+      int target = local_world_ranks[t];
+      using Write = IPCWrite<FabricBench, NumBlocks, NumThreads>;
+      using Verify = IPCVerify<FabricBench, NumBlocks, NumThreads>;
+
+      peer.Warmup(bufs, bufs, Write{target}, Verify{target}, opts.warmup);
+      auto r = peer.Bench(name, bufs, bufs, Write{target}, Verify{target}, opts.repeat, 0, ipc_bw);
+      sum_bw += r.bw_gbps;
+      sum_time += r.time_us;
     }
 
     int npairs = local_size - 1;
     double avg_bw = sum_bw / npairs;
-    double link_bw = ipc_bw_bps / 1e9;
-    double bus_bw = (link_bw > 0) ? (avg_bw / link_bw) * 100.0 : 0;
-    return {size, sum_time / npairs, avg_bw, bus_bw};
+    double link_bw = ipc_bw / 1e9;
+    return {size, sum_time / npairs, avg_bw, (link_bw > 0) ? (avg_bw / link_bw) * 100.0 : 0};
   }
 };
 
 // Test configurations demonstrating parallelism impact on NVLink utilization
-using Test1x256 = TestConfig<SymmetricDMAMemory, 1, 256>;      // Low parallelism baseline
-using Test1x1024 = TestConfig<SymmetricDMAMemory, 1, 1024>;    // Single block, max threads
-using Test16x256 = TestConfig<SymmetricDMAMemory, 16, 256>;    // Multi-block, moderate
-using Test128x256 = TestConfig<SymmetricDMAMemory, 128, 256>;  // High parallelism
+using Test1x256 = Test<1, 256>;    // Low parallelism baseline
+using Test1x1024 = Test<1, 1024>;  // Single block, max threads
+using Test16x256 = Test<16, 256>;  // Multi-block, moderate
+using Test128x256 = Test<128, 256>;
 
 int main(int argc, char* argv[]) {
   auto opts = parse_args(argc, argv);

src/include/bench/modules/ipc.cuh

@@ -45,23 +45,22 @@ __global__ void IPCReadKernel(void* const* __restrict__ ipc_ptrs, int source, si
 }
 
 /**
- * @brief IPC Write benchmark functor
+ * @brief IPC Write benchmark functor with configurable parallelism
  */
-template <typename Peer>
+template <typename Peer, unsigned int NumBlocks = 1, unsigned int NumThreads = 256>
 struct IPCWrite {
-  int iters, target;
+  int target;
 
   template <typename T>
-  void operator()(Peer& peer, typename Peer::template Buffers<T>& a, typename Peer::template Buffers<T>& b) {
+  void operator()(Peer& peer, typename Peer::template Buffers<T>& a, typename Peer::template Buffers<T>&) {
     int rank = peer.mpi.GetWorldRank();
     if (rank != 0) return;
 
-    auto& mem = *a[rank];
-    auto ctx = mem.GetContext();
-    size_t len = mem.Size() / sizeof(int);
+    auto ctx = a[rank]->GetContext();
+    size_t len = a[rank]->Size() / sizeof(int);
 
-    cudaLaunchConfig_t cfg{.gridDim = {1, 1, 1}, .blockDim = {256, 1, 1}, .stream = peer.stream};
-    LAUNCH_KERNEL(&cfg, IPCWriteKernel, ctx.ipc_ptrs, target, len, iters);
+    cudaLaunchConfig_t cfg{.gridDim = {NumBlocks, 1, 1}, .blockDim = {NumThreads, 1, 1}, .stream = peer.stream};
+    LAUNCH_KERNEL(&cfg, IPCWriteKernel, ctx.ipc_ptrs, target, len, 1);
     CUDA_CHECK(cudaStreamSynchronize(peer.stream));
   }
 };
@@ -89,20 +88,28 @@ struct IPCRead {
 };
 
 /**
- * @brief IPC Verify functor - verifies data on receiver side
+ * @brief IPC Verify functor with configurable parallelism (GPU-side verification)
  */
+template <typename Peer, unsigned int NumBlocks = 1, unsigned int NumThreads = 256>
 struct IPCVerify {
   int target;
 
-  template <typename P, typename Buffers>
-  void operator()(P& peer, Buffers& b) const {
+  template <typename T>
+  void operator()(Peer& peer, typename Peer::template Buffers<T>& bufs) {
     int rank = peer.mpi.GetWorldRank();
     if (rank != target) return;
 
-    auto& mem = *b[rank];
-    size_t num_ints = mem.Size() / sizeof(int);
-    std::vector<int> host_buf(num_ints);
-    CUDA_CHECK(cudaMemcpy(host_buf.data(), mem.Data(), mem.Size(), cudaMemcpyDeviceToHost));
-    if (VerifyBufferData(host_buf, rank, rank, 0) > 0) throw std::runtime_error("IPC verification failed");
+    auto* data = static_cast<int*>(bufs[rank]->Data());
+    size_t len = bufs[rank]->Size() / sizeof(int);
+
+    int* result;
+    CUDA_CHECK(cudaMallocManaged(&result, sizeof(int)));
+    *result = 1;
+    cudaLaunchConfig_t cfg{.gridDim = {NumBlocks, 1, 1}, .blockDim = {NumThreads, 1, 1}, .stream = peer.stream};
+    LAUNCH_KERNEL(&cfg, IPCVerifyKernel, data, target, len, result);
+    CUDA_CHECK(cudaStreamSynchronize(peer.stream));
+    bool ok = (*result == 1);
+    CUDA_CHECK(cudaFree(result));
+    if (!ok) throw std::runtime_error("IPC verification failed");
   }
 };

src/include/bench/mpi/fabric.cuh

@@ -221,6 +221,30 @@ class FabricBench : public Peer {
     return buffers;
   }
 
+  /**
+   * @brief Allocate buffer for IPC (intra-node) communication
+   *
+   * Unlike Alloc() which creates buffers for each peer, this allocates a single
+   * buffer at bufs[rank] that will be shared via CUDA IPC with local ranks.
+   * Use with Handshake(bufs, std::true_type{}) to exchange IPC handles.
+   *
+   * @tparam T Buffer type (typically SymmetricDMAMemory)
+   * @param size Buffer size in bytes
+   * @param init_value Initial value for buffer (-1 uses rank as value)
+   * @return Vector with buffer only at bufs[rank], others are nullptr
+   */
+  template <typename T>
+  Buffers<T> AllocIPC(size_t size, int init_value = 0) {
+    const auto world_size = mpi.GetWorldSize();
+    const auto rank = mpi.GetWorldRank();
+    Buffers<T> buffers(world_size);
+    buffers[rank] = MakeBuffer<T>(channels[rank], device, size, world_size);
+    const size_t num_ints = size / sizeof(int);
+    int value = (init_value == -1) ? rank : init_value;
+    InitBuffer(buffers[rank].get(), num_ints, value, stream);
+    return buffers;
+  }
+
   /**
    * @brief Allocate send/recv buffer pair
    * @param size Buffer size in bytes
@@ -262,7 +286,11 @@ class FabricBench : public Peer {
   BenchResult
   Bench(std::string_view name, Buffers<T>& a, Buffers<T>& b, F&& func, V&& verify, int iters, size_t progress_bytes = 0, size_t progress_bw = 0) {
     const auto rank = mpi.GetWorldRank();
-    const size_t buf_size = a[rank == 0 ? 1 : 0]->Size();
+    const auto world_size = mpi.GetWorldSize();
+    size_t buf_size = 0;
+    for (int i = 0; i < world_size && buf_size == 0; ++i) {
+      if (a[i]) buf_size = a[i]->Size();
+    }
     const size_t bytes = progress_bytes > 0 ? progress_bytes : buf_size;
     const size_t bw = progress_bw > 0 ? progress_bw : GetBandwidth(0);
     Progress progress(iters, bw, name);

src/include/bootstrap/mpi/fabric.h

@@ -98,17 +98,21 @@ class Peer : private NoCopy {
    *
    * @tparam T SymmetricMemory type (must be DeviceDMABuffer-based)
    * @param bufs Buffer set to exchange IPC handles for
+   * @return Vector mapping local rank index to world rank
    */
   template <typename T>
-  void Handshake(Buffers<T>& bufs, std::true_type /* ipc */) {
+  std::vector<int> Handshake(Buffers<T>& bufs, std::true_type /* ipc */) {
     const int local_size = mpi.GetLocalSize();
     const int local_rank = mpi.GetLocalRank();
     const int world_rank = mpi.GetWorldRank();
 
-    if (local_size <= 1) return;
+    std::vector<int> local_world_ranks(local_size);
+    if (local_size <= 1) {
+      local_world_ranks[0] = world_rank;
+      return local_world_ranks;
+    }
 
     std::vector<cudaIpcMemHandle_t> all_handles(local_size);
-    std::vector<int> local_world_ranks(local_size);
 
     cudaIpcMemHandle_t local_handle;
     CUDA_CHECK(cudaIpcGetMemHandle(&local_handle, bufs[world_rank]->Data()));
@@ -119,6 +123,7 @@ class Peer : private NoCopy {
     MPI_Allgather(&local_handle, hsz, MPI_BYTE, all_handles.data(), hsz, MPI_BYTE, local);
 
     bufs[world_rank]->OpenIPCHandles(all_handles, local_world_ranks, local_rank);
+    return local_world_ranks;
   }
 
   /**