1/**
2 * Copyright (c) 2019-present, Facebook, Inc.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9#pragma once
10
11#include "gloo/context.h"
12#include "gloo/transport/unbound_buffer.h"
13
14namespace gloo {
15
16class AllgathervOptions {
17 public:
18 explicit AllgathervOptions(const std::shared_ptr<Context>& context)
19 : context(context), timeout(context->getTimeout()) {}
20
21 template <typename T>
22 void setInput(std::unique_ptr<transport::UnboundBuffer> buf) {
23 setInput(std::move(buf), sizeof(T));
24 }
25
26 template <typename T>
27 void setInput(T* ptr, size_t elements) {
28 setInput(static_cast<void*>(ptr), elements, sizeof(T));
29 }
30
31 template <typename T>
32 void setOutput(
33 std::unique_ptr<transport::UnboundBuffer> buf,
34 std::vector<size_t> elements) {
35 setOutput(std::move(buf), std::move(elements), sizeof(T));
36 }
37
38 template <typename T>
39 void setOutput(T* ptr, std::vector<size_t> elements) {
40 setOutput(static_cast<void*>(ptr), std::move(elements), sizeof(T));
41 }
42
43 void setTag(uint32_t tag) {
44 this->tag = tag;
45 }
46
47 void setTimeout(std::chrono::milliseconds timeout) {
48 this->timeout = timeout;
49 }
50
51 protected:
52 std::shared_ptr<Context> context;
53 std::unique_ptr<transport::UnboundBuffer> in;
54 std::unique_ptr<transport::UnboundBuffer> out;
55
56 // Number of elements per rank in the output.
57 std::vector<size_t> elements;
58
59 // Number of bytes per element.
60 size_t elementSize = 0;
61
62 // Tag for this operation.
63 // Must be unique across operations executing in parallel.
64 uint32_t tag = 0;
65
66 // End-to-end timeout for this operation.
67 std::chrono::milliseconds timeout;
68
69 // Set element size, or check the argument is equal to the current value.
70 void setElementSize(size_t elementSize);
71
72 // Untemplated implementation of setInput on unbound buffer.
73 void setInput(
74 std::unique_ptr<transport::UnboundBuffer> buf,
75 size_t elementSize);
76
77 // Untemplated implementation of setInput on opaque pointer.
78 void setInput(void* ptr, size_t elements, size_t elementSize);
79
80 // Untemplated implementation of setOutput on unbound buffer.
81 void setOutput(
82 std::unique_ptr<transport::UnboundBuffer> buf,
83 std::vector<size_t> elements,
84 size_t elementSize);
85
86 // Untemplated implementation of setOutput on opaque pointer.
87 void setOutput(void* ptr, std::vector<size_t> elements, size_t elementSize);
88
89 friend void allgatherv(AllgathervOptions&);
90};
91
92void allgatherv(AllgathervOptions& opts);
93
94} // namespace gloo
95