gloo/ibverbs: support put/get operations

Differential Revision: D75919074

Pull Request resolved: #450

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_;