c10d/gloo: add ibverbs backend (#437)

Summary:
X-link: pytorch/pytorch#153015


This provides a new "UnboundBuffer" implementation for Gloo ibverbs backend so it can be used with PyTorch.

This currently is passing basic tests such as `reduce_test` and `send_recv_test` but there are a number of failures. Putting this up for review so the follow up fixes are less of a mega PR and also so we can start doing some initial testing with this E2E with PyTorch.

Known issues:

* support recv from any is not supported
* AllreduceBcubeBase2 is failing

Reviewed By: fduwjj

Differential Revision: D73291471

Author: d4l3k

gloo/allreduce_ring_chunked.h

@@ -145,14 +145,19 @@ class AllreduceRingChunked : public Algorithm {
         fn_->call(&ptrs_[0][offset], inbox_[chunkOffset & 1], length);
       }
 
+      GLOO_DEBUG("START sendNotification round=", round);
+
       // Send notification to node on the left that
       // this node is ready for an inbox write.
       sendNotificationBuf_->send();
 
+      GLOO_DEBUG("START recvNotification round=", round);
       // Wait for notification from node on the right
       // to be sure this node can start an inbox write.
       recvNotificationBuf_->waitRecv();
 
+      GLOO_DEBUG("DONE recvNotification round=", round);
+
       // Copy accumulated chunk
       copyChunkAtOffset(chunkOffset);
     }

gloo/common/logging.h

@@ -15,10 +15,20 @@
 #include <limits>
 #include <vector>
 
