include/gloo/allgather.h

0001 /**
0002  * Copyright (c) 2018-present, Facebook, Inc.
0003  * All rights reserved.
0004  *
0005  * This source code is licensed under the BSD-style license found in the
0006  * LICENSE file in the root directory of this source tree.
0007  */
0008
0009 #pragma once
0010
0011 #include "gloo/context.h"
0012 #include "gloo/transport/unbound_buffer.h"
0013
0014 namespace gloo {
0015
0016 class AllgatherOptions {
0017  public:
0018   explicit AllgatherOptions(const std::shared_ptr<Context>& context)
0019       : context(context), timeout(context->getTimeout()) {}
0020
0021   template <typename T>
0022   void setInput(std::unique_ptr<transport::UnboundBuffer> buf) {
0023     elementSize = sizeof(T);
0024     in = std::move(buf);
0025   }
0026
0027   template <typename T>
0028   void setInput(T* ptr, size_t elements) {
0029     elementSize = sizeof(T);
0030     in = context->createUnboundBuffer(ptr, elements * sizeof(T));
0031   }
0032
0033   template <typename T>
0034   void setOutput(std::unique_ptr<transport::UnboundBuffer> buf) {
0035     elementSize = sizeof(T);
0036     out = std::move(buf);
0037   }
0038
0039   template <typename T>
0040   void setOutput(T* ptr, size_t elements) {
0041     elementSize = sizeof(T);
0042     out = context->createUnboundBuffer(ptr, elements * sizeof(T));
0043   }
0044
0045   void setTag(uint32_t tag) {
0046     this->tag = tag;
0047   }
0048
0049   void setTimeout(std::chrono::milliseconds timeout) {
0050     this->timeout = timeout;
0051   }
0052
0053  protected:
0054   std::shared_ptr<Context> context;
0055   std::unique_ptr<transport::UnboundBuffer> in;
0056   std::unique_ptr<transport::UnboundBuffer> out;
0057
0058   // Number of bytes per element.
0059   size_t elementSize = 0;
0060
0061   // Tag for this operation.
0062   // Must be unique across operations executing in parallel.
0063   uint32_t tag = 0;
0064
0065   // End-to-end timeout for this operation.
0066   std::chrono::milliseconds timeout;
0067
0068   friend void allgather(AllgatherOptions&);
0069 };
0070
0071 void allgather(AllgatherOptions& opts);
0072
0073 } // namespace gloo