Eigen-unsupported  3.4.90 (git rev 67eeba6e720c5745abc77ae6c92ce0a44aa7b7ae)
TensorCustomOp.h
1 // This file is part of Eigen, a lightweight C++ template library
2 // for linear algebra.
3 //
4 // Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
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_CUSTOM_OP_H
11 #define EIGEN_CXX11_TENSOR_TENSOR_CUSTOM_OP_H
12 
13 #include "./InternalHeaderCheck.h"
14 
15 namespace Eigen {
16 
24 namespace internal {
25 template<typename CustomUnaryFunc, typename XprType>
26 struct traits<TensorCustomUnaryOp<CustomUnaryFunc, XprType> >
27 {
28  typedef typename XprType::Scalar Scalar;
29  typedef typename XprType::StorageKind StorageKind;
30  typedef typename XprType::Index Index;
31  typedef typename XprType::Nested Nested;
32  typedef std::remove_reference_t<Nested> Nested_;
33  static constexpr int NumDimensions = traits<XprType>::NumDimensions;
34  static constexpr int Layout = traits<XprType>::Layout;
35  typedef typename traits<XprType>::PointerType PointerType;
36 };
37 
38 template<typename CustomUnaryFunc, typename XprType>
39 struct eval<TensorCustomUnaryOp<CustomUnaryFunc, XprType>, Eigen::Dense>
40 {
41  typedef const TensorCustomUnaryOp<CustomUnaryFunc, XprType>EIGEN_DEVICE_REF type;
42 };
43 
44 template<typename CustomUnaryFunc, typename XprType>
45 struct nested<TensorCustomUnaryOp<CustomUnaryFunc, XprType> >
46 {
47  typedef TensorCustomUnaryOp<CustomUnaryFunc, XprType> type;
48 };
49 
50 } // end namespace internal
51 
52 
53 
54 template<typename CustomUnaryFunc, typename XprType>
55 class TensorCustomUnaryOp : public TensorBase<TensorCustomUnaryOp<CustomUnaryFunc, XprType>, ReadOnlyAccessors>
56 {
57  public:
58  typedef typename internal::traits<TensorCustomUnaryOp>::Scalar Scalar;
59  typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
60  typedef typename XprType::CoeffReturnType CoeffReturnType;
61  typedef typename internal::nested<TensorCustomUnaryOp>::type Nested;
62  typedef typename internal::traits<TensorCustomUnaryOp>::StorageKind StorageKind;
63  typedef typename internal::traits<TensorCustomUnaryOp>::Index Index;
64 
65  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorCustomUnaryOp(const XprType& expr, const CustomUnaryFunc& func)
66  : m_expr(expr), m_func(func) {}
67 
68  EIGEN_DEVICE_FUNC
69  const CustomUnaryFunc& func() const { return m_func; }
70 
71  EIGEN_DEVICE_FUNC
72  const internal::remove_all_t<typename XprType::Nested>&
73  expression() const { return m_expr; }
74 
75  protected:
76  typename XprType::Nested m_expr;
77  const CustomUnaryFunc m_func;
78 };
79 
80 
81 // Eval as rvalue
82 template<typename CustomUnaryFunc, typename XprType, typename Device>
83 struct TensorEvaluator<const TensorCustomUnaryOp<CustomUnaryFunc, XprType>, Device>
84 {
85  typedef TensorCustomUnaryOp<CustomUnaryFunc, XprType> ArgType;
86  typedef typename internal::traits<ArgType>::Index Index;
87  static constexpr int NumDims = internal::traits<ArgType>::NumDimensions;
88  typedef DSizes<Index, NumDims> Dimensions;
89  typedef std::remove_const_t<typename ArgType::Scalar> Scalar;
90  typedef std::remove_const_t<typename XprType::CoeffReturnType> CoeffReturnType;
91  typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
92  static constexpr int PacketSize = PacketType<CoeffReturnType, Device>::size;
93  typedef typename Eigen::internal::traits<XprType>::PointerType TensorPointerType;
94  typedef StorageMemory<CoeffReturnType, Device> Storage;
95  typedef typename Storage::Type EvaluatorPointerType;
96 
97  static constexpr int Layout = TensorEvaluator<XprType, Device>::Layout;
98  enum {
99  IsAligned = false,
100  PacketAccess = (PacketType<CoeffReturnType, Device>::size > 1),
101  BlockAccess = false,
102  PreferBlockAccess = false,
103  CoordAccess = false, // to be implemented
104  RawAccess = false
105  };
106 
107  //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
108  typedef internal::TensorBlockNotImplemented TensorBlock;
109  //===--------------------------------------------------------------------===//
110 
111  EIGEN_STRONG_INLINE TensorEvaluator(const ArgType& op, const Device& device)
112  : m_op(op), m_device(device), m_result(NULL)
113  {
114  m_dimensions = op.func().dimensions(op.expression());
115  }
116 
117  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_dimensions; }
118 
119  EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(EvaluatorPointerType data) {
120  if (data) {
121  evalTo(data);
122  return false;
123  } else {
124  m_result = static_cast<EvaluatorPointerType>(m_device.get( (CoeffReturnType*)
125  m_device.allocate_temp(dimensions().TotalSize() * sizeof(Scalar))));
126  evalTo(m_result);
127  return true;
128  }
129  }
130 
131  EIGEN_STRONG_INLINE void cleanup() {
132  if (m_result) {
133  m_device.deallocate_temp(m_result);
134  m_result = NULL;
135  }
136  }
137 
138  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const {
139  return m_result[index];
140  }
141 
142  template<int LoadMode>
143  EIGEN_DEVICE_FUNC PacketReturnType packet(Index index) const {
144  return internal::ploadt<PacketReturnType, LoadMode>(m_result + index);
145  }
146 
147  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool vectorized) const {
148  // TODO(rmlarsen): Extend CustomOp API to return its cost estimate.
149  return TensorOpCost(sizeof(CoeffReturnType), 0, 0, vectorized, PacketSize);
150  }
151 
152  EIGEN_DEVICE_FUNC EvaluatorPointerType data() const { return m_result; }
153 
154 #ifdef EIGEN_USE_SYCL
155  // binding placeholder accessors to a command group handler for SYCL
156  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void bind(cl::sycl::handler &cgh) const {
157  m_result.bind(cgh);
158  }
159 #endif
160 
161  protected:
162  void evalTo(EvaluatorPointerType data) {
163  TensorMap<Tensor<CoeffReturnType, NumDims, Layout, Index> > result(m_device.get(data), m_dimensions);
164  m_op.func().eval(m_op.expression(), result, m_device);
165  }
166 
167  Dimensions m_dimensions;
168  const ArgType m_op;
169  const Device EIGEN_DEVICE_REF m_device;
170  EvaluatorPointerType m_result;
171 };
172 
173 
174 
182 namespace internal {
183 template<typename CustomBinaryFunc, typename LhsXprType, typename RhsXprType>
184 struct traits<TensorCustomBinaryOp<CustomBinaryFunc, LhsXprType, RhsXprType> >
185 {
186  typedef typename internal::promote_storage_type<typename LhsXprType::Scalar,
187  typename RhsXprType::Scalar>::ret Scalar;
188  typedef typename internal::promote_storage_type<typename LhsXprType::CoeffReturnType,
189  typename RhsXprType::CoeffReturnType>::ret CoeffReturnType;
190  typedef typename promote_storage_type<typename traits<LhsXprType>::StorageKind,
191  typename traits<RhsXprType>::StorageKind>::ret StorageKind;
192  typedef typename promote_index_type<typename traits<LhsXprType>::Index,
193  typename traits<RhsXprType>::Index>::type Index;
194  typedef typename LhsXprType::Nested LhsNested;
195  typedef typename RhsXprType::Nested RhsNested;
196  typedef std::remove_reference_t<LhsNested> LhsNested_;
197  typedef std::remove_reference_t<RhsNested> RhsNested_;
198  static constexpr int NumDimensions = traits<LhsXprType>::NumDimensions;
199  static constexpr int Layout = traits<LhsXprType>::Layout;
200  typedef std::conditional_t<Pointer_type_promotion<typename LhsXprType::Scalar, Scalar>::val,
201  typename traits<LhsXprType>::PointerType, typename traits<RhsXprType>::PointerType> PointerType;
202 };
203 
204 template<typename CustomBinaryFunc, typename LhsXprType, typename RhsXprType>
205 struct eval<TensorCustomBinaryOp<CustomBinaryFunc, LhsXprType, RhsXprType>, Eigen::Dense>
206 {
207  typedef const TensorCustomBinaryOp<CustomBinaryFunc, LhsXprType, RhsXprType>& type;
208 };
209 
210 template<typename CustomBinaryFunc, typename LhsXprType, typename RhsXprType>
211 struct nested<TensorCustomBinaryOp<CustomBinaryFunc, LhsXprType, RhsXprType> >
212 {
213  typedef TensorCustomBinaryOp<CustomBinaryFunc, LhsXprType, RhsXprType> type;
214 };
215 
216 } // end namespace internal
217 
218 
219 
220 template<typename CustomBinaryFunc, typename LhsXprType, typename RhsXprType>
221 class TensorCustomBinaryOp : public TensorBase<TensorCustomBinaryOp<CustomBinaryFunc, LhsXprType, RhsXprType>, ReadOnlyAccessors>
222 {
223  public:
224  typedef typename internal::traits<TensorCustomBinaryOp>::Scalar Scalar;
225  typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
226  typedef typename internal::traits<TensorCustomBinaryOp>::CoeffReturnType CoeffReturnType;
227  typedef typename internal::nested<TensorCustomBinaryOp>::type Nested;
228  typedef typename internal::traits<TensorCustomBinaryOp>::StorageKind StorageKind;
229  typedef typename internal::traits<TensorCustomBinaryOp>::Index Index;
230 
231  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorCustomBinaryOp(const LhsXprType& lhs, const RhsXprType& rhs, const CustomBinaryFunc& func)
232 
233  : m_lhs_xpr(lhs), m_rhs_xpr(rhs), m_func(func) {}
234 
235  EIGEN_DEVICE_FUNC
236  const CustomBinaryFunc& func() const { return m_func; }
237 
238  EIGEN_DEVICE_FUNC
239  const internal::remove_all_t<typename LhsXprType::Nested>&
240  lhsExpression() const { return m_lhs_xpr; }
241 
242  EIGEN_DEVICE_FUNC
243  const internal::remove_all_t<typename RhsXprType::Nested>&
244  rhsExpression() const { return m_rhs_xpr; }
245 
246  protected:
247  typename LhsXprType::Nested m_lhs_xpr;
248  typename RhsXprType::Nested m_rhs_xpr;
249  const CustomBinaryFunc m_func;
250 };
251 
252 
253 // Eval as rvalue
254 template<typename CustomBinaryFunc, typename LhsXprType, typename RhsXprType, typename Device>
255 struct TensorEvaluator<const TensorCustomBinaryOp<CustomBinaryFunc, LhsXprType, RhsXprType>, Device>
256 {
257  typedef TensorCustomBinaryOp<CustomBinaryFunc, LhsXprType, RhsXprType> XprType;
258  typedef typename internal::traits<XprType>::Index Index;
259  static constexpr int NumDims = internal::traits<XprType>::NumDimensions;
260  typedef DSizes<Index, NumDims> Dimensions;
261  typedef typename XprType::Scalar Scalar;
262  typedef std::remove_const_t<typename XprType::CoeffReturnType> CoeffReturnType;
263  typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
264  static constexpr int PacketSize = PacketType<CoeffReturnType, Device>::size;
265 
266  typedef typename Eigen::internal::traits<XprType>::PointerType TensorPointerType;
267  typedef StorageMemory<CoeffReturnType, Device> Storage;
268  typedef typename Storage::Type EvaluatorPointerType;
269 
270  static constexpr int Layout = TensorEvaluator<LhsXprType, Device>::Layout;
271  enum {
272  IsAligned = false,
273  PacketAccess = (PacketType<CoeffReturnType, Device>::size > 1),
274  BlockAccess = false,
275  PreferBlockAccess = false,
276  CoordAccess = false, // to be implemented
277  RawAccess = false
278  };
279 
280  //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
281  typedef internal::TensorBlockNotImplemented TensorBlock;
282  //===--------------------------------------------------------------------===//
283 
284  EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
285  : m_op(op), m_device(device), m_result(NULL)
286  {
287  m_dimensions = op.func().dimensions(op.lhsExpression(), op.rhsExpression());
288  }
289 
290  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_dimensions; }
291 
292  EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(EvaluatorPointerType data) {
293  if (data) {
294  evalTo(data);
295  return false;
296  } else {
297  m_result = static_cast<EvaluatorPointerType>(m_device.get( (CoeffReturnType*)
298  m_device.allocate_temp(dimensions().TotalSize() * sizeof(CoeffReturnType))));
299  evalTo(m_result);
300  return true;
301  }
302  }
303 
304  EIGEN_STRONG_INLINE void cleanup() {
305  if (m_result != NULL) {
306  m_device.deallocate_temp(m_result);
307  m_result = NULL;
308  }
309  }
310 
311  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const {
312  return m_result[index];
313  }
314 
315  template<int LoadMode>
316  EIGEN_DEVICE_FUNC PacketReturnType packet(Index index) const {
317  return internal::ploadt<PacketReturnType, LoadMode>(m_result + index);
318  }
319 
320  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool vectorized) const {
321  // TODO(rmlarsen): Extend CustomOp API to return its cost estimate.
322  return TensorOpCost(sizeof(CoeffReturnType), 0, 0, vectorized, PacketSize);
323  }
324 
325  EIGEN_DEVICE_FUNC EvaluatorPointerType data() const { return m_result; }
326 
327 #ifdef EIGEN_USE_SYCL
328  // binding placeholder accessors to a command group handler for SYCL
329  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void bind(cl::sycl::handler &cgh) const {
330  m_result.bind(cgh);
331  }
332 #endif
333 
334  protected:
335  void evalTo(EvaluatorPointerType data) {
336  TensorMap<Tensor<CoeffReturnType, NumDims, Layout> > result(m_device.get(data), m_dimensions);
337  m_op.func().eval(m_op.lhsExpression(), m_op.rhsExpression(), result, m_device);
338  }
339 
340  Dimensions m_dimensions;
341  const XprType m_op;
342  const Device EIGEN_DEVICE_REF m_device;
343  EvaluatorPointerType m_result;
344 };
345 
346 
347 } // end namespace Eigen
348 
349 #endif // EIGEN_CXX11_TENSOR_TENSOR_CUSTOM_OP_H
Eigen::TensorBase
The tensor base class.
Definition: TensorForwardDeclarations.h:58
Eigen::TensorCustomBinaryOp
Tensor custom class.
Definition: TensorCustomOp.h:222
Eigen::TensorCustomUnaryOp
Tensor custom class.
Definition: TensorCustomOp.h:56
Eigen
Namespace containing all symbols from the Eigen library.
Eigen::Index
EIGEN_DEFAULT_DENSE_INDEX_TYPE Index
Eigen::TensorEvaluator
A cost model used to limit the number of threads used for evaluating tensor expression.
Definition: TensorEvaluator.h:31