include/gloo/broadcast.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 BroadcastOptions {
0017  public:
0018   explicit BroadcastOptions(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     this->elements = buf->size / sizeof(T);
0024     this->elementSize = sizeof(T);
0025     this->in = std::move(buf);
0026   }
0027
0028   template <typename T>
0029   void setInput(T* ptr, size_t elements) {
0030     this->elements = elements;
0031     this->elementSize = sizeof(T);
0032     this->in = context->createUnboundBuffer(ptr, elements * sizeof(T));
0033   }
0034
0035   template <typename T>
0036   void setOutput(std::unique_ptr<transport::UnboundBuffer> buf) {
0037     this->elements = buf->size / sizeof(T);
0038     this->elementSize = sizeof(T);
0039     this->out = std::move(buf);
0040   }
0041
0042   template <typename T>
0043   void setOutput(T* ptr, size_t elements) {
0044     this->elements = elements;
0045     this->elementSize = sizeof(T);
0046     this->out = context->createUnboundBuffer(ptr, elements * sizeof(T));
0047   }
0048
0049   void setRoot(int root) {
0050     this->root = root;
0051   }
0052
0053   void setTag(uint32_t tag) {
0054     this->tag = tag;
0055   }
0056
0057   void setTimeout(std::chrono::milliseconds timeout) {
0058     this->timeout = timeout;
0059   }
0060
0061  protected:
0062   std::shared_ptr<Context> context;
0063
0064   // Broadcast has an optional input buffer for the root.
0065   std::unique_ptr<transport::UnboundBuffer> in;
0066
0067   // Broadcast has a mandatory output buffer for all ranks.
0068   std::unique_ptr<transport::UnboundBuffer> out;
0069
0070   // Number of elements.
0071   size_t elements = 0;
0072
0073   // Number of bytes per element.
0074   size_t elementSize = 0;
0075
0076   // Rank of process to broadcast from.
0077   int root = -1;
0078
0079   // Tag for this operation.
0080   // Must be unique across operations executing in parallel.
0081   uint32_t tag = 0;
0082
0083   // End-to-end timeout for this operation.
0084   std::chrono::milliseconds timeout;
0085
0086   friend void broadcast(BroadcastOptions&);
0087 };
0088
0089 void broadcast(BroadcastOptions& opts);
0090
0091 } // namespace gloo