gloo/rendezvous/context.h

0001 /**
0002  * Copyright (c) 2017-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 <memory>
0012 #include <string>
0013 #include <vector>
0014
0015 #include "gloo/common/error.h"
0016 #include "gloo/common/store.h"
0017 #include "gloo/context.h"
0018 #include "gloo/rendezvous/store.h"
0019 #include "gloo/transport/address.h"
0020 #include "gloo/transport/device.h"
0021
0022 namespace gloo {
0023 namespace rendezvous {
0024
0025 class ContextFactory;
0026
0027 class Context : public ::gloo::Context {
0028  public:
0029   Context(int rank, int size, int base = 2);
0030   virtual ~Context();
0031
0032   void connectFullMesh(
0033       Store& store,
0034       std::shared_ptr<transport::Device>& dev);
0035
0036  protected:
0037   friend class ContextFactory;
0038 };
0039
0040 class ContextFactory {
0041  public:
0042   static constexpr auto kMaxAddressSize =
0043       ::gloo::transport::Address::kMaxByteSize;
0044
0045   explicit ContextFactory(std::shared_ptr<::gloo::Context> backingContext);
0046
0047   std::shared_ptr<::gloo::Context> makeContext(
0048     std::shared_ptr<transport::Device>& dev);
0049
0050  protected:
0051   std::shared_ptr<::gloo::Context> backingContext_;
0052
0053   std::vector<std::vector<char>> recvData_;
0054   std::vector<std::vector<char>> sendData_;
0055
0056   std::vector<std::unique_ptr<transport::Buffer>> recvBuffers_;
0057   std::vector<std::unique_ptr<transport::Buffer>> sendBuffers_;
0058
0059   std::vector<int> recvNotificationData_;
0060   std::vector<std::unique_ptr<transport::Buffer>> recvNotificationBuffers_;
0061
0062   std::vector<int> sendNotificationData_;
0063   std::vector<std::unique_ptr<transport::Buffer>> sendNotificationBuffers_;
0064 };
0065
0066
0067 } // namespace rendezvous
0068
0069 } // namespace gloo