TensorRearranger.hpp

// Copyright (c) 2020 Graphcore Ltd. All rights reserved.
 
#ifndef poplar_TensorRearranger_hpp
#define poplar_TensorRearranger_hpp
#include <poplar/ArrayRef.hpp>
#include <poplar/Interval.hpp>
#include <poplar/Tensor.hpp>
 
namespace poplar {
namespace core {
class TensorRearranger;
}
 
class TensorRearranger {
public:
  TensorRearranger();
  TensorRearranger(const TensorRearranger &other);
  TensorRearranger(TensorRearranger &&other) noexcept;
  TensorRearranger &operator=(const TensorRearranger &other) &;
  TensorRearranger &operator=(TensorRearranger &&other) &noexcept;
  ~TensorRearranger();
 
  Tensor rearrange(const Tensor &t) const;
 
  Tensor undoRearrangement(const Tensor &t) const;
 
  std::vector<Interval> rearrange(ArrayRef<Interval> is) const;
 
  std::vector<Interval> undoRearrangement(ArrayRef<Interval> is) const;
 
  // Implementation
  TensorRearranger(std::unique_ptr<core::TensorRearranger>);
  core::TensorRearranger &getImpl() const { return *impl; }
  std::unique_ptr<core::TensorRearranger> *getPImpl() { return &impl; }
  bool valid() const { return impl.get() != nullptr; }
 
private:
  std::unique_ptr<core::TensorRearranger> impl;
};
 
} // end namespace poplar
 
#endif // poplar_TensorRearranger_hpp