+#include "gloo/common/error.h"
 #include "gloo/common/string.h"
 
 namespace gloo {
 
+#define GLOO_LOG_MSG(level, ...)   \
+  std::cerr << ::gloo::MakeString( \
+      "[", __FILE__, ":", __LINE__, "] ", level, " ", __VA_ARGS__, "\n")
+
+#define GLOO_INFO(...) GLOO_LOG_MSG("INFO", __VA_ARGS__)
+#define GLOO_ERROR(...) GLOO_LOG_MSG("ERROR", __VA_ARGS__)
+#define GLOO_WARN(...) GLOO_LOG_MSG("WARN", __VA_ARGS__)
+#define GLOO_DEBUG(...) // GLOO_LOG_MSG("DEBUG", __VA_ARGS__)
+
 class EnforceNotMet : public std::exception {
  public:
   EnforceNotMet(
@@ -157,7 +167,4 @@ BINARY_COMP_HELPER(LessEquals, <=)
 #define GLOO_ENFORCE_GT(x, y, ...) \
   GLOO_ENFORCE_THAT_IMPL(Greater((x), (y)), #x " > " #y, __VA_ARGS__)
 
-#define GLOO_ERROR(...) \
-  std::cerr << "Gloo error: " << ::gloo::MakeString(__VA_ARGS__) << std::endl
-
 } // namespace gloo

gloo/test/base_test.cc

@@ -16,23 +16,35 @@ const char* kDefaultDevice = "localhost";
 
 // Transports that instantiated algorithms can be tested against.
 const std::vector<Transport> kTransportsForClassAlgorithms = {
+#if GLOO_HAVE_TRANSPORT_TCP
     Transport::TCP,
     Transport::TCP_LAZY,
+#endif
 #if GLOO_HAVE_TRANSPORT_TCP_TLS
     Transport::TCP_TLS,
 #endif
+#if GLOO_HAVE_TRANSPORT_IBVERBS
+    Transport::IBVERBS,
+#endif
 };
 
 // Transports that function algorithms can be tested against.
 // This is the new style of calling collectives and must be
 // preferred over the instantiated style.
 const std::vector<Transport> kTransportsForFunctionAlgorithms = {
+#if GLOO_HAVE_TRANSPORT_TCP
     Transport::TCP,
     Transport::TCP_LAZY,
+#endif
 #if GLOO_HAVE_TRANSPORT_TCP_TLS
     Transport::TCP_TLS,
 #endif
+#if GLOO_HAVE_TRANSPORT_UV
     Transport::UV,
+#endif
+#if GLOO_HAVE_TRANSPORT_IBVERBS
+    Transport::IBVERBS,
+#endif
 };
 
 std::shared_ptr<::gloo::transport::Device> createDevice(Transport transport) {
@@ -59,6 +71,17 @@ std::shared_ptr<::gloo::transport::Device> createDevice(Transport transport) {
     return ::gloo::transport::uv::CreateDevice(kDefaultDevice);
 #endif
   }
+#endif
+#if GLOO_HAVE_TRANSPORT_IBVERBS
+  if (transport == Transport::IBVERBS) {
+    gloo::transport::ibverbs::attr attr;
+    attr.port = 1;
+    try {
+      return ::gloo::transport::ibverbs::CreateDevice(attr);
+    } catch (const InvalidOperationException& e) {
+      GLOO_INFO("IBVERBS not available: ", e.what());
+    }
+  }
 #endif
   return nullptr;
 }

gloo/test/base_test.h

@@ -33,6 +33,10 @@
 #include "gloo/transport/uv/device.h"
 #endif
 
+#if GLOO_HAVE_TRANSPORT_IBVERBS
+#include "gloo/transport/ibverbs/device.h"
+#endif
+
 namespace gloo {
 namespace test {
 
@@ -64,6 +68,7 @@ enum Transport {
   TCP_TLS,
 #endif
   UV,
+  IBVERBS,
 };
 
 extern const std::vector<Transport> kTransportsForClassAlgorithms;
@@ -117,6 +122,7 @@ class BaseTest : public ::testing::Test {
       // socket address.
       auto device = device_creator(transport);
       if (!device) {
+        GTEST_SKIP() << "Skipping test: transport not available";
         return;
       }
       context->connectFullMesh(store, device);

gloo/test/send_recv_test.cc

@@ -12,8 +12,6 @@
 #include <array>
 #include <unordered_set>
 
-#include "gloo/transport/tcp/unbound_buffer.h"
-
 namespace gloo {
 namespace test {
 namespace {
@@ -515,7 +513,7 @@ INSTANTIATE_TEST_CASE_P(
     SendRecvDefault,
     SendRecvTest,
     ::testing::Combine(
-        ::testing::Values(Transport::TCP, Transport::UV),
+        ::testing::ValuesIn(kTransportsForFunctionAlgorithms),
         ::testing::Values(2, 3, 4, 5, 6, 7, 8),
         ::testing::Values(1)));
 

gloo/test/tcp_test.cc

@@ -1,5 +1,7 @@
 #include <gtest/gtest.h>
 
+#if GLOO_HAVE_TRANSPORT_TCP
+
 #include <gloo/transport/tcp/helpers.h>
 #include <gloo/transport/tcp/loop.h>
 
@@ -34,3 +36,5 @@ TEST(TcpTest, ConnectTimeout) {
 } // namespace tcp
 } // namespace transport
 } // namespace gloo
+
+#endif // GLOO_HAVE_TRANSPORT_TCP

gloo/transport/ibverbs/CMakeLists.txt

@@ -5,6 +5,7 @@ list(APPEND GLOO_TRANSPORT_SRCS
   "${CMAKE_CURRENT_SOURCE_DIR}/device.cc"
   "${CMAKE_CURRENT_SOURCE_DIR}/memory_region.cc"
   "${CMAKE_CURRENT_SOURCE_DIR}/pair.cc"
+  "${CMAKE_CURRENT_SOURCE_DIR}/unbound_buffer.cc"
   )
 
 list(APPEND GLOO_TRANSPORT_HDRS
@@ -14,6 +15,7 @@ list(APPEND GLOO_TRANSPORT_HDRS
   "${CMAKE_CURRENT_SOURCE_DIR}/device.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/memory_region.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/pair.h"
+  "${CMAKE_CURRENT_SOURCE_DIR}/unbound_buffer.h"
   )
 
 set(GLOO_TRANSPORT_SRCS ${GLOO_TRANSPORT_SRCS} PARENT_SCOPE)

gloo/transport/ibverbs/buffer.cc

@@ -30,8 +30,8 @@ Buffer::Buffer(Pair* pair, int slot, void* ptr, size_t size)
       ex_(nullptr) {
   mr_ = ibv_reg_mr(
       pair_->dev_->pd_,
-      ptr_,
-      size_,
+      size == 0 ? static_cast<void*>(&emptyBuf_) : ptr,
+      size == 0 ? sizeof(emptyBuf_) : size,
       IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE);
 
   // Provide hint if the error is EFAULT and nv_peer_mem is not loaded
@@ -59,7 +59,12 @@ Buffer::Buffer(Pair* pair, int slot, void* ptr, size_t size)
 }
 
 Buffer::~Buffer() {
-  GLOO_ENFORCE_EQ(sendPending_, 0, "Destructing buffer expecting completions");
+  std::lock_guard<std::mutex> lock(m_);
+  if (sendPending_ > 0) {
+    GLOO_WARN(
+        "Destructing buffer with pending sends, sendPending_=", sendPending_);
+  }
+
   ibv_dereg_mr(mr_);
 }
 
@@ -167,36 +172,40 @@ void Buffer::waitSend() {
 }
 
 void Buffer::send(size_t offset, size_t length, size_t roffset) {
-  // Can't assert on roffset, since we don't know the size of
-  // the remote buffer. Refactor of initialization code needed
-  // to support this.
-  GLOO_ENFORCE_LE(offset + length, size_);
-
   {
     std::unique_lock<std::mutex> lock(m_);
+
+    // Can't assert on roffset, since we don't know the size of
+    // the remote buffer. Refactor of initialization code needed
+    // to support this.
+    GLOO_ENFORCE_LE(offset + length, size_);
+
     checkErrorState();
-  }
 
-  if (debug_) {
-    std::cout << "[" << getpid() << "] ";
-    std::cout << "send " << length << " bytes";
-    std::cout << std::endl;
+    if (debug_) {
+      std::cout << "[" << getpid() << "] ";
+      std::cout << "send " << length << " bytes";
+      std::cout << std::endl;
+    }
+
+    // Increment number of sends in flight
+    sendPending_++;
   }
 
-  // Increment number of sends in flight
-  sendPending_++;
+  // Release lock before calling into the pair to avoid deadlock.
 
   pair_->send(this, offset, length, roffset);
 }
 
-void Buffer::handleCompletion(struct ibv_wc* wc) {
+void Buffer::handleCompletion(int rank, struct ibv_wc* wc) {
+  std::unique_lock<std::mutex> lock(m_);
+
   if (wc->opcode & IBV_WC_RECV) {
     if (debug_) {
       std::cout << "[" << getpid() << "] ";
       std::cout << "recv " << wc->byte_len << " bytes";
       std::cout << std::endl;
     }
-    std::unique_lock<std::mutex> lock(m_);
     recvCompletions_++;
     recvCv_.notify_one();
   } else if (wc->opcode == IBV_WC_RDMA_WRITE) {
@@ -205,7 +214,6 @@ void Buffer::handleCompletion(struct ibv_wc* wc) {
       std::cout << "send complete";
       std::cout << std::endl;
     }
-    std::unique_lock<std::mutex> lock(m_);
     sendCompletions_++;
     sendPending_--;
     sendCv_.notify_one();

gloo/transport/ibverbs/buffer.h

@@ -24,7 +24,7 @@ namespace gloo {
 namespace transport {
 namespace ibverbs {
 
-class Buffer : public ::gloo::transport::Buffer {
+class Buffer : public ::gloo::transport::Buffer, public BufferHandler {
  public:
   virtual ~Buffer();
 
@@ -33,18 +33,26 @@ class Buffer : public ::gloo::transport::Buffer {
   virtual void waitRecv() override;
   virtual void waitSend() override;
 
-  void handleCompletion(struct ibv_wc* wc);
+  void handleCompletion(int rank, struct ibv_wc* wc) override;
 
-  void signalError(const std::exception_ptr& ex);
+  void signalError(const std::exception_ptr& ex) override;
   void checkErrorState();
 
+  bool isPeristentHandler() override {
+    return true;
+  }
+
  protected:
   // May only be constructed from helper function in pair.cc
   Buffer(Pair* pair, int slot, void* ptr, size_t size);
 
   Pair* pair_;
 
+  // Empty buffer to use when a nullptr buffer is created.
+  char emptyBuf_[1];
+
   struct ibv_mr* mr_;
+  std::unique_ptr<struct ibv_mr> peerMr_;
 
   std::mutex m_;
   std::condition_variable recvCv_;

gloo/transport/ibverbs/context.cc

@@ -11,6 +11,7 @@
 #include "gloo/common/error.h"
 #include "gloo/transport/ibverbs/device.h"
 #include "gloo/transport/ibverbs/pair.h"
+#include "gloo/transport/ibverbs/unbound_buffer.h"
 
 namespace gloo {
 namespace transport {
@@ -23,16 +24,26 @@ Context::~Context() {}
 
 std::unique_ptr<transport::Pair>& Context::createPair(int rank) {
   pairs_[rank] = std::unique_ptr<transport::Pair>(
-      new ibverbs::Pair(device_, getTimeout()));
+      new ibverbs::Pair(rank, device_, getTimeout()));
   return pairs_[rank];
 }
 
 std::unique_ptr<transport::UnboundBuffer> Context::createUnboundBuffer(
     void* ptr,
     size_t size) {
-  GLOO_THROW_INVALID_OPERATION_EXCEPTION(
-      "Unbound buffers not supported yet for ibverbs transport");
-  return std::unique_ptr<transport::UnboundBuffer>();
+  return std::make_unique<UnboundBuffer>(this->shared_from_this(), ptr, size);
+}
+
+void Context::signalException(const std::string& msg) {
+  // The `pairs_` vector is logically constant. After the context and
+  // all of its pairs have been created it is not mutated until the
+  // context is destructed. Therefore, we don't need to acquire this
+  // context's instance lock before looping over `pairs_`.
+  for (auto& pair : pairs_) {
+    if (pair) {
+      reinterpret_cast<ibverbs::Pair*>(pair.get())->signalIoFailure(msg);
+    }
+  }
 }
 
 } // namespace ibverbs

gloo/transport/ibverbs/context.h

@@ -33,10 +33,16 @@ class Context : public ::gloo::transport::Context,
       void* ptr,
       size_t size) override;
 
+  // Set exception on every pair in this context. This is called when
+  // waiting for a send or recv operation on an unbound buffer times
+  // out. All pairs should be signaled and closed in that event.
+  void signalException(const std::string& msg);
+
  protected:
   std::shared_ptr<Device> device_;
 
   friend class Pair;
+  friend class UnboundBuffer;
 };
 
 } // namespace ibverbs