CollectiveBalancedReorder.hpp

Go to the documentation of this file.

// Copyright (c) 2021 Graphcore Ltd. All rights reserved.
#ifndef INCLUDE_GCL_COLLECTIVEBALANCEDREORDER_HPP
#define INCLUDE_GCL_COLLECTIVEBALANCEDREORDER_HPP
 
#include <poplar/DebugContext.hpp>
#include <poplar/Graph.hpp>
#include <poplar/Interval.hpp>
#include <poplar/Tensor.hpp>
#include <poplar/TensorRearranger.hpp>
#include <poplar/Type.hpp>
#include <cstdint>
#include <string>
#include <vector>
 
#if __cplusplus >= 201603L
#define GCL_NO_DISCARD [[nodiscard]]
#else
#define GCL_NO_DISCARD
#endif
 
namespace gcl {
 
class CollectiveBalancedHostRearrangement {
private:
  template <typename ElementType>
  void rearrangeImpl(const ElementType *in, std::size_t inLen, ElementType *out,
                     std::size_t outLen, bool refToGathered) const;
 
public:
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
  // In the deprecation period we need to ensure none of the
  // default implementations warn about using the deprecated vars.
  CollectiveBalancedHostRearrangement() = default;
  ~CollectiveBalancedHostRearrangement() = default;
  CollectiveBalancedHostRearrangement(
      const CollectiveBalancedHostRearrangement &) = default;
  CollectiveBalancedHostRearrangement(
      CollectiveBalancedHostRearrangement &&) noexcept = default;
  CollectiveBalancedHostRearrangement &
  operator=(const CollectiveBalancedHostRearrangement &) = default;
  CollectiveBalancedHostRearrangement &
  operator=(CollectiveBalancedHostRearrangement &&) noexcept = default;
#pragma GCC diagnostic pop
 
  // clang-format off
  GC_DEPRECATED_MSG("Use rearrangeForCollective(const void*, std::size_t, void*, std::size_t, std::size_t) instead ")
  // clang-format on
  void rearrangeForCollective(const void *in, void *out,
                              int64_t elemByteSize) const;
 
  void rearrangeForCollective(const void *in, std::size_t inSize, void *out,
                              std::size_t outSize,
                              std::size_t elemByteSize) const;
 
  template <typename T>
  void rearrangeForCollective(const std::vector<T> &in, std::vector<T> &out,
                              std::size_t elemByteSize = sizeof(T)) const {
    rearrangeForCollective(in.data(), in.size() * sizeof(T), out.data(),
                           out.size() * sizeof(T), elemByteSize);
  }
 
  // clang-format off
  GC_DEPRECATED_MSG("Use undoRearrangeForCollective(const void*, std::size_t, void*, std::size_t, std::size_t) instead ")
  // clang-format on
  void undoRearrangeForCollective(const void *in, void *out,
                                  int64_t elemByteSize) const;
 
  template <typename T>
  void undoRearrangeForCollective(const std::vector<T> &in, std::vector<T> &out,
                                  std::size_t elemByteSize = sizeof(T)) const {
    undoRearrangeForCollective(in.data(), in.size() * sizeof(T), out.data(),
                               out.size() * sizeof(T), elemByteSize);
  }
 
  void undoRearrangeForCollective(const void *in, std::size_t inSize, void *out,
                                  std::size_t outSize,
                                  std::size_t elemByteSize) const;
 
  // NOLINTNEXTLINE
  GCL_NO_DISCARD size_t getNumRearrangedTensorElems() const;
 
  // clang-format off
  GC_DEPRECATED_MSG("Use rearrangeForCollective or undoRearrangeForCollective instead")
  // clang-format on
  void rearrange(const void *in, void *out, int64_t elemByteSize,
                 bool refToGathered) const;
 
  GCL_NO_DISCARD unsigned getReplicationFactor() const;
 
  void setReplicationFactor(unsigned replicationFactor);
 
  GCL_NO_DISCARD std::size_t getTotalElementsPerReplica() const;
 
  void setTotalElementsPerReplica(std::size_t totalElementsPerReplica);
 
  GCL_NO_DISCARD const std::vector<poplar::Interval> &
  getGatheredToRefSlices() const;
 
  void setGatheredToRefSlices(std::vector<poplar::Interval> slices);
 
  GCL_NO_DISCARD const std::vector<uint32_t> &getElementMap() const;
 
  GC_DEPRECATED_MSG("Use get/setReplicationFactor() instead")
  // NOLINTNEXTLINE(readability-identifier-naming)
  unsigned replicationFactor = 0;
 
  GC_DEPRECATED_MSG("Use get/setTotalElementsPerReplica() instead")
  // NOLINTNEXTLINE(readability-identifier-naming)
  std::size_t totalElementsPerReplica = 0;
 
  GC_DEPRECATED_MSG("Use set/getGatheredToRefSlices() instead")
  // NOLINTNEXTLINE(readability-identifier-naming)
  std::vector<poplar::Interval> gatheredToRefSlices;
 
  GC_DEPRECATED_MSG("Use getElementMap() instead")
  // NOLINTNEXTLINE(readability-identifier-naming)
  std::vector<uint32_t> elementMap;
 
private:
  void rearrange(const void *in, std::size_t inSize, void *out,
                 std::size_t outSize, std::size_t elemByteSize,
                 bool refToGathered) const;
 
  void update(
      bool allowElementMap, size_t numGatheredToRefSlicesVecEntries,
      const std::vector<std::vector<poplar::Interval>> &gatheredToRefSlicesVec);
 
  friend class CollectiveBalancedReorder;
};
 
class CollectiveBalancedReorder {
public:
  CollectiveBalancedReorder(poplar::Graph &graph_, poplar::Tensor tensor_,
                            unsigned replicationFactor_,
                            const poplar::DebugNameAndId &dnai_,
                            bool allowElementMap = false,
                            unsigned grainSize = 1);
 
  GCL_NO_DISCARD poplar::Tensor createReplicaSlice(const poplar::Type &type);
 
  GCL_NO_DISCARD poplar::Tensor
  createCollectivesTensor(const poplar::Type &type,
                          const std::string &debugPrefix);
 
  // NOLINTNEXTLINE(modernize-use-nodiscard)
  GCL_NO_DISCARD poplar::Tensor
  undoRearrangeForCollective(const poplar::Tensor &tensor) const;
 
  // NOLINTNEXTLINE(modernize-use-nodiscard)
  GCL_NO_DISCARD std::vector<std::size_t> getReferenceShape() const {
    return mReferenceTensor.shape();
  }
 
  // NOLINTNEXTLINE(modernize-use-nodiscard)
  const CollectiveBalancedHostRearrangement &getHostRearrangement() const {
    return mHostRearrangement;
  }
 
  void zeroPaddingInCollectiveTensor(poplar::Tensor &collectiveTensor,
                                     poplar::program::Sequence &prog) const;
 
private:
  void rearrange(const void *in, void *out, std::size_t elemByteSize,
                 bool refToGathered) const;
 
  poplar::Graph &mGraph;
 
  unsigned mReplicationFactor;
 
  std::vector<std::size_t> mNumReplicaElementsPerTile;
  std::vector<std::size_t> mElementsPerTile;
  std::vector<poplar::Interval> mGatheredToSimplifiedRefSlices;
  poplar::Tensor mReferenceTensor;
  poplar::TensorRearranger mSimplifier;
 
  CollectiveBalancedHostRearrangement mHostRearrangement;
  const poplar::DebugNameAndId mDnai;
};
 
} // namespace gcl
#endif // INCLUDE_GCL_COLLECTIVEBALANCEDREORDER_HPP