include/gloo/allreduce_local.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/algorithm.h"
0012 #include "gloo/context.h"
0013
0014 namespace gloo {
0015
0016 template <typename T>
0017 class AllreduceLocal : public Algorithm {
0018  public:
0019   AllreduceLocal(
0020       const std::shared_ptr<Context>& context,
0021       const std::vector<T*>& ptrs,
0022       const int count,
0023       const ReductionFunction<T>* fn = ReductionFunction<T>::sum);
0024
0025   virtual ~AllreduceLocal() = default;
0026
0027   virtual void run() override;
0028
0029  protected:
0030   std::vector<T*> ptrs_;
0031   const int count_;
0032   const int bytes_;
0033   const ReductionFunction<T>* fn_;
0034 };
0035
0036 } // namespace gloo