Add new style allgather collective (#137)

Summary:
Like gather in the previous commit, this is a single function and
doesn't hold on to any state. This implements a ring allgather.

Rebase needed after #136 is merged.
Pull Request resolved: #137

Reviewed By: teng-li

Differential Revision: D10376419

Pulled By: pietern

fbshipit-source-id: bc561bff4c8952531bada14dfe0bc2e9d9c68f78

Author: pietern

gloo/CMakeLists.txt

@@ -7,6 +7,7 @@ set(GLOO_HDRS)
 # Compiled sources in root directory
 list(APPEND GLOO_SRCS
   "${CMAKE_CURRENT_SOURCE_DIR}/algorithm.cc"
+  "${CMAKE_CURRENT_SOURCE_DIR}/allgather.cc"
   "${CMAKE_CURRENT_SOURCE_DIR}/allreduce_local.cc"
   "${CMAKE_CURRENT_SOURCE_DIR}/context.cc"
   "${CMAKE_CURRENT_SOURCE_DIR}/gather.cc"
@@ -15,6 +16,7 @@ list(APPEND GLOO_SRCS
 
 list(APPEND GLOO_HDRS
   "${CMAKE_CURRENT_SOURCE_DIR}/algorithm.h"
+  "${CMAKE_CURRENT_SOURCE_DIR}/allgather.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/allgather_ring.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/allreduce_halving_doubling.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/allreduce_bcube.h"

gloo/allgather.cc

@@ -0,0 +1,118 @@
+/**
+ * Copyright (c) 2018-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. An additional grant
+ * of patent rights can be found in the PATENTS file in the same directory.
+ */
+
+#include "gloo/allgather.h"
+
+#include <cstring>
+
+#include "gloo/common/logging.h"
+#include "gloo/types.h"
+
+namespace gloo {
+
+void allgather(
+    const std::shared_ptr<Context>& context,
+    AllgatherOptions& opts) {
+  std::unique_ptr<transport::UnboundBuffer> tmpInBuffer;
+  std::unique_ptr<transport::UnboundBuffer> tmpOutBuffer;
+  transport::UnboundBuffer* in = nullptr;
+  transport::UnboundBuffer* out = nullptr;
+  const auto slot = Slot::build(kAllgatherSlotPrefix, opts.tag);
+
+  // Sanity checks
+  GLOO_ENFORCE(opts.elementSize > 0);
+  const auto recvRank = (context->size + context->rank - 1) % context->size;
+  GLOO_ENFORCE(
+      context->getPair(recvRank),
+      "missing connection between rank " + std::to_string(context->rank) +
+          " (this process) and rank " + std::to_string(recvRank));
+  const auto sendRank = (context->size + context->rank + 1) % context->size;
+  GLOO_ENFORCE(
+      context->getPair(sendRank),
+      "missing connection between rank " + std::to_string(context->rank) +
+          " (this process) and rank " + std::to_string(sendRank));
+
+  // Figure out pointer to input buffer
+  if (opts.inBuffer) {
+    in = opts.inBuffer.get();
+  } else if (opts.inPtr != nullptr) {
+    GLOO_ENFORCE(opts.inElements > 0);
+    tmpInBuffer = context->createUnboundBuffer(
+        opts.inPtr, opts.inElements * opts.elementSize);
+    in = tmpInBuffer.get();
+  }
+
+  // Figure out pointer to output buffer
+  if (opts.outBuffer) {
+    out = opts.outBuffer.get();
+  } else {
+    GLOO_ENFORCE(opts.outPtr != nullptr);
+    GLOO_ENFORCE(opts.outElements > 0);
+    tmpOutBuffer = context->createUnboundBuffer(
+        opts.outPtr, opts.outElements * opts.elementSize);
+    out = tmpOutBuffer.get();
+  }
+
+  if (in != nullptr) {
+    GLOO_ENFORCE_EQ(out->size, in->size * context->size);
+  } else {
+    GLOO_ENFORCE_EQ(out->size % context->size, 0);
+  }
+
+  const size_t inBytes = out->size / context->size;
+  const size_t outBytes = out->size;
+
+  // If the input buffer is specified, this is NOT an in place operation,
+  // and the output buffer needs to be primed with the input.
+  if (in != nullptr) {
+    memcpy(
+        static_cast<uint8_t*>(out->ptr) + context->rank * in->size,
+        static_cast<uint8_t*>(in->ptr),
+        in->size);
+  }
+
+  // The chunk size may not be divisible by 2; use dynamic lookup.
+  std::array<size_t, 2> chunkSize;
+  chunkSize[0] = inBytes / 2;
+  chunkSize[1] = inBytes - chunkSize[0];
+  std::array<size_t, 2> chunkOffset;
+  chunkOffset[0] = 0;
+  chunkOffset[1] = chunkSize[0];
+
+  for (auto i = 0; i < (context->size - 1) * 2; i++) {
+    const size_t sendSegment = context->size + context->rank - (i / 2);
+    const size_t recvSegment = sendSegment - 1;
+    size_t sendOffset =
+        ((sendSegment * inBytes) + chunkOffset[i & 0x1]) % outBytes;
+    size_t recvOffset =
+        ((recvSegment * inBytes) + chunkOffset[i & 0x1]) % outBytes;
+    size_t size = chunkSize[i & 0x1];
+    if (i < 2) {
+      out->send(sendRank, slot, sendOffset, size);
+      out->recv(recvRank, slot, recvOffset, size);
+      continue;
+    }
+
+    // Wait for pending operations to complete to synchronize with the
+    // previous iteration. Because we kick off two operations before
+    // getting here we always wait for the next-to-last operation.
+    out->waitSend();
+    out->waitRecv();
+    out->send(sendRank, slot, sendOffset, size);
+    out->recv(recvRank, slot, recvOffset, size);
+  }
+
+  // Wait for completes
+  for (auto i = 0; i < 2; i++) {
+    out->waitSend();
+    out->waitRecv();
+  }
+}
+
+} // namespace gloo

gloo/allgather.h

@@ -0,0 +1,42 @@
+/**
+ * Copyright (c) 2018-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. An additional grant
+ * of patent rights can be found in the PATENTS file in the same directory.
+ */
+
+#pragma once
+
+#include "gloo/context.h"
+#include "gloo/transport/unbound_buffer.h"
+
+namespace gloo {
+
+struct AllgatherOptions {
+  // The input and output can either be specified as a unbound buffer
+  // (that can be cached and reused by the caller), or a literal
+  // pointer and number of elements stored at that pointer.
+  //
+  // The operation is executed in place on the output if the input is
+  // set to null. The input for this process is assumed to be at the
+  // location in the output buffer where it would otherwise be.
+  std::unique_ptr<transport::UnboundBuffer> inBuffer;
+  void* inPtr;
+  size_t inElements;
+  std::unique_ptr<transport::UnboundBuffer> outBuffer;
+  void* outPtr;
+  size_t outElements;
+
+  // Number of bytes per element.
+  size_t elementSize;
+
+  // Tag for this gather operation.
+  // Must be unique across operations executing in parallel.
+  uint32_t tag;
+};
+
+void allgather(const std::shared_ptr<Context>& context, AllgatherOptions& opts);
+
+} // namespace gloo

gloo/test/CMakeLists.txt

@@ -1,4 +1,5 @@
 set(GLOO_TEST_SRCS
+  "${CMAKE_CURRENT_SOURCE_DIR}/allgather_test.cc"
   "${CMAKE_CURRENT_SOURCE_DIR}/allreduce_builder_test.cc"
   "${CMAKE_CURRENT_SOURCE_DIR}/allreduce_test.cc"
   "${CMAKE_CURRENT_SOURCE_DIR}/barrier_test.cc"

gloo/test/allgather_test.cc

@@ -11,6 +11,7 @@
 #include <thread>
 #include <vector>
 
+#include "gloo/allgather.h"
 #include "gloo/allgather_ring.h"
 #include "gloo/common/common.h"
 #include "gloo/test/base_test.h"
@@ -109,6 +110,67 @@ INSTANTIATE_TEST_CASE_P(
         ::testing::ValuesIn(genMemorySizes()),
         ::testing::Range(1, 4)));
 
+using NewParam = std::tuple<int, int, bool>;
+
+class AllgatherNewTest : public BaseTest,
+                         public ::testing::WithParamInterface<NewParam> {};
+
+TEST_P(AllgatherNewTest, Default) {
+  auto contextSize = std::get<0>(GetParam());
+  auto dataSize = std::get<1>(GetParam());
+  auto passBuffers = std::get<2>(GetParam());
+
+  auto validate = [dataSize](
+      const std::shared_ptr<Context>& context,
+      Fixture<uint64_t>& output) {
+    const auto ptr = output.getPointer();
+    const auto stride = context->size;
+    for (auto j = 0; j < context->size; j++) {
+      for (auto k = 0; k < dataSize; k++) {
+        ASSERT_EQ(j + k * stride, ptr[k + j * dataSize])
+          << "Mismatch at index " << (k + j * dataSize);
+      }
+    }
+  };
+
+  spawn(contextSize, [&](std::shared_ptr<Context> context) {
+      auto input = Fixture<uint64_t>(context, 1, dataSize);
+      auto output = Fixture<uint64_t>(context, 1, contextSize * dataSize);
+
+      AllgatherOptions opts;
+      opts.elementSize = sizeof(uint64_t);
+
+      if (passBuffers) {
+        // Run with (optionally cached) unbound buffers in options
+        opts.inBuffer = context->createUnboundBuffer(
+            input.getPointer(),
+            dataSize * sizeof(uint64_t));
+        opts.outBuffer = context->createUnboundBuffer(
+            output.getPointer(),
+            contextSize * dataSize * sizeof(uint64_t));
+      } else {
+        // Run with raw pointers and sizes in options
+        opts.inPtr = input.getPointer();
+        opts.inElements = dataSize;
+        opts.outPtr = output.getPointer();
+        opts.outElements = contextSize * dataSize;
+      }
+
+      input.assignValues();
+      allgather(context, opts);
+      validate(context, output);
+    });
+}
+
+INSTANTIATE_TEST_CASE_P(
+    AllgatherNewDefault,
+    AllgatherNewTest,
+    ::testing::Combine(
+        ::testing::Values(2, 4, 7),
+        ::testing::ValuesIn(genMemorySizes()),
+        ::testing::Values(false, true)));
+
+
 } // namespace
 } // namespace test
 } // namespace gloo

gloo/test/base_test.h

@@ -136,6 +136,14 @@ class Fixture {
     }
   }
 
+  void clear() {
+    for (auto i = 0; i < srcs.size(); i++) {
+      for (auto j = 0; j < count; j++) {
+        srcs[i][j] = 0;
+      }
+    }
+  }
+
   void checkBroadcastResult(Fixture<T>& fixture, int root, int rootPointer) {
     // Expected is set to the expected value at ptr[0]
     const auto expected = root * fixture.srcs.size() + rootPointer;

gloo/types.h

@@ -55,6 +55,7 @@ namespace gloo {
 //
 
 constexpr uint8_t kGatherSlotPrefix = 0x01;
+constexpr uint8_t kAllgatherSlotPrefix = 0x02;
 
 class Slot {
  public: