EIC code displayed by LXR master/​include/​gloo/​allreduce.h	Source navigation Diff markup Identifier search General search
	Version: master test Revert   Change

File indexing completed on 2025-01-18 10:00:10

 /**
  * Copyright (c) 2018-present, Facebook, Inc.
  * All rights reserved.
  *
  * This source code is licensed under the BSD-style license found in the
  * LICENSE file in the root directory of this source tree.
  */
 
 #pragma once
 
 #include <functional>
 #include <memory>
 #include <vector>
 
 #include "gloo/context.h"
 #include "gloo/transport/unbound_buffer.h"
 
 namespace gloo {
 
 namespace detail {
 
 struct AllreduceOptionsImpl {
   // This type describes the function to use for element wise reduction.
   //
   // Its arguments are:
   //   1. non-const output pointer
   //   2. const input pointer 1 (may be equal to 1)
   //   3. const input pointer 2 (may be equal to 1)
   //   4. number of elements to reduce.
   //
   // Note that this function is not strictly typed and takes void pointers.
   // This is specifically done to avoid the need for a templated options class
   // and templated algorithm implementations. We found this adds very little
   // value for the increase in compilation time and code size.
   //
   using Func = std::function<void(void*, const void*, const void*, size_t)>;
 
   enum Algorithm {
     UNSPECIFIED = 0,
     RING = 1,
     BCUBE = 2,
   };
 
   explicit AllreduceOptionsImpl(const std::shared_ptr<Context>& context)
       : context(context),
         timeout(context->getTimeout()),
         algorithm(UNSPECIFIED) {}
 
   std::shared_ptr<Context> context;
 
   // End-to-end timeout for this operation.
   std::chrono::milliseconds timeout;
 
   // Algorithm selection.
   Algorithm algorithm;
 
   // Input and output buffers.
   // The output is used as input if input is not specified.
   std::vector<std::unique_ptr<transport::UnboundBuffer>> in;
   std::vector<std::unique_ptr<transport::UnboundBuffer>> out;
 
   // Number of elements.
   size_t elements = 0;
 
   // Number of bytes per element.
   size_t elementSize = 0;
 
   // Reduction function.
   Func reduce;
 
   // Tag for this operation.
   // Must be unique across operations executing in parallel.
   uint32_t tag = 0;
 
   // This is the maximum size of each I/O operation (send/recv) of which
   // two are in flight at all times. A smaller value leads to more
   // overhead and a larger value leads to poor cache behavior.
   static constexpr size_t kMaxSegmentSize = 1024 * 1024;
 
   // Internal use only. This is used to exercise code paths where we
   // have more than 2 segments per rank without making the tests slow
   // (because they would require millions of elements if the default
   // were not configurable).
   size_t maxSegmentSize = kMaxSegmentSize;
 };
 
 } // namespace detail
 
 class AllreduceOptions {
  public:
   using Func = detail::AllreduceOptionsImpl::Func;
   using Algorithm = detail::AllreduceOptionsImpl::Algorithm;
 
   explicit AllreduceOptions(const std::shared_ptr<Context>& context)
       : impl_(context) {}
 
   void setAlgorithm(Algorithm algorithm) {
     impl_.algorithm = algorithm;
   }
 
   template <typename T>
   void setInput(std::unique_ptr<transport::UnboundBuffer> buf) {
     std::vector<std::unique_ptr<transport::UnboundBuffer>> bufs(1);
     bufs[0] = std::move(buf);
     setInputs<T>(std::move(bufs));
   }
 
   template <typename T>
   void setInputs(std::vector<std::unique_ptr<transport::UnboundBuffer>> bufs) {
     impl_.elements = bufs[0]->size / sizeof(T);
     impl_.elementSize = sizeof(T);
     impl_.in = std::move(bufs);
   }
 
   template <typename T>
   void setInput(T* ptr, size_t elements) {
     setInputs(&ptr, 1, elements);
   }
 
   template <typename T>
   void setInputs(std::vector<T*> ptrs, size_t elements) {
     setInputs(ptrs.data(), ptrs.size(), elements);
   }
 
   template <typename T>
   void setInputs(T** ptrs, size_t len, size_t elements) {
     impl_.elements = elements;
     impl_.elementSize = sizeof(T);
     impl_.in.reserve(len);
     for (size_t i = 0; i < len; i++) {
       impl_.in.push_back(
           impl_.context->createUnboundBuffer(ptrs[i], elements * sizeof(T)));
     }
   }
 
   template <typename T>
   void setOutput(std::unique_ptr<transport::UnboundBuffer> buf) {
     std::vector<std::unique_ptr<transport::UnboundBuffer>> bufs(1);
     bufs[0] = std::move(buf);
     setOutputs<T>(std::move(bufs));
   }
 
   template <typename T>
   void setOutputs(std::vector<std::unique_ptr<transport::UnboundBuffer>> bufs) {
     impl_.elements = bufs[0]->size / sizeof(T);
     impl_.elementSize = sizeof(T);
     impl_.out = std::move(bufs);
   }
 
   template <typename T>
   void setOutput(T* ptr, size_t elements) {
     setOutputs(&ptr, 1, elements);
   }
 
   template <typename T>
   void setOutputs(std::vector<T*> ptrs, size_t elements) {
     setOutputs(ptrs.data(), ptrs.size(), elements);
   }
 
   template <typename T>
   void setOutputs(T** ptrs, size_t len, size_t elements) {
     impl_.elements = elements;
     impl_.elementSize = sizeof(T);
     impl_.out.reserve(len);
     for (size_t i = 0; i < len; i++) {
       impl_.out.push_back(
           impl_.context->createUnboundBuffer(ptrs[i], elements * sizeof(T)));
     }
   }
 
   void setReduceFunction(Func fn) {
     impl_.reduce = fn;
   }
 
   void setTag(uint32_t tag) {
     impl_.tag = tag;
   }
 
   void setMaxSegmentSize(size_t maxSegmentSize) {
     impl_.maxSegmentSize = maxSegmentSize;
   }
 
   void setTimeout(std::chrono::milliseconds timeout) {
     impl_.timeout = timeout;
   }
 
  protected:
   detail::AllreduceOptionsImpl impl_;
 
   friend void allreduce(const AllreduceOptions&);
 };
 
 void allreduce(const AllreduceOptions& opts);
 
 } // namespace gloo

This page was automatically generated by the 2.3.7 LXR engine. The LXR team