xref: /aosp_15_r20/external/eigen/unsupported/Eigen/CXX11/src/Tensor/TensorEvalTo.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_EVAL_TO_H
11 #define EIGEN_CXX11_TENSOR_TENSOR_EVAL_TO_H
12 
13 namespace Eigen {
14 
15 /** \class TensorForcedEval
16   * \ingroup CXX11_Tensor_Module
17   *
18   * \brief Tensor reshaping class.
19   *
20   *
21   */
22 namespace internal {
23 template<typename XprType, template <class> class MakePointer_>
24 struct traits<TensorEvalToOp<XprType, MakePointer_> >
25 {
26   // Type promotion to handle the case where the types of the lhs and the rhs are different.
27   typedef typename XprType::Scalar Scalar;
28   typedef traits<XprType> XprTraits;
29   typedef typename XprTraits::StorageKind StorageKind;
30   typedef typename XprTraits::Index Index;
31   typedef typename XprType::Nested Nested;
32   typedef typename remove_reference<Nested>::type _Nested;
33   static const int NumDimensions = XprTraits::NumDimensions;
34   static const int Layout = XprTraits::Layout;
35   typedef typename MakePointer_<Scalar>::Type PointerType;
36 
37   enum {
38     Flags = 0
39   };
40   template <class T>
41   struct MakePointer {
42     // Intermediate typedef to workaround MSVC issue.
43     typedef MakePointer_<T> MakePointerT;
44     typedef typename MakePointerT::Type Type;
45 
46 
47   };
48 };
49 
50 template<typename XprType, template <class> class MakePointer_>
51 struct eval<TensorEvalToOp<XprType, MakePointer_>, Eigen::Dense>
52 {
53   typedef const TensorEvalToOp<XprType, MakePointer_>& type;
54 };
55 
56 template<typename XprType, template <class> class MakePointer_>
57 struct nested<TensorEvalToOp<XprType, MakePointer_>, 1, typename eval<TensorEvalToOp<XprType, MakePointer_> >::type>
58 {
59   typedef TensorEvalToOp<XprType, MakePointer_> type;
60 };
61 
62 }  // end namespace internal
63 
64 
65 
66 
67 template<typename XprType, template <class> class MakePointer_>
68 class TensorEvalToOp : public TensorBase<TensorEvalToOp<XprType, MakePointer_>, ReadOnlyAccessors>
69 {
70   public:
71   typedef typename Eigen::internal::traits<TensorEvalToOp>::Scalar Scalar;
72   typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
73   typedef typename internal::remove_const<typename XprType::CoeffReturnType>::type CoeffReturnType;
74   typedef typename MakePointer_<CoeffReturnType>::Type PointerType;
75   typedef typename Eigen::internal::nested<TensorEvalToOp>::type Nested;
76   typedef typename Eigen::internal::traits<TensorEvalToOp>::StorageKind StorageKind;
77   typedef typename Eigen::internal::traits<TensorEvalToOp>::Index Index;
78 
79   static const int NumDims = Eigen::internal::traits<TensorEvalToOp>::NumDimensions;
80 
81   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorEvalToOp(PointerType buffer, const XprType& expr)
82       : m_xpr(expr), m_buffer(buffer) {}
83 
84     EIGEN_DEVICE_FUNC
85     const typename internal::remove_all<typename XprType::Nested>::type&
86     expression() const { return m_xpr; }
87 
88     EIGEN_DEVICE_FUNC PointerType buffer() const { return m_buffer; }
89 
90   protected:
91     typename XprType::Nested m_xpr;
92     PointerType m_buffer;
93 };
94 
95 
96 
97 template<typename ArgType, typename Device, template <class> class MakePointer_>
98 struct TensorEvaluator<const TensorEvalToOp<ArgType, MakePointer_>, Device>
99 {
100   typedef TensorEvalToOp<ArgType, MakePointer_> XprType;
101   typedef typename ArgType::Scalar Scalar;
102   typedef typename TensorEvaluator<ArgType, Device>::Dimensions Dimensions;
103   typedef typename XprType::Index Index;
104   typedef typename internal::remove_const<typename XprType::CoeffReturnType>::type CoeffReturnType;
105   typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
106   static const int PacketSize = PacketType<CoeffReturnType, Device>::size;
107   typedef typename Eigen::internal::traits<XprType>::PointerType TensorPointerType;
108   typedef StorageMemory<CoeffReturnType, Device> Storage;
109   typedef typename Storage::Type EvaluatorPointerType;
110   enum {
111     IsAligned         = TensorEvaluator<ArgType, Device>::IsAligned,
112     PacketAccess      = TensorEvaluator<ArgType, Device>::PacketAccess,
113     BlockAccess       = true,
114     PreferBlockAccess = false,
115     Layout            = TensorEvaluator<ArgType, Device>::Layout,
116     CoordAccess       = false,  // to be implemented
117     RawAccess         = true
118   };
119 
120   static const int NumDims = internal::traits<ArgType>::NumDimensions;
121 
122   //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
123   typedef internal::TensorBlockDescriptor<NumDims, Index> TensorBlockDesc;
124   typedef internal::TensorBlockScratchAllocator<Device> TensorBlockScratch;
125 
126   typedef typename TensorEvaluator<const ArgType, Device>::TensorBlock
127       ArgTensorBlock;
128 
129   typedef internal::TensorBlockAssignment<
130       CoeffReturnType, NumDims, typename ArgTensorBlock::XprType, Index>
131       TensorBlockAssignment;
132   //===--------------------------------------------------------------------===//
133 
134   EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
135       : m_impl(op.expression(), device), m_buffer(device.get(op.buffer())), m_expression(op.expression()){}
136 
137 
138   EIGEN_STRONG_INLINE ~TensorEvaluator() {
139   }
140 
141 
142   EIGEN_DEVICE_FUNC const Dimensions& dimensions() const { return m_impl.dimensions(); }
143 
144   EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(EvaluatorPointerType scalar) {
145     EIGEN_UNUSED_VARIABLE(scalar);
146     eigen_assert(scalar == NULL);
147     return m_impl.evalSubExprsIfNeeded(m_buffer);
148   }
149 
150 #ifdef EIGEN_USE_THREADS
151   template <typename EvalSubExprsCallback>
152   EIGEN_STRONG_INLINE void evalSubExprsIfNeededAsync(
153       EvaluatorPointerType scalar, EvalSubExprsCallback done) {
154     EIGEN_UNUSED_VARIABLE(scalar);
155     eigen_assert(scalar == NULL);
156     m_impl.evalSubExprsIfNeededAsync(m_buffer, std::move(done));
157   }
158 #endif
159 
160   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalScalar(Index i) {
161     m_buffer[i] = m_impl.coeff(i);
162   }
163   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalPacket(Index i) {
164     internal::pstoret<CoeffReturnType, PacketReturnType, Aligned>(m_buffer + i, m_impl.template packet<TensorEvaluator<ArgType, Device>::IsAligned ? Aligned : Unaligned>(i));
165   }
166 
167   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
168   internal::TensorBlockResourceRequirements getResourceRequirements() const {
169     return m_impl.getResourceRequirements();
170   }
171 
172   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalBlock(
173       TensorBlockDesc& desc, TensorBlockScratch& scratch) {
174     // Add `m_buffer` as destination buffer to the block descriptor.
175     desc.template AddDestinationBuffer<Layout>(
176         /*dst_base=*/m_buffer + desc.offset(),
177         /*dst_strides=*/internal::strides<Layout>(m_impl.dimensions()));
178 
179     ArgTensorBlock block =
180         m_impl.block(desc, scratch, /*root_of_expr_ast=*/true);
181 
182     // If block was evaluated into a destination buffer, there is no need to do
183     // an assignment.
184     if (block.kind() != internal::TensorBlockKind::kMaterializedInOutput) {
185       TensorBlockAssignment::Run(
186           TensorBlockAssignment::target(
187               desc.dimensions(), internal::strides<Layout>(m_impl.dimensions()),
188               m_buffer, desc.offset()),
189           block.expr());
190     }
191     block.cleanup();
192   }
193 
194   EIGEN_STRONG_INLINE void cleanup() {
195     m_impl.cleanup();
196   }
197 
198   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const
199   {
200     return m_buffer[index];
201   }
202 
203   template<int LoadMode>
204   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packet(Index index) const
205   {
206     return internal::ploadt<PacketReturnType, LoadMode>(m_buffer + index);
207   }
208 
209   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool vectorized) const {
210     // We assume that evalPacket or evalScalar is called to perform the
211     // assignment and account for the cost of the write here.
212     return m_impl.costPerCoeff(vectorized) +
213         TensorOpCost(0, sizeof(CoeffReturnType), 0, vectorized, PacketSize);
214   }
215 
216   EIGEN_DEVICE_FUNC EvaluatorPointerType data() const { return m_buffer; }
217   ArgType expression() const { return m_expression; }
218   #ifdef EIGEN_USE_SYCL
219   // binding placeholder accessors to a command group handler for SYCL
220   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void bind(cl::sycl::handler &cgh) const {
221     m_impl.bind(cgh);
222     m_buffer.bind(cgh);
223   }
224   #endif
225 
226 
227  private:
228   TensorEvaluator<ArgType, Device> m_impl;
229   EvaluatorPointerType m_buffer;
230   const ArgType m_expression;
231 };
232 
233 
234 } // end namespace Eigen
235 
236 #endif // EIGEN_CXX11_TENSOR_TENSOR_EVAL_TO_H
237