TensorMap.h

// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0
 
#ifndef EIGEN_TENSOR_TENSOR_MAP_H
#define EIGEN_TENSOR_TENSOR_MAP_H
 
// IWYU pragma: private
#include "./InternalHeaderCheck.h"
 
namespace Eigen {
 
// FIXME: Use proper doxygen documentation (e.g. \tparam MakePointer_).

template <typename PlainObjectType, int Options_, template <class> class MakePointer_>
class TensorMap : public TensorBase<TensorMap<PlainObjectType, Options_, MakePointer_> > {
 public:
  typedef TensorMap<PlainObjectType, Options_, MakePointer_> Self;
  typedef TensorBase<TensorMap<PlainObjectType, Options_, MakePointer_> > Base;
#ifdef EIGEN_USE_SYCL
  typedef std::remove_reference_t<typename Eigen::internal::ref_selector<Self>::type> Nested;
#else
  typedef typename Eigen::internal::ref_selector<Self>::type Nested;
#endif
  typedef typename internal::traits<PlainObjectType>::StorageKind StorageKind;
  typedef typename internal::traits<PlainObjectType>::Index Index;
  typedef typename internal::traits<PlainObjectType>::Scalar Scalar;
  typedef typename NumTraits<Scalar>::Real RealScalar;
  typedef typename PlainObjectType::Base::CoeffReturnType CoeffReturnType;
 
  typedef typename MakePointer_<Scalar>::Type PointerType;
  typedef typename MakePointer_<Scalar>::ConstType PointerConstType;
 
  // WARN: PointerType still can be a pointer to const (const Scalar*), for
  // example in TensorMap<Tensor<const Scalar, ...>> expression. This type of
  // expression should be illegal, but adding this restriction is not possible
  // in practice (see https://bitbucket.org/eigen/eigen/pull-requests/488).
  typedef std::conditional_t<bool(internal::is_lvalue<PlainObjectType>::value),
                             PointerType,      // use simple pointer in lvalue expressions
                             PointerConstType  // use const pointer in rvalue expressions
                             >
      StoragePointerType;
 
  // If TensorMap was constructed over rvalue expression (e.g. const Tensor),
  // we should return a reference to const from operator() (and others), even
  // if TensorMap itself is not const.
  typedef std::conditional_t<bool(internal::is_lvalue<PlainObjectType>::value), Scalar&, const Scalar&> StorageRefType;
 
  static constexpr int Options = Options_;
 
  static constexpr Index NumIndices = PlainObjectType::NumIndices;
  typedef typename PlainObjectType::Dimensions Dimensions;
 
  static constexpr int Layout = PlainObjectType::Layout;
  enum { IsAligned = ((int(Options_) & Aligned) == Aligned), CoordAccess = true, RawAccess = true };
 
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorMap(StoragePointerType dataPtr) : m_data(dataPtr), m_dimensions() {
    // The number of dimensions used to construct a tensor must be equal to the rank of the tensor.
    EIGEN_STATIC_ASSERT((0 == NumIndices || NumIndices == Dynamic), YOU_MADE_A_PROGRAMMING_MISTAKE)
  }
 
  template <typename... IndexTypes>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorMap(StoragePointerType dataPtr, Index firstDimension,
                                                  IndexTypes... otherDimensions)
      : m_data(dataPtr), m_dimensions(firstDimension, otherDimensions...) {
    // The number of dimensions used to construct a tensor must be equal to the rank of the tensor.
    EIGEN_STATIC_ASSERT((sizeof...(otherDimensions) + 1 == NumIndices || NumIndices == Dynamic),
                        YOU_MADE_A_PROGRAMMING_MISTAKE)
  }
 
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorMap(StoragePointerType dataPtr,
                                                  const array<Index, NumIndices>& dimensions)
      : m_data(dataPtr), m_dimensions(dimensions) {}
 
