Tensor.hpp

// Copyright (c) 2016 Graphcore Ltd. All rights reserved.
 
#ifndef poplar_Tensor_hpp
#define poplar_Tensor_hpp
#include <memory>
#include <poplar/ArrayRef.hpp>
#include <poplar/Interval.hpp>
#include <poplar/Type.hpp>
#include <poplar/TypeTraits.hpp>
#include <vector>
 
namespace poplar {
 
struct VariableInterval;
namespace core {
class Tensor;
}
 
enum class UpsampleMethod {
  REPEAT 
};
 
class Tensor {
public:
  Tensor();
  Tensor(const Tensor &other);
  Tensor(Tensor &&other) noexcept;
  const Tensor &operator=(const Tensor &other) &;
  Tensor &operator=(Tensor &&other) &noexcept;
  ~Tensor();
 
  Type elementType() const;
 
  Tensor operator[](std::size_t i) const &;
  Tensor &&operator[](std::size_t i) &&;
 
  Tensor slice(std::size_t begin, std::size_t end, unsigned dimension) const &;
  Tensor &&slice(std::size_t begin, std::size_t end, unsigned dimension) &&;
 
  Tensor slice(std::size_t begin, std::size_t end) const {
    return slice(begin, end, 0);
  }
 
  Tensor slice(const Interval &region, unsigned dimension = 0) const {
    return slice(region.begin(), region.end(), dimension);
  }
 
  Tensor slice(ArrayRef<std::size_t> begin, ArrayRef<std::size_t> end) const;
 
  std::vector<Tensor> slices(ArrayRef<Interval> intervals,
                             unsigned dimension = 0) const;
 
  std::vector<Tensor>
  slices(const std::vector<std::vector<Interval>> &intervals,
         unsigned dimension = 0) const;
 
  std::vector<Tensor> slices(const poplar::ArrayRef<unsigned> &indices,
                             unsigned dimension = 0) const;
 
  Tensor index(ArrayRef<std::size_t> indices) const;
 
  Tensor flatten() const;
 
  Tensor flatten(unsigned dimBegin, unsigned dimEnd) const;
 
  Tensor reshape(ArrayRef<std::size_t> shape) const;
 
  Tensor dimShuffle(ArrayRef<unsigned> permutation) const;
 
  Tensor dimShufflePartial(ArrayRef<unsigned> source,
                           ArrayRef<unsigned> destination) const;
 
  Tensor dimRoll(unsigned dimIdx, unsigned newIdx = 0) const {
    return dimShufflePartial({dimIdx}, {newIdx});
  }
 
  Tensor reshapePartial(unsigned beginIndex, unsigned endIndex,
                        ArrayRef<std::size_t> newDims) const;
 
  Tensor expand(ArrayRef<std::size_t> indices) const;
 
  Tensor squeeze(ArrayRef<std::size_t> indices) const;
 
  Tensor transpose() const { return dimShuffle({1, 0}); }
 
  Tensor subSample(unsigned stride, unsigned dimension) const;
 
  Tensor upsample(unsigned scale, unsigned dimension,
                  UpsampleMethod method) const;
 
  Tensor broadcast(unsigned N, unsigned dimension) const;
 
  Tensor reinterpret(const Type &type) const;
 
  Tensor reverse(unsigned dimension) const;
 
  std::size_t numElements() const;
 
  std::size_t dim(unsigned i) const;
 
  std::vector<std::size_t> shape() const;
 
  unsigned rank() const;
 
  bool isContiguous() const;
 
  bool containsAliases() const;
 
  bool containsConstant() const;
 
  bool isParallelWriteable() const;
 
  const std::vector<Interval> getContiguousRegions() const;
 
  const std::vector<VariableInterval> getVarRegions() const;
 
  template <typename T> bool getConstantValue(T *val) const {
    return getConstantData(val, TypeTraits::make<T>());
  }
 
  bool intersectsWith(const Tensor &other) const;
 
  std::ostream &output(std::ostream &os) const;
 
  std::ostream &outputRegions(std::ostream &os) const;
 
  std::string shapeToString() const;
 
  void dump() const;
 
  void dumpRegions() const;
 
  std::string getVarStr() const;
  std::string getDebugStr() const;
 
  // Implementation
  Tensor(std::unique_ptr<core::Tensor>);
  core::Tensor &getImpl() const { return *impl; }
  std::unique_ptr<core::Tensor> *getPImpl() { return &impl; }
  bool valid() const { return impl.get() != nullptr; }
 
  Tensor getMetadata() const;
 
  bool hasMetadata() const;
 
private:
  std::unique_ptr<core::Tensor> impl;
  bool getConstantData(void *dst, const TypeTraits &traits) const;
};
 
bool operator==(const Tensor &a, const Tensor &b);
bool inline operator!=(const Tensor &a, const Tensor &b) { return !(a == b); }
 
Tensor concat(ArrayRef<Tensor> ts, unsigned dimension = 0);
 
inline Tensor concat(const Tensor &first, const Tensor &second,
                     unsigned dimension = 0) {
  return concat({first, second}, dimension);
}
 
Tensor append(const Tensor &first, const Tensor &second, unsigned dimension);
inline Tensor append(const Tensor &first, const Tensor &second) {
  return append(first, second, 0);
}
 
std::ostream &operator<<(std::ostream &os, const Tensor &tensor);
 
} // namespace poplar
 
#endif // poplar_Tensor_hpp