gloo/ibverbs: support put/get operations (#450)

Summary:
Pull Request resolved: #450

This adds put/get operations to the gloo ibverbs backend. This is done by adding new put/get operations to UnboundBuffer/Pair. These operate using an opaque "RemoteKey" class that is backend specific.

We should be able to emulate remote operations using the TCP backend as well though we only have support in the ibverbs backend.

Reviewed By: tianfengfrank, fduwjj

Differential Revision: D75919074

Author: d4l3k

gloo/context.h

@@ -12,6 +12,7 @@
 #include <memory>
 #include <vector>
 
+#include <gloo/transport/context.h>
 #include <gloo/transport/pair.h>
 
 namespace gloo {
@@ -51,6 +52,11 @@ class Context {
 
   std::chrono::milliseconds getTimeout() const;
 
+  std::unique_ptr<transport::RemoteKey> deserializeRemoteKey(
+      const std::string& serialized) {
+    return transportContext_->deserializeRemoteKey(serialized);
+  }
+
  protected:
   std::shared_ptr<transport::Device> device_;
   std::shared_ptr<transport::Context> transportContext_;

gloo/test/CMakeLists.txt

@@ -14,6 +14,7 @@ set(GLOO_TEST_SRCS
   "${CMAKE_CURRENT_SOURCE_DIR}/main.cc"
   "${CMAKE_CURRENT_SOURCE_DIR}/memory_test.cc"
   "${CMAKE_CURRENT_SOURCE_DIR}/reduce_test.cc"
+  "${CMAKE_CURRENT_SOURCE_DIR}/remote_key_test.cc"
   "${CMAKE_CURRENT_SOURCE_DIR}/send_recv_test.cc"
   )
 set(GLOO_TEST_LIBRARIES)

gloo/test/base_test.cc

@@ -26,6 +26,12 @@ const std::vector<Transport> kTransportsForClassAlgorithms = {
 #endif
 };
 
+const std::vector<Transport> kTransportsForRDMA = {
+#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.

gloo/test/base_test.h

@@ -73,6 +73,7 @@ enum Transport {
 
 extern const std::vector<Transport> kTransportsForClassAlgorithms;
 extern const std::vector<Transport> kTransportsForFunctionAlgorithms;
+extern const std::vector<Transport> kTransportsForRDMA;
 
 std::shared_ptr<::gloo::transport::Device> createDevice(Transport transport);
 

gloo/test/remote_key_test.cc

@@ -0,0 +1,176 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <array>
+#include <functional>
+#include <vector>
+
+#include "gloo/allgather.h"
+#include "gloo/common/common.h"
+#include "gloo/test/base_test.h"
+
+namespace gloo {
+namespace test {
+namespace {
+
+// Test parameterization.
+using Param = std::tuple<Transport, int, int>;
+
+std::vector<std::unique_ptr<gloo::transport::RemoteKey>> exchangeKeys(
+    std::shared_ptr<Context>& context,
+    const std::unique_ptr<gloo::transport::RemoteKey>& key) {
+  auto selfRemoteKey = key->serialize();
+  GLOO_ENFORCE_GT(selfRemoteKey.size(), 0);
+
+  std::array<char, 1024> keyBuf;
+  for (int i = 0; i < selfRemoteKey.size(); ++i) {
+    keyBuf[i] = selfRemoteKey[i];
+  }
+  keyBuf[selfRemoteKey.size()] = '\0';
+
+  std::unique_ptr<char[]> outPtr{
+      gloo::make_unique<char[]>(keyBuf.size() * context->size)};
+
+  AllgatherOptions opts(context);
+  opts.setInput(&keyBuf[0], keyBuf.size());
+  opts.setOutput(outPtr.get(), keyBuf.size() * context->size);
+  allgather(opts);
+
+  std::vector<std::unique_ptr<gloo::transport::RemoteKey>> remoteKeys;
+
+  for (int i = 0; i < context->size; i++) {
+    std::string remoteKeyStr{&outPtr[i * keyBuf.size()]};
+    GLOO_ENFORCE_GT(remoteKeyStr.size(), 0);
+    auto remoteKey = context->deserializeRemoteKey(remoteKeyStr);
+    GLOO_ENFORCE(remoteKey.get() != nullptr);
+
+    remoteKeys.push_back(std::move(remoteKey));
+  }
+
+  return remoteKeys;
+}
+
+// Test fixture.
+class RemoteKeyTest : public BaseTest,
+                      public ::testing::WithParamInterface<Param> {};
+
+TEST_P(RemoteKeyTest, Get) {
+  const auto transport = std::get<0>(GetParam());
+  const auto contextSize = std::get<1>(GetParam());
+  const auto dataSize = std::get<2>(GetParam());
+
+  Barrier barrier(contextSize);
+
+  spawn(transport, contextSize, [&](std::shared_ptr<Context> context) {
+    int rank = context->rank;
+    std::unique_ptr<char[]> sharedPtr{gloo::make_unique<char[]>(dataSize)};
+    for (int i = 0; i < dataSize; ++i) {
+      sharedPtr[i] = rank;
+    }
+    auto sharedBuf = context->createUnboundBuffer(sharedPtr.get(), dataSize);
+
+    std::unique_ptr<char[]> localPtr{gloo::make_unique<char[]>(dataSize)};
+    auto localBuf = context->createUnboundBuffer(localPtr.get(), dataSize);
+
+    auto remoteKeys = exchangeKeys(context, sharedBuf->getRemoteKey());
+
+    for (int i = 0; i < contextSize; ++i) {
+      if (i == rank) {
+        continue;
+      }
+
+      localBuf->get(*remoteKeys.at(i), context->nextSlot(), 0, 0, dataSize);
+      localBuf->waitSend();
+
+      for (int j = 0; j < dataSize; ++j) {
+        ASSERT_EQ(localPtr[j], i);
+      }
+    }
+
+    barrier.wait();
+
+    // Do bound checking
+    auto testRank = (rank + 1) % contextSize;
+    auto& testKey = remoteKeys.at(testRank);
+    EXPECT_THROW(
+        localBuf->get(*testKey, context->nextSlot(), 1000000000, 0, 1),
+        gloo::EnforceNotMet);
+    EXPECT_THROW(
+        localBuf->get(*testKey, context->nextSlot(), 0, 1000000000, 1),
+        gloo::EnforceNotMet);
+    EXPECT_THROW(
+        localBuf->get(*testKey, context->nextSlot(), 0, 0, 1000000000),
+        gloo::EnforceNotMet);
+  });
+}
+
+TEST_P(RemoteKeyTest, Put) {
+  const auto transport = std::get<0>(GetParam());
+  const auto contextSize = std::get<1>(GetParam());
+
+  Barrier barrier(contextSize);
+
+  spawn(transport, contextSize, [&](std::shared_ptr<Context> context) {
+    int rank = context->rank;
+    std::unique_ptr<char[]> exportPtr{gloo::make_unique<char[]>(contextSize)};
+    auto exportBuf = context->createUnboundBuffer(exportPtr.get(), contextSize);
+
+    std::unique_ptr<char[]> localPtr{gloo::make_unique<char[]>(contextSize)};
+    for (int i = 0; i < contextSize; ++i) {
+      localPtr[i] = rank;
+    }
+    auto localBuf = context->createUnboundBuffer(localPtr.get(), contextSize);
+
+    auto remoteKeys = exchangeKeys(context, exportBuf->getRemoteKey());
+
+    for (int i = 0; i < contextSize; ++i) {
+      if (i == rank) {
+        continue;
+      }
+
+      localBuf->put(*remoteKeys.at(i), context->nextSlot(), rank, rank, 1);
+      localBuf->waitSend();
+    }
+
+    barrier.wait();
+
+    for (int j = 0; j < contextSize; ++j) {
+      if (j == rank) {
+        continue;
+      }
+      ASSERT_EQ(exportPtr[j], j);
+    }
+
+    // Do bound checking
+    auto testRank = (rank + 1) % contextSize;
+    auto& testKey = remoteKeys.at(testRank);
+    EXPECT_THROW(
+        localBuf->put(*testKey, context->nextSlot(), 1000000000, 0, 1),
+        gloo::EnforceNotMet);
+    EXPECT_THROW(
+        localBuf->put(*testKey, context->nextSlot(), 0, 1000000000, 1),
+        gloo::EnforceNotMet);
+    EXPECT_THROW(
+        localBuf->put(*testKey, context->nextSlot(), 0, 0, 1000000000),
+        gloo::EnforceNotMet);
+  });
+}
+
+INSTANTIATE_TEST_CASE_P(
+    RemoteKeyTestBasics,
+    RemoteKeyTest,
+    ::testing::Combine(
+        ::testing::ValuesIn(kTransportsForRDMA),
+        ::testing::Values(2, 4),
+        ::testing::Values(0, 1024, 1000000)));
+
+GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(RemoteKeyTest);
+
+} // namespace
+} // namespace test
+} // namespace gloo

gloo/transport/CMakeLists.txt

@@ -13,6 +13,7 @@ set(GLOO_TRANSPORT_HDRS
   "${CMAKE_CURRENT_SOURCE_DIR}/context.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/device.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/pair.h"
+  "${CMAKE_CURRENT_SOURCE_DIR}/remote_key.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/unbound_buffer.h"
   )
 

gloo/transport/context.h

@@ -68,6 +68,11 @@ class Context {
     return timeout_;
   }
 
+  virtual std::unique_ptr<RemoteKey> deserializeRemoteKey(
+      const std::string& serialized) {
+    throw std::runtime_error("Not implemented");
+  }
+
  protected:
   // Protects access to the pending operations and expected
   // notifications vectors. These vectors can only be mutated by an

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}/remote_key.cc"
   "${CMAKE_CURRENT_SOURCE_DIR}/unbound_buffer.cc"
   )
 
@@ -15,6 +16,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}/remote_key.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/unbound_buffer.h"
   )
 

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/remote_key.h"
 #include "gloo/transport/ibverbs/unbound_buffer.h"
 
 namespace gloo {
@@ -46,6 +47,11 @@ void Context::signalException(const std::string& msg) {
   }
 }
 
+std::unique_ptr<::gloo::transport::RemoteKey> Context::deserializeRemoteKey(
+    const std::string& serialized) {
+  return RemoteKey::deserialize(serialized);
+}
+
 } // namespace ibverbs
 } // namespace transport
 } // namespace gloo

gloo/transport/ibverbs/context.h

@@ -38,6 +38,9 @@ class Context : public ::gloo::transport::Context,
   // out. All pairs should be signaled and closed in that event.
   void signalException(const std::string& msg);
 
+  std::unique_ptr<::gloo::transport::RemoteKey> deserializeRemoteKey(
+      const std::string& serialized) override;
+
  protected:
   std::shared_ptr<Device> device_;