include/gloo/barrier_all_to_all.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 "gloo/barrier.h"
0012
0013 namespace gloo {
0014
0015 class BarrierAllToAll : public Barrier {
0016  public:
0017   explicit BarrierAllToAll(const std::shared_ptr<Context>& context)
0018       : Barrier(context) {
0019     // Create send/recv buffers for every peer
0020     auto slot = this->context_->nextSlot();
0021     for (auto i = 0; i < this->contextSize_; i++) {
0022       // Skip self
0023       if (i == this->contextRank_) {
0024         continue;
0025       }
0026
0027       auto& pair = this->getPair(i);
0028       auto sdata = std::unique_ptr<int>(new int);
0029       auto sbuf = pair->createSendBuffer(slot, sdata.get(), sizeof(int));
0030       sendBuffersData_.push_back(std::move(sdata));
0031       sendBuffers_.push_back(std::move(sbuf));
0032       auto rdata = std::unique_ptr<int>(new int);
0033       auto rbuf = pair->createRecvBuffer(slot, rdata.get(), sizeof(int));
0034       recvBuffersData_.push_back(std::move(rdata));
0035       recvBuffers_.push_back(std::move(rbuf));
0036     }
0037   }
0038
0039   void run() {
0040     // Notify peers
0041     for (auto& buffer : sendBuffers_) {
0042       buffer->send();
0043     }
0044     // Wait for notification from peers
0045     for (auto& buffer : recvBuffers_) {
0046       buffer->waitRecv();
0047     }
0048   }
0049
0050  protected:
0051   std::vector<std::unique_ptr<int>> sendBuffersData_;
0052   std::vector<std::unique_ptr<transport::Buffer>> sendBuffers_;
0053   std::vector<std::unique_ptr<int>> recvBuffersData_;
0054   std::vector<std::unique_ptr<transport::Buffer>> recvBuffers_;
0055 };
0056
0057 } // namespace gloo