add pipelining send/recv

Summary:
Instead of having 1 databuffer (or staging buffer), make it `pipelineDepth` buffers.
```
   * Pipeline Flow:
   * Time │ Sender              │ Receiver
   * ─────┼─────────────────────┼──────────────────
   *   0  │ Step 0 → Slot[0]    │ -
   *   1  │ Step 1 → Slot[1]    │ Step 0 ← Slot[0]
   *   2  │ Step 2 → Slot[2]    │ Step 1 ← Slot[1]
   *   3  │ Step 3 → Slot[3]    │ Step 2 ← Slot[2]
   *   4  │ Step 4 → Slot[0]    │ Step 3 ← Slot[3]
   *   5  │ Step 5 → Slot[1]    │ Step 4 ← Slot[0]
   *
   * The sender can be ahead by up to 3 steps, hiding NVLink latency and
   * improving throughput

```

Reviewed By: cenzhaometa

Differential Revision: D88463702

fbshipit-source-id: d1414c50e605be6c8a6585b3b0019150d3ce266b

Author: siyengar

comms/pipes/P2pNvlTransport.cc

@@ -15,26 +15,30 @@ P2pNvlTransport::P2pNvlTransport(
       nRanks_(nRanks),
       mpiBootstrap_(mpiBootstrap),
       config_(p2pNvlTransportConfig) {
+  // Calculate total buffer sizes with pipelining
+  const std::size_t totalDataBufferSize =
+      config_.pipelineDepth * config_.dataBufferSize;
+
+  // Calculate state buffer size based on chunk size and pipeline depth
+  const std::size_t numChunksPerStep =
+      (config_.dataBufferSize + config_.chunkSize - 1) / config_.chunkSize;
+  const std::size_t totalNumChunks = config_.pipelineDepth * numChunksPerStep;
+  const std::size_t stateBufferSize = totalNumChunks * sizeof(ChunkState<int>);
+
   dataBuffer_d_ =
-      std::make_unique<meta::comms::DeviceBuffer>(config_.dataBufferSize);
+      std::make_unique<meta::comms::DeviceBuffer>(totalDataBufferSize);
   dataBufferHandler_ = std::make_unique<meta::comms::IpcMemHandler>(
       mpiBootstrap_, myRank, nRanks_);
   dataBufferHandler_->addSelfDeviceMemPtr(dataBuffer_d_->get());
 
-  // Calculate state buffer size based on chunk size
-  // stateBufferSize = dataBufferSize / chunkSize * sizeof(ChunkState<int>)
-  const std::size_t numChunks =
-      (config_.dataBufferSize + config_.chunkSize - 1) / config_.chunkSize;
-  const std::size_t stateBufferSize = numChunks * sizeof(ChunkState<int>);
-
   stateBuffer_d_ = std::make_unique<meta::comms::DeviceBuffer>(stateBufferSize);
   stateBufferHandler_ = std::make_unique<meta::comms::IpcMemHandler>(
       mpiBootstrap_, myRank, nRanks_);
   stateBufferHandler_->addSelfDeviceMemPtr(stateBuffer_d_->get());
 
-  // Initialize state buffer to -1
+  // Initialize state buffer to -1 for all pipeline slots
   auto statePtr = static_cast<ChunkState<int>*>(stateBuffer_d_->get());
-  std::vector<ChunkState<int>> initStates(numChunks, ChunkState<int>(-1));
+  std::vector<ChunkState<int>> initStates(totalNumChunks, ChunkState<int>(-1));
   auto cudaErr = cudaMemcpy(
       statePtr, initStates.data(), stateBufferSize, cudaMemcpyDefault);
   if (cudaErr != cudaSuccess) {
@@ -50,7 +54,9 @@ void P2pNvlTransport::exchange() {
 
 P2pNvlTransportDevice P2pNvlTransport::getTransportDevice(int peerRank) {
   P2pNvlTransportOptions options{
-      .dataBufferSize = config_.dataBufferSize, .chunkSize = config_.chunkSize};
+      .dataBufferSize = config_.dataBufferSize,
+      .chunkSize = config_.chunkSize,
+      .pipelineDepth = config_.pipelineDepth};
 
   LocalState localState{
       .dataBuffer = static_cast<char*>(dataBuffer_d_->get()),

comms/pipes/P2pNvlTransport.h

@@ -13,6 +13,7 @@ namespace comms::pipes {
 struct P2pNvlTransportConfig {
   std::size_t dataBufferSize{0};
   std::size_t chunkSize{0};
+  std::size_t pipelineDepth{0};
 };
 
 // Host-side P2P NVL transport that exchanges IPC buffer handles between ranks