include/gloo/gather.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 GatherOptions {
0017  public:
0018   explicit GatherOptions(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 setRoot(int root) {
0046     this->root = root;
0047   }
0048
0049   void setTag(uint32_t tag) {
0050     this->tag = tag;
0051   }
0052
0053   void setTimeout(std::chrono::milliseconds timeout) {
0054     this->timeout = timeout;
0055   }
0056
0057  protected:
0058   std::shared_ptr<Context> context;
0059   std::unique_ptr<transport::UnboundBuffer> in;
0060   std::unique_ptr<transport::UnboundBuffer> out;
0061
0062   // Number of bytes per element.
0063   size_t elementSize = 0;
0064
0065   // Rank of receiving process.
0066   int root = -1;
0067
0068   // Tag for this operation.
0069   // Must be unique across operations executing in parallel.
0070   uint32_t tag = 0;
0071
0072   // End-to-end timeout for this operation.
0073   std::chrono::milliseconds timeout;
0074
0075   friend void gather(GatherOptions&);
0076 };
0077
0078 void gather(GatherOptions& opts);
0079
0080 } // namespace gloo