  template <typename Dimensions>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorMap(StoragePointerType dataPtr, const Dimensions& dimensions)
      : m_data(dataPtr), m_dimensions(dimensions) {}
 
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorMap(PlainObjectType& tensor)
      : m_data(tensor.data()), m_dimensions(tensor.dimensions()) {}
 
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index rank() const { return m_dimensions.rank(); }
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index dimension(Index n) const { return m_dimensions[n]; }
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_dimensions; }
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index size() const { return m_dimensions.TotalSize(); }
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE StoragePointerType data() { return m_data; }
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE StoragePointerType data() const { return m_data; }
 
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE StorageRefType operator()(const array<Index, NumIndices>& indices) const {
    EIGEN_IF_CONSTEXPR (PlainObjectType::Options & RowMajor) {
      const Index index = m_dimensions.IndexOfRowMajor(indices);
      return m_data[index];
    } else {
      const Index index = m_dimensions.IndexOfColMajor(indices);
      return m_data[index];
    }
  }
 
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE StorageRefType operator()() const {
    EIGEN_STATIC_ASSERT(NumIndices == 0, YOU_MADE_A_PROGRAMMING_MISTAKE)
    return m_data[0];
  }
 
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE StorageRefType operator()(Index index) const {
    eigen_internal_assert(index >= 0 && index < size());
    return m_data[index];
  }
 
  template <typename... IndexTypes>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE StorageRefType operator()(Index firstIndex, Index secondIndex,
                                                                  IndexTypes... otherIndices) const {
    EIGEN_STATIC_ASSERT(sizeof...(otherIndices) + 2 == NumIndices, YOU_MADE_A_PROGRAMMING_MISTAKE)
    eigen_assert(internal::indices_fit<Index>(otherIndices...));
    EIGEN_IF_CONSTEXPR (PlainObjectType::Options & RowMajor) {
      const Index index = m_dimensions.IndexOfRowMajor(
          array<Index, NumIndices>{{firstIndex, secondIndex, static_cast<Index>(otherIndices)...}});
      return m_data[index];
    } else {
      const Index index = m_dimensions.IndexOfColMajor(
          array<Index, NumIndices>{{firstIndex, secondIndex, static_cast<Index>(otherIndices)...}});
      return m_data[index];
    }
  }
 
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE StorageRefType operator()(const array<Index, NumIndices>& indices) {
    EIGEN_IF_CONSTEXPR (PlainObjectType::Options & RowMajor) {
      const Index index = m_dimensions.IndexOfRowMajor(indices);
      return m_data[index];
    } else {
      const Index index = m_dimensions.IndexOfColMajor(indices);
      return m_data[index];
    }
  }
 
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE StorageRefType operator()() {
    EIGEN_STATIC_ASSERT(NumIndices == 0, YOU_MADE_A_PROGRAMMING_MISTAKE)
    return m_data[0];
  }
 
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE StorageRefType operator()(Index index) {
    eigen_internal_assert(index >= 0 && index < size());
    return m_data[index];
  }
 
  template <typename... IndexTypes>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE StorageRefType operator()(Index firstIndex, Index secondIndex,
                                                                  IndexTypes... otherIndices) {
    static_assert(sizeof...(otherIndices) + 2 == NumIndices || NumIndices == Dynamic,
                  "Number of indices used to access a tensor coefficient must be equal to the rank of the tensor.");
    eigen_assert(internal::indices_fit<Index>(otherIndices...));
    const std::size_t NumDims = sizeof...(otherIndices) + 2;
    EIGEN_IF_CONSTEXPR (PlainObjectType::Options & RowMajor) {
      const Index index = m_dimensions.IndexOfRowMajor(
          array<Index, NumDims>{{firstIndex, secondIndex, static_cast<Index>(otherIndices)...}});
      return m_data[index];
    } else {
      const Index index = m_dimensions.IndexOfColMajor(
          array<Index, NumDims>{{firstIndex, secondIndex, static_cast<Index>(otherIndices)...}});
      return m_data[index];
    }
  }
 
  EIGEN_INHERIT_ASSIGNMENT_OPERATORS(TensorMap)
 
 private:
  StoragePointerType m_data;
  Dimensions m_dimensions;
};
 
}  // end namespace Eigen
 
#endif  // EIGEN_TENSOR_TENSOR_MAP_H