include/gloo/allgatherv.h

0001 /**
0002  * Copyright (c) 2019-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 AllgathervOptions {
0017  public:
0018   explicit AllgathervOptions(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     setInput(std::move(buf), sizeof(T));
0024   }
0025
0026   template <typename T>
0027   void setInput(T* ptr, size_t elements) {
0028     setInput(static_cast<void*>(ptr), elements, sizeof(T));
0029   }
0030
0031   template <typename T>
0032   void setOutput(
0033       std::unique_ptr<transport::UnboundBuffer> buf,
0034       std::vector<size_t> elements) {
0035     setOutput(std::move(buf), std::move(elements), sizeof(T));
0036   }
0037
0038   template <typename T>
0039   void setOutput(T* ptr, std::vector<size_t> elements) {
0040     setOutput(static_cast<void*>(ptr), std::move(elements), sizeof(T));
0041   }
0042
0043   void setTag(uint32_t tag) {
0044     this->tag = tag;
0045   }
0046
0047   void setTimeout(std::chrono::milliseconds timeout) {
0048     this->timeout = timeout;
0049   }
0050
0051  protected:
0052   std::shared_ptr<Context> context;
0053   std::unique_ptr<transport::UnboundBuffer> in;
0054   std::unique_ptr<transport::UnboundBuffer> out;
0055
0056   // Number of elements per rank in the output.
0057   std::vector<size_t> elements;
0058
0059   // Number of bytes per element.
0060   size_t elementSize = 0;
0061
0062   // Tag for this operation.
0063   // Must be unique across operations executing in parallel.
0064   uint32_t tag = 0;
0065
0066   // End-to-end timeout for this operation.
0067   std::chrono::milliseconds timeout;
0068
0069   // Set element size, or check the argument is equal to the current value.
0070   void setElementSize(size_t elementSize);
0071
0072   // Untemplated implementation of setInput on unbound buffer.
0073   void setInput(
0074       std::unique_ptr<transport::UnboundBuffer> buf,
0075       size_t elementSize);
0076
0077   // Untemplated implementation of setInput on opaque pointer.
0078   void setInput(void* ptr, size_t elements, size_t elementSize);
0079
0080   // Untemplated implementation of setOutput on unbound buffer.
0081   void setOutput(
0082       std::unique_ptr<transport::UnboundBuffer> buf,
0083       std::vector<size_t> elements,
0084       size_t elementSize);
0085
0086   // Untemplated implementation of setOutput on opaque pointer.
0087   void setOutput(void* ptr, std::vector<size_t> elements, size_t elementSize);
0088
0089   friend void allgatherv(AllgathervOptions&);
0090 };
0091
0092 void allgatherv(AllgathervOptions& opts);
0093
0094 } // namespace gloo