xref: /aosp_15_r20/external/eigen/unsupported/Eigen/CXX11/src/Tensor/TensorLayoutSwap.h (revision bf2c37156dfe67e5dfebd6d394bad8b2ab5804d4)
1 // This file is part of Eigen, a lightweight C++ template library
2 // for linear algebra.
3 //
4 // Copyright (C) 2014 Benoit Steiner <[email protected]>
5 //
6 // This Source Code Form is subject to the terms of the Mozilla
7 // Public License v. 2.0. If a copy of the MPL was not distributed
8 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
9 
10 #ifndef EIGEN_CXX11_TENSOR_TENSOR_LAYOUT_SWAP_H
11 #define EIGEN_CXX11_TENSOR_TENSOR_LAYOUT_SWAP_H
12 
13 namespace Eigen {
14 
15 /** \class TensorLayoutSwap
16   * \ingroup CXX11_Tensor_Module
17   *
18   * \brief Swap the layout from col-major to row-major, or row-major
19   * to col-major, and invert the order of the dimensions.
20   *
21   * Beware: the dimensions are reversed by this operation. If you want to
22   * preserve the ordering of the dimensions, you need to combine this
23   * operation with a shuffle.
24   *
25   * \example:
26   * Tensor<float, 2, ColMajor> input(2, 4);
27   * Tensor<float, 2, RowMajor> output = input.swap_layout();
28   * eigen_assert(output.dimension(0) == 4);
29   * eigen_assert(output.dimension(1) == 2);
30   *
31   * array<int, 2> shuffle(1, 0);
32   * output = input.swap_layout().shuffle(shuffle);
33   * eigen_assert(output.dimension(0) == 2);
34   * eigen_assert(output.dimension(1) == 4);
35   *
36   */
37 namespace internal {
38 template<typename XprType>
39 struct traits<TensorLayoutSwapOp<XprType> > : public traits<XprType>
40 {
41   typedef typename XprType::Scalar Scalar;
42   typedef traits<XprType> XprTraits;
43   typedef typename XprTraits::StorageKind StorageKind;
44   typedef typename XprTraits::Index Index;
45   typedef typename XprType::Nested Nested;
46   typedef typename remove_reference<Nested>::type _Nested;
47   static const int NumDimensions = traits<XprType>::NumDimensions;
48   static const int Layout = (traits<XprType>::Layout == ColMajor) ? RowMajor : ColMajor;
49   typedef typename XprTraits::PointerType PointerType;
50 };
51 
52 template<typename XprType>
53 struct eval<TensorLayoutSwapOp<XprType>, Eigen::Dense>
54 {
55   typedef const TensorLayoutSwapOp<XprType>& type;
56 };
57 
58 template<typename XprType>
59 struct nested<TensorLayoutSwapOp<XprType>, 1, typename eval<TensorLayoutSwapOp<XprType> >::type>
60 {
61   typedef TensorLayoutSwapOp<XprType> type;
62 };
63 
64 }  // end namespace internal
65 
66 
67 
68 template<typename XprType>
69 class TensorLayoutSwapOp : public TensorBase<TensorLayoutSwapOp<XprType>, WriteAccessors>
70 {
71   public:
72     typedef TensorBase<TensorLayoutSwapOp<XprType>, WriteAccessors> Base;
73     typedef typename Eigen::internal::traits<TensorLayoutSwapOp>::Scalar Scalar;
74     typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
75     typedef typename internal::remove_const<typename XprType::CoeffReturnType>::type CoeffReturnType;
76     typedef typename Eigen::internal::nested<TensorLayoutSwapOp>::type Nested;
77     typedef typename Eigen::internal::traits<TensorLayoutSwapOp>::StorageKind StorageKind;
78     typedef typename Eigen::internal::traits<TensorLayoutSwapOp>::Index Index;
79 
80     EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorLayoutSwapOp(const XprType& expr)
81         : m_xpr(expr) {}
82 
83     EIGEN_DEVICE_FUNC
84     const typename internal::remove_all<typename XprType::Nested>::type&
85     expression() const { return m_xpr; }
86 
87     EIGEN_TENSOR_INHERIT_ASSIGNMENT_OPERATORS(TensorLayoutSwapOp)
88   protected:
89     typename XprType::Nested m_xpr;
90 };
91 
92 
93 // Eval as rvalue
94 template<typename ArgType, typename Device>
95 struct TensorEvaluator<const TensorLayoutSwapOp<ArgType>, Device>
96 {
97   typedef TensorLayoutSwapOp<ArgType> XprType;
98   typedef typename XprType::Index Index;
99   static const int NumDims = internal::array_size<typename TensorEvaluator<ArgType, Device>::Dimensions>::value;
100   typedef DSizes<Index, NumDims> Dimensions;
101 
102   enum {
103     IsAligned = TensorEvaluator<ArgType, Device>::IsAligned,
104     PacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess,
105     BlockAccess = false,
106     PreferBlockAccess = TensorEvaluator<ArgType, Device>::PreferBlockAccess,
107     Layout = (static_cast<int>(TensorEvaluator<ArgType, Device>::Layout) == static_cast<int>(ColMajor)) ? RowMajor : ColMajor,
108     CoordAccess = false,  // to be implemented
109     RawAccess = TensorEvaluator<ArgType, Device>::RawAccess
110   };
111 
112   //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
113   typedef internal::TensorBlockNotImplemented TensorBlock;
114   //===--------------------------------------------------------------------===//
115 
116   EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
117       : m_impl(op.expression(), device)
118   {
119     for(int i = 0; i < NumDims; ++i) {
120       m_dimensions[i] = m_impl.dimensions()[NumDims-1-i];
121     }
122   }
123 
124 #ifdef EIGEN_USE_SYCL
125   // binding placeholder accessors to a command group handler for SYCL
126   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void bind(cl::sycl::handler &cgh) const {
127     m_impl.bind(cgh);
128   }
129 #endif
130 
131   typedef typename XprType::Scalar Scalar;
132   typedef typename XprType::CoeffReturnType CoeffReturnType;
133   typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
134   typedef StorageMemory<CoeffReturnType, Device> Storage;
135   typedef typename Storage::Type EvaluatorPointerType;
136 
137   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_dimensions; }
138 
139   EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(EvaluatorPointerType data) {
140     return m_impl.evalSubExprsIfNeeded(data);
141   }
142   EIGEN_STRONG_INLINE void cleanup() {
143     m_impl.cleanup();
144   }
145 
146   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const
147   {
148     return m_impl.coeff(index);
149   }
150 
151   template<int LoadMode>
152   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packet(Index index) const
153   {
154     return m_impl.template packet<LoadMode>(index);
155   }
156 
157   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool vectorized) const {
158     return m_impl.costPerCoeff(vectorized);
159   }
160 
161   EIGEN_DEVICE_FUNC typename Storage::Type data() const {
162     return constCast(m_impl.data());
163   }
164 
165   const TensorEvaluator<ArgType, Device>& impl() const { return m_impl; }
166 
167  protected:
168   TensorEvaluator<ArgType, Device> m_impl;
169   Dimensions m_dimensions;
170 };
171 
172 
173 // Eval as lvalue
174 template<typename ArgType, typename Device>
175   struct TensorEvaluator<TensorLayoutSwapOp<ArgType>, Device>
176   : public TensorEvaluator<const TensorLayoutSwapOp<ArgType>, Device>
177 {
178   typedef TensorEvaluator<const TensorLayoutSwapOp<ArgType>, Device> Base;
179   typedef TensorLayoutSwapOp<ArgType> XprType;
180 
181   enum {
182     IsAligned = TensorEvaluator<ArgType, Device>::IsAligned,
183     PacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess,
184     BlockAccess = false,
185     PreferBlockAccess = TensorEvaluator<ArgType, Device>::PreferBlockAccess,
186     Layout = (static_cast<int>(TensorEvaluator<ArgType, Device>::Layout) == static_cast<int>(ColMajor)) ? RowMajor : ColMajor,
187     CoordAccess = false  // to be implemented
188   };
189 
190   //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
191   typedef internal::TensorBlockNotImplemented TensorBlock;
192   //===--------------------------------------------------------------------===//
193 
194   EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
195     : Base(op, device)
196   { }
197 
198   typedef typename XprType::Index Index;
199   typedef typename XprType::Scalar Scalar;
200   typedef typename XprType::CoeffReturnType CoeffReturnType;
201   typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
202 
203   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType& coeffRef(Index index)
204   {
205     return this->m_impl.coeffRef(index);
206   }
207   template <int StoreMode> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
208   void writePacket(Index index, const PacketReturnType& x)
209   {
210     this->m_impl.template writePacket<StoreMode>(index, x);
211   }
212 };
213 
214 } // end namespace Eigen
215 
216 #endif // EIGEN_CXX11_TENSOR_TENSOR_LAYOUT_SWAP_H
217