Add bulk-buffer registration leases for IBGDA (meta-pytorch#3414)

Summary:

Add move-only bulk-buffer registration leases with stable generations and non-owning contained-range views. Bulk registrations request PCIe Relaxed Ordering, share the existing per-allocation MR cache, and reject invalid, overflowing, stale, or released ranges.

Expose the lease operations through `MultiPeerTransport` for the IBGDA zero-copy AllReduce path. This change does not alter collective behavior or enable zero-copy Ring or Tree.

Add focused lifetime, containment, overlap, ordering, and re-registration coverage. IBRC exercises the shared transport-base implementation as regression coverage; IBGDA remains the product scope.

Reviewed By: rmahidhar

Differential Revision: D114554503

Author: Santosh Narayankhedkar

comms/prims/tests/MultiSegmentRegistrationTest.cc

@@ -5,6 +5,7 @@
 
 #include <folly/init/Init.h>
 #include <cstddef>
+#include <limits>
 #include <memory>
 #include <vector>
 
@@ -151,12 +152,39 @@ struct TransportHandle {
       ibrc->deregisterBuffer(ptr);
     }
   }
+
+  IbBufferRegistrationLease registerIbBulkBuffer(void* ptr, std::size_t size) {
+    return ibgda ? ibgda->registerIbBulkBuffer(ptr, size)
+                 : ibrc->registerIbBulkBuffer(ptr, size);
+  }
+
+  std::optional<IbBufferRegistrationView> lookupIbBulkBuffer(
+      const IbBufferRegistrationLease& lease,
+      void* ptr,
+      std::size_t size) const {
+    return ibgda ? ibgda->lookupIbBulkBuffer(lease, ptr, size)
+                 : ibrc->lookupIbBulkBuffer(lease, ptr, size);
+  }
+
+  void deregisterIbBulkBuffer(IbBufferRegistrationLease& lease) {
+    if (ibgda) {
+      ibgda->deregisterIbBulkBuffer(lease);
+    } else {
+      ibrc->deregisterIbBulkBuffer(lease);
+    }
+  }
+
+  bool isIbBulkBufferViewActive(const IbBufferRegistrationView& view) const {
+    return ibgda ? ibgda->isIbBulkBufferViewActive(view)
+                 : ibrc->isIbBulkBufferViewActive(view);
+  }
 };
 
-TransportHandle createTransport(IbTestBackend backend) {
+TransportHandle createTransport(
+    IbTestBackend backend,
+    const MultipeerIbTransportConfig& config) {
   auto bootstrap = std::make_shared<
       testing::NiceMock<meta::comms::testing::MockBootstrap>>();
-  auto config = makeConfig();
   TransportHandle handle;
   if (backend == IbTestBackend::Ibgda) {
     handle.ibgda = std::make_unique<MultipeerIbgdaTransport>(
@@ -168,6 +196,10 @@ TransportHandle createTransport(IbTestBackend backend) {
   return handle;
 }
 
+TransportHandle createTransport(IbTestBackend backend) {
+  return createTransport(backend, makeConfig());
+}
+
 class MultiSegmentRegistrationTest
     : public ::testing::TestWithParam<IbTestBackend> {};
 
@@ -251,6 +283,161 @@ TEST_P(MultiSegmentRegistrationTest, ContiguousBufferRegistration) {
   CUDACHECK_TEST(cudaFree(devPtr));
 }
 
+TEST_P(MultiSegmentRegistrationTest, BulkLeaseBoundsContainedViews) {
+  CUDACHECK_TEST(cudaSetDevice(0));
+
+  TransportHandle transport;
+  try {
+    transport = createTransport(GetParam());
+  } catch (const std::exception& e) {
+    GTEST_SKIP() << backendName(GetParam())
+                 << " transport not available: " << e.what();
+  }
+
+  constexpr std::size_t kAllocationSize = 4 * 1024 * 1024;
+  constexpr std::size_t kLeaseOffset = 512 * 1024;
+  constexpr std::size_t kLeaseSize = 2 * 1024 * 1024;
+  constexpr std::size_t kViewOffset = 128 * 1024;
+  constexpr std::size_t kViewSize = 256 * 1024;
+  void* allocation = nullptr;
+  CUDACHECK_TEST(cudaMalloc(&allocation, kAllocationSize));
+  auto* const leasePtr = static_cast<char*>(allocation) + kLeaseOffset;
+
+  EXPECT_THROW(
+      transport.registerIbBulkBuffer(leasePtr, 0), std::invalid_argument);
+  auto lease = transport.registerIbBulkBuffer(leasePtr, kLeaseSize);
+  EXPECT_THROW(
+      transport.lookupIbBulkBuffer(lease, leasePtr, 0), std::invalid_argument);
+  EXPECT_THROW(
+      transport.lookupIbBulkBuffer(
+          lease, leasePtr, std::numeric_limits<std::size_t>::max()),
+      std::invalid_argument);
+  auto exact = transport.lookupIbBulkBuffer(lease, leasePtr, kLeaseSize);
+  auto contained =
+      transport.lookupIbBulkBuffer(lease, leasePtr + kViewOffset, kViewSize);
+  auto tail = transport.lookupIbBulkBuffer(lease, leasePtr + kLeaseSize - 1, 1);
+  auto before = transport.lookupIbBulkBuffer(lease, leasePtr - 1, kViewSize);
+  auto after = transport.lookupIbBulkBuffer(
+      lease, leasePtr + kLeaseSize - kViewSize + 1, kViewSize);
+
+  ASSERT_TRUE(exact.has_value());
+  ASSERT_TRUE(contained.has_value());
+  ASSERT_TRUE(tail.has_value());
+  EXPECT_EQ(exact->localBuffer.ptr, leasePtr);
+  EXPECT_EQ(contained->localBuffer.ptr, leasePtr + kViewOffset);
+  EXPECT_EQ(contained->size, kViewSize);
+  EXPECT_EQ(contained->leaseGeneration, lease.generation());
+  EXPECT_FALSE(before.has_value());
+  EXPECT_FALSE(after.has_value());
+  EXPECT_TRUE(transport.isIbBulkBufferViewActive(*contained));
+
+  transport.deregisterIbBulkBuffer(lease);
+  EXPECT_FALSE(lease.valid());
+  EXPECT_FALSE(
+      transport.lookupIbBulkBuffer(lease, leasePtr, kViewSize).has_value());
+  EXPECT_FALSE(transport.isIbBulkBufferViewActive(*contained));
+  EXPECT_THROW(transport.deregisterIbBulkBuffer(lease), std::invalid_argument);
+  CUDACHECK_TEST(cudaFree(allocation));
+}
+
+TEST_P(MultiSegmentRegistrationTest, BulkLeaseReportsEffectiveStrictOrdering) {
+  CUDACHECK_TEST(cudaSetDevice(0));
+
+  auto config = makeConfig();
+  config.enablePciRelaxedOrdering =
+      MultipeerIbTransportConfig::PciRelaxedOrderingMode::Disabled;
+  TransportHandle transport;
+  try {
+    transport = createTransport(GetParam(), config);
+  } catch (const std::exception& e) {
+    GTEST_SKIP() << backendName(GetParam())
+                 << " transport not available: " << e.what();
+  }
+
+  constexpr std::size_t kSize = 2 * 1024 * 1024;
+  void* allocation = nullptr;
+  CUDACHECK_TEST(cudaMalloc(&allocation, kSize));
+
+  auto lease = transport.registerIbBulkBuffer(allocation, kSize);
+  auto view = transport.lookupIbBulkBuffer(lease, allocation, kSize);
+  ASSERT_TRUE(view.has_value());
+  EXPECT_FALSE(view->relaxedOrdering);
+
+  transport.deregisterIbBulkBuffer(lease);
+  CUDACHECK_TEST(cudaFree(allocation));
+}
+
+TEST_P(MultiSegmentRegistrationTest, OverlappingBulkLeasesRemainDistinct) {
+  CUDACHECK_TEST(cudaSetDevice(0));
+
+  TransportHandle transport;
+  try {
+    transport = createTransport(GetParam());
+  } catch (const std::exception& e) {
+    GTEST_SKIP() << backendName(GetParam())
+                 << " transport not available: " << e.what();
+  }
+
+  constexpr std::size_t kAllocationSize = 4 * 1024 * 1024;
+  constexpr std::size_t kOuterSize = 3 * 1024 * 1024;
+  constexpr std::size_t kInnerOffset = 1024 * 1024;
+  constexpr std::size_t kInnerSize = 1024 * 1024;
+  void* allocation = nullptr;
+  CUDACHECK_TEST(cudaMalloc(&allocation, kAllocationSize));
+  auto* const base = static_cast<char*>(allocation);
+
+  auto outer = transport.registerIbBulkBuffer(base, kOuterSize);
+  auto inner = transport.registerIbBulkBuffer(base + kInnerOffset, kInnerSize);
+  auto outerView =
+      transport.lookupIbBulkBuffer(outer, base + kInnerOffset, kInnerSize);
+  auto innerView =
+      transport.lookupIbBulkBuffer(inner, base + kInnerOffset, kInnerSize);
+
+  ASSERT_TRUE(outerView.has_value());
+  ASSERT_TRUE(innerView.has_value());
+  EXPECT_NE(outer.generation(), inner.generation());
+  EXPECT_EQ(outerView->leaseGeneration, outer.generation());
+  EXPECT_EQ(innerView->leaseGeneration, inner.generation());
+
+  transport.deregisterIbBulkBuffer(inner);
+  EXPECT_FALSE(transport.isIbBulkBufferViewActive(*innerView));
+  EXPECT_TRUE(transport.isIbBulkBufferViewActive(*outerView));
+  transport.deregisterIbBulkBuffer(outer);
+  CUDACHECK_TEST(cudaFree(allocation));
+}
+
+TEST_P(MultiSegmentRegistrationTest, ReregistrationChangesLeaseGeneration) {
+  CUDACHECK_TEST(cudaSetDevice(0));
+
+  TransportHandle transport;
+  try {
+    transport = createTransport(GetParam());
+  } catch (const std::exception& e) {
+    GTEST_SKIP() << backendName(GetParam())
+                 << " transport not available: " << e.what();
+  }
+
+  constexpr std::size_t kSize = 2 * 1024 * 1024;
+  void* allocation = nullptr;
+  CUDACHECK_TEST(cudaMalloc(&allocation, kSize));
+
+  auto first = transport.registerIbBulkBuffer(allocation, kSize);
+  auto firstView = transport.lookupIbBulkBuffer(first, allocation, kSize);
+  ASSERT_TRUE(firstView.has_value());
+  const uint64_t firstGeneration = first.generation();
+  transport.deregisterIbBulkBuffer(first);
+
+  auto second = transport.registerIbBulkBuffer(allocation, kSize);
+  auto secondView = transport.lookupIbBulkBuffer(second, allocation, kSize);
+  ASSERT_TRUE(secondView.has_value());
+  EXPECT_NE(firstGeneration, second.generation());
+  EXPECT_FALSE(transport.isIbBulkBufferViewActive(*firstView));
+  EXPECT_TRUE(transport.isIbBulkBufferViewActive(*secondView));
+
+  transport.deregisterIbBulkBuffer(second);
+  CUDACHECK_TEST(cudaFree(allocation));
+}
+
 } // namespace comms::prims::tests
 
 int main(int argc, char** argv) {

comms/prims/transport/MultiPeerIbTransport.cc

@@ -248,6 +248,27 @@ void checkSendRecvSignalAlignment(const void* ptr, const char* label) {
         fmt::format("{} must be {}-byte aligned", label, alignof(SignalState)));
   }
 }
+
+bool rangeContains(
+    uintptr_t outerBegin,
+    std::size_t outerSize,
+    uintptr_t innerBegin,
+    std::size_t innerSize) {
+  return innerSize != 0 && innerBegin >= outerBegin && innerSize <= outerSize &&
+      innerBegin - outerBegin <= outerSize - innerSize;
+}
+
+uintptr_t
+checkedRangeEnd(const void* ptr, std::size_t size, const char* operation) {
+  const auto begin = reinterpret_cast<uintptr_t>(ptr);
+  if (ptr == nullptr || size == 0 ||
+      size > std::numeric_limits<uintptr_t>::max() - begin) {
+    throw std::invalid_argument(
+        fmt::format(
+            "{}: invalid buffer range ptr={} size={}", operation, ptr, size));
+  }
+  return begin + size;
+}
 } // namespace
 
 MultiPeerIbTransportBase::MultiPeerIbTransportBase(
@@ -1081,6 +1102,10 @@ IbgdaLocalBuffer MultiPeerIbTransportBase::registerBuffer(
     throw std::invalid_argument("Invalid buffer pointer or size");
   }
 
+  const auto addr = reinterpret_cast<uintptr_t>(ptr);
+  [[maybe_unused]] const auto requestedEnd =
+      checkedRangeEnd(ptr, size, "registerBuffer");
+
   // Resolve the effective Relaxed Ordering once, up front: the caller's request
   // gated by config (NCCL_IB_PCI_RELAXED_ORDERING) AND by NIC capability probed
   // during openNics. Gating on capability means a NIC whose driver rejects
@@ -1094,11 +1119,10 @@ IbgdaLocalBuffer MultiPeerIbTransportBase::registerBuffer(
   // Fast path: containment lookup — if [ptr, ptr+size) falls entirely within an
   // existing registration with the same effective ordering, return the cached
   // per-NIC lkeys with no driver call.
-  const auto addr = reinterpret_cast<uintptr_t>(ptr);
   auto it = registeredBuffers_.upper_bound(addr);
   if (it != registeredBuffers_.begin()) {
     --it;
-    if (addr + size <= it->first + it->second.allocSize) {
+    if (rangeContains(it->first, it->second.allocSize, addr, size)) {
       // The cache holds one MR set per allocation; its access flags (including
       // Relaxed Ordering) are fixed at registration, so the effective ordering
       // is part of the cache key. A containment hit resolving to different
@@ -1155,7 +1179,6 @@ IbgdaLocalBuffer MultiPeerIbTransportBase::registerBuffer(
   // covers the full contiguous VA — ibv_reg_dmabuf_mr handles the underlying
   // physical discontinuity transparently.
   {
-    const auto requestedEnd = reinterpret_cast<uintptr_t>(ptr) + size;
     const auto allocEnd = static_cast<uintptr_t>(allocBase) + allocSize;
     if (requestedEnd > allocEnd) {
       allocBase = reinterpret_cast<CUdeviceptr>(ptr);
@@ -1319,6 +1342,113 @@ IbgdaLocalBuffer MultiPeerIbTransportBase::registerBuffer(
   return IbgdaLocalBuffer(ptr, keys);
 }
 
+IbBufferRegistrationLease MultiPeerIbTransportBase::registerIbBulkBuffer(
+    void* ptr,
+    std::size_t size) {
+  checkedRangeEnd(ptr, size, "registerIbBulkBuffer");
+  const auto localBuffer = registerBuffer(ptr, size, /*relaxedOrdering=*/true);
+
+  try {
+    const auto addr = reinterpret_cast<uintptr_t>(ptr);
+    auto mrIt = registeredBuffers_.upper_bound(addr);
+    if (mrIt == registeredBuffers_.begin()) {
+      throw std::logic_error(
+          "registerIbBulkBuffer: MR missing after registration");
+    }
+    --mrIt;
+    if (!rangeContains(mrIt->first, mrIt->second.allocSize, addr, size)) {
+      throw std::logic_error(
+          "registerIbBulkBuffer: MR does not contain registered range");
+    }
+    if (nextBulkBufferGeneration_ == std::numeric_limits<uint64_t>::max()) {
+      throw std::overflow_error(
+          "registerIbBulkBuffer: lease generation exhausted");
+    }
+
+    const uint64_t generation = nextBulkBufferGeneration_++;
+    bulkBufferRegistrations_.emplace(
+        generation,
+        BulkBufferRegistration{
+            .ptr = ptr,
+            .size = size,
+            .localBuffer = localBuffer,
+            .relaxedOrdering = mrIt->second.relaxedOrdering,
+        });
+    return IbBufferRegistrationLease(generation);
+  } catch (...) {
+    deregisterBuffer(ptr);
+    throw;
+  }
+}
+
+std::optional<IbBufferRegistrationView>
+MultiPeerIbTransportBase::lookupIbBulkBuffer(
+    const IbBufferRegistrationLease& lease,
+    void* ptr,
+    std::size_t size) const {
+  checkedRangeEnd(ptr, size, "lookupIbBulkBuffer");
+  if (!lease.valid()) {
+    return std::nullopt;
+  }
+
+  const auto registration = bulkBufferRegistrations_.find(lease.generation());
+  if (registration == bulkBufferRegistrations_.end()) {
+    return std::nullopt;
+  }
+
+  const auto& active = registration->second;
+  const auto requestedBegin = reinterpret_cast<uintptr_t>(ptr);
+  const auto registeredBegin = reinterpret_cast<uintptr_t>(active.ptr);
+  if (!rangeContains(registeredBegin, active.size, requestedBegin, size)) {
+    return std::nullopt;
+  }
+
+  return IbBufferRegistrationView{
+      .leaseGeneration = lease.generation(),
+      .localBuffer =
+          active.localBuffer.subBuffer(requestedBegin - registeredBegin),
+      .size = size,
+      .relaxedOrdering = active.relaxedOrdering,
+  };
+}
+
+void MultiPeerIbTransportBase::deregisterIbBulkBuffer(
+    IbBufferRegistrationLease& lease) {
+  if (!lease.valid()) {
+    throw std::invalid_argument(
+        "deregisterIbBulkBuffer: invalid registration lease");
+  }
+
+  const auto registration = bulkBufferRegistrations_.find(lease.generation());
+  if (registration == bulkBufferRegistrations_.end()) {
+    throw std::runtime_error(
+        "deregisterIbBulkBuffer: stale registration lease");
+  }
+
+  void* const ptr = registration->second.ptr;
+  bulkBufferRegistrations_.erase(registration);
+  deregisterBuffer(ptr);
+  lease.reset();
+}
+
+bool MultiPeerIbTransportBase::isIbBulkBufferViewActive(
+    const IbBufferRegistrationView& view) const {
+  if (!view.valid()) {
+    return false;
+  }
+  const auto registration = bulkBufferRegistrations_.find(view.leaseGeneration);
+  if (registration == bulkBufferRegistrations_.end()) {
+    return false;
+  }
+
+  const auto& active = registration->second;
+  return rangeContains(
+      reinterpret_cast<uintptr_t>(active.ptr),
+      active.size,
+      reinterpret_cast<uintptr_t>(view.localBuffer.ptr),
+      view.size);
+}
+
 void MultiPeerIbTransportBase::deregisterBuffer(void* ptr) {
   // Containment lookup on the ordered map avoids resolving the allocation range
   // again (which fails once the underlying memory is freed).