Please, help us to better know about our user community by answering the following short survey: https://forms.gle/wpyrxWi18ox9Z5ae9
TensorReverse.h
1 // This file is part of Eigen, a lightweight C++ template library
2 // for linear algebra.
3 //
4 // Copyright (C) 2014 Navdeep Jaitly <ndjaitly@google.com>
5 // Benoit Steiner <benoit.steiner.goog@gmail.com>
6 //
7 // This Source Code Form is subject to the terms of the Mozilla
8 // Public License v. 2.0. If a copy of the MPL was not distributed
9 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
10 
11 #ifndef EIGEN_CXX11_TENSOR_TENSOR_REVERSE_H
12 #define EIGEN_CXX11_TENSOR_TENSOR_REVERSE_H
13 namespace Eigen {
14 
21 namespace internal {
22 template<typename ReverseDimensions, typename XprType>
23 struct traits<TensorReverseOp<ReverseDimensions,
24  XprType> > : public traits<XprType>
25 {
26  typedef typename XprType::Scalar Scalar;
27  typedef traits<XprType> XprTraits;
28  typedef typename XprTraits::StorageKind StorageKind;
29  typedef typename XprTraits::Index Index;
30  typedef typename XprType::Nested Nested;
31  typedef typename remove_reference<Nested>::type _Nested;
32  static const int NumDimensions = XprTraits::NumDimensions;
33  static const int Layout = XprTraits::Layout;
34  typedef typename XprTraits::PointerType PointerType;
35 };
36 
37 template<typename ReverseDimensions, typename XprType>
38 struct eval<TensorReverseOp<ReverseDimensions, XprType>, Eigen::Dense>
39 {
40  typedef const TensorReverseOp<ReverseDimensions, XprType>& type;
41 };
42 
43 template<typename ReverseDimensions, typename XprType>
44 struct nested<TensorReverseOp<ReverseDimensions, XprType>, 1,
45  typename eval<TensorReverseOp<ReverseDimensions, XprType> >::type>
46 {
47  typedef TensorReverseOp<ReverseDimensions, XprType> type;
48 };
49 
50 } // end namespace internal
51 
52 template<typename ReverseDimensions, typename XprType>
53 class TensorReverseOp : public TensorBase<TensorReverseOp<ReverseDimensions,
54  XprType>, WriteAccessors>
55 {
56  public:
57  typedef TensorBase<TensorReverseOp<ReverseDimensions, XprType>, WriteAccessors>Base;
58  typedef typename Eigen::internal::traits<TensorReverseOp>::Scalar Scalar;
59  typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
60  typedef typename XprType::CoeffReturnType CoeffReturnType;
61  typedef typename Eigen::internal::nested<TensorReverseOp>::type Nested;
62  typedef typename Eigen::internal::traits<TensorReverseOp>::StorageKind
63  StorageKind;
64  typedef typename Eigen::internal::traits<TensorReverseOp>::Index Index;
65 
66  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorReverseOp(
67  const XprType& expr, const ReverseDimensions& reverse_dims)
68  : m_xpr(expr), m_reverse_dims(reverse_dims) { }
69 
70  EIGEN_DEVICE_FUNC
71  const ReverseDimensions& reverse() const { return m_reverse_dims; }
72 
73  EIGEN_DEVICE_FUNC
74  const typename internal::remove_all<typename XprType::Nested>::type&
75  expression() const { return m_xpr; }
76 
77  EIGEN_TENSOR_INHERIT_ASSIGNMENT_OPERATORS(TensorReverseOp)
78 
79 
80  protected:
81  typename XprType::Nested m_xpr;
82  const ReverseDimensions m_reverse_dims;
83 };
84 
85 // Eval as rvalue
86 template<typename ReverseDimensions, typename ArgType, typename Device>
87 struct TensorEvaluator<const TensorReverseOp<ReverseDimensions, ArgType>, Device>
88 {
89  typedef TensorReverseOp<ReverseDimensions, ArgType> XprType;
90  typedef typename XprType::Index Index;
91  static const int NumDims = internal::array_size<ReverseDimensions>::value;
92  typedef DSizes<Index, NumDims> Dimensions;
93  typedef typename XprType::Scalar Scalar;
94  typedef typename XprType::CoeffReturnType CoeffReturnType;
95  typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
96  static const int PacketSize = PacketType<CoeffReturnType, Device>::size;
97  typedef StorageMemory<CoeffReturnType, Device> Storage;
98  typedef typename Storage::Type EvaluatorPointerType;
99 
100  enum {
101  IsAligned = false,
102  PacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess,
103  BlockAccess = NumDims > 0,
104  PreferBlockAccess = true,
105  Layout = TensorEvaluator<ArgType, Device>::Layout,
106  CoordAccess = false, // to be implemented
107  RawAccess = false
108  };
109 
110  typedef internal::TensorIntDivisor<Index> IndexDivisor;
111 
112  //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
113  typedef internal::TensorBlockDescriptor<NumDims, Index> TensorBlockDesc;
114  typedef internal::TensorBlockScratchAllocator<Device> TensorBlockScratch;
115 
116  typedef typename TensorEvaluator<const ArgType, Device>::TensorBlock
117  ArgTensorBlock;
118 
119  typedef typename internal::TensorMaterializedBlock<CoeffReturnType, NumDims,
120  Layout, Index>
121  TensorBlock;
122  //===--------------------------------------------------------------------===//
123 
124  EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
125  : m_impl(op.expression(), device),
126  m_reverse(op.reverse()),
127  m_device(device)
128  {
129  // Reversing a scalar isn't supported yet. It would be a no-op anyway.
130  EIGEN_STATIC_ASSERT((NumDims > 0), YOU_MADE_A_PROGRAMMING_MISTAKE);
131 
132  // Compute strides
133  m_dimensions = m_impl.dimensions();
134  if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
135  m_strides[0] = 1;
136  for (int i = 1; i < NumDims; ++i) {
137  m_strides[i] = m_strides[i-1] * m_dimensions[i-1];
138  if (m_strides[i] > 0) m_fastStrides[i] = IndexDivisor(m_strides[i]);
139  }
140  } else {
141  m_strides[NumDims-1] = 1;
142  for (int i = NumDims - 2; i >= 0; --i) {
143  m_strides[i] = m_strides[i+1] * m_dimensions[i+1];
144  if (m_strides[i] > 0) m_fastStrides[i] = IndexDivisor(m_strides[i]);
145  }
146  }
147  }
148 
149  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
150  const Dimensions& dimensions() const { return m_dimensions; }
151 
152  EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(EvaluatorPointerType) {
153  m_impl.evalSubExprsIfNeeded(NULL);
154  return true;
155  }
156 
157 #ifdef EIGEN_USE_THREADS
158  template <typename EvalSubExprsCallback>
159  EIGEN_STRONG_INLINE void evalSubExprsIfNeededAsync(
160  EvaluatorPointerType, EvalSubExprsCallback done) {
161  m_impl.evalSubExprsIfNeededAsync(nullptr, [done](bool) { done(true); });
162  }
163 #endif // EIGEN_USE_THREADS
164 
165  EIGEN_STRONG_INLINE void cleanup() {
166  m_impl.cleanup();
167  }
168 
169  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index reverseIndex(
170  Index index) const {
171  eigen_assert(index < dimensions().TotalSize());
172  Index inputIndex = 0;
173  if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
174  EIGEN_UNROLL_LOOP
175  for (int i = NumDims - 1; i > 0; --i) {
176  Index idx = index / m_fastStrides[i];
177  index -= idx * m_strides[i];
178  if (m_reverse[i]) {
179  idx = m_dimensions[i] - idx - 1;
180  }
181  inputIndex += idx * m_strides[i] ;
182  }
183  if (m_reverse[0]) {
184  inputIndex += (m_dimensions[0] - index - 1);
185  } else {
186  inputIndex += index;
187  }
188  } else {
189  EIGEN_UNROLL_LOOP
190  for (int i = 0; i < NumDims - 1; ++i) {
191  Index idx = index / m_fastStrides[i];
192  index -= idx * m_strides[i];
193  if (m_reverse[i]) {
194  idx = m_dimensions[i] - idx - 1;
195  }
196  inputIndex += idx * m_strides[i] ;
197  }
198  if (m_reverse[NumDims-1]) {
199  inputIndex += (m_dimensions[NumDims-1] - index - 1);
200  } else {
201  inputIndex += index;
202  }
203  }
204  return inputIndex;
205  }
206 
207  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(
208  Index index) const {
209  return m_impl.coeff(reverseIndex(index));
210  }
211 
212  template<int LoadMode>
213  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
214  PacketReturnType packet(Index index) const
215  {
216  EIGEN_STATIC_ASSERT((PacketSize > 1), YOU_MADE_A_PROGRAMMING_MISTAKE)
217  eigen_assert(index+PacketSize-1 < dimensions().TotalSize());
218 
219  // TODO(ndjaitly): write a better packing routine that uses
220  // local structure.
221  EIGEN_ALIGN_MAX typename internal::remove_const<CoeffReturnType>::type
222  values[PacketSize];
223  EIGEN_UNROLL_LOOP
224  for (int i = 0; i < PacketSize; ++i) {
225  values[i] = coeff(index+i);
226  }
227  PacketReturnType rslt = internal::pload<PacketReturnType>(values);
228  return rslt;
229  }
230 
231  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
232  internal::TensorBlockResourceRequirements getResourceRequirements() const {
233  const size_t target_size = m_device.lastLevelCacheSize();
234  // Block evaluation reads underlying memory in reverse order, and default
235  // cost model does not properly catch this in bytes stored/loaded.
236  return internal::TensorBlockResourceRequirements::skewed<Scalar>(
237  target_size)
238  .addCostPerCoeff({0, 0, 24});
239  }
240 
241  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBlock
242  block(TensorBlockDesc& desc, TensorBlockScratch& scratch,
243  bool /*root_of_expr_ast*/ = false) const {
244  // TODO(ezhulenev): If underlying tensor expression supports and prefers
245  // block evaluation we must use it. Currently we use coeff and packet
246  // access into the underlying tensor expression.
247  // static const bool useBlockAccessForArgType =
248  // TensorEvaluator<ArgType, Device>::BlockAccess &&
249  // TensorEvaluator<ArgType, Device>::PreferBlockAccess;
250 
251  static const bool isColMajor =
252  static_cast<int>(Layout) == static_cast<int>(ColMajor);
253 
254  static const Index inner_dim_idx = isColMajor ? 0 : NumDims - 1;
255  const bool inner_dim_reversed = m_reverse[inner_dim_idx];
256 
257  // Offset in the output block.
258  Index block_offset = 0;
259 
260  // Offset in the input Tensor.
261  Index input_offset = reverseIndex(desc.offset());
262 
263  // Initialize output block iterator state. Dimension in this array are
264  // always in inner_most -> outer_most order (col major layout).
265  array<BlockIteratorState, NumDims> it;
266  for (int i = 0; i < NumDims; ++i) {
267  const int dim = isColMajor ? i : NumDims - 1 - i;
268  it[i].size = desc.dimension(dim);
269  it[i].count = 0;
270  it[i].reverse = m_reverse[dim];
271 
272  it[i].block_stride =
273  i == 0 ? 1 : (it[i - 1].size * it[i - 1].block_stride);
274  it[i].block_span = it[i].block_stride * (it[i].size - 1);
275 
276  it[i].input_stride = m_strides[dim];
277  it[i].input_span = it[i].input_stride * (it[i].size - 1);
278 
279  if (it[i].reverse) {
280  it[i].input_stride = -1 * it[i].input_stride;
281  it[i].input_span = -1 * it[i].input_span;
282  }
283  }
284 
285  // If multiple inner dimensions have the same reverse flag, check if we can
286  // merge them into a single virtual inner dimension.
287  int effective_inner_dim = 0;
288  for (int i = 1; i < NumDims; ++i) {
289  if (it[i].reverse != it[effective_inner_dim].reverse) break;
290  if (it[i].block_stride != it[effective_inner_dim].size) break;
291  if (it[i].block_stride != numext::abs(it[i].input_stride)) break;
292 
293  it[i].size = it[effective_inner_dim].size * it[i].size;
294 
295  it[i].block_stride = 1;
296  it[i].input_stride = (inner_dim_reversed ? -1 : 1);
297 
298  it[i].block_span = it[i].block_stride * (it[i].size - 1);
299  it[i].input_span = it[i].input_stride * (it[i].size - 1);
300 
301  effective_inner_dim = i;
302  }
303 
304  eigen_assert(it[effective_inner_dim].block_stride == 1);
305  eigen_assert(it[effective_inner_dim].input_stride ==
306  (inner_dim_reversed ? -1 : 1));
307 
308  const Index inner_dim_size = it[effective_inner_dim].size;
309 
310  // Prepare storage for the materialized reverse result.
311  const typename TensorBlock::Storage block_storage =
312  TensorBlock::prepareStorage(desc, scratch);
313  CoeffReturnType* block_buffer = block_storage.data();
314 
315  while (it[NumDims - 1].count < it[NumDims - 1].size) {
316  // Copy inner-most dimension data from reversed location in input.
317  Index dst = block_offset;
318  Index src = input_offset;
319 
320  // NOTE(ezhulenev): Adding vectorized path with internal::preverse showed
321  // worse results in benchmarks than a simple coefficient loop.
322  if (inner_dim_reversed) {
323  for (Index i = 0; i < inner_dim_size; ++i) {
324  block_buffer[dst] = m_impl.coeff(src);
325  ++dst;
326  --src;
327  }
328  } else {
329  for (Index i = 0; i < inner_dim_size; ++i) {
330  block_buffer[dst] = m_impl.coeff(src);
331  ++dst;
332  ++src;
333  }
334  }
335 
336  // For the 1d tensor we need to generate only one inner-most dimension.
337  if ((NumDims - effective_inner_dim) == 1) break;
338 
339  // Update offset.
340  for (Index i = effective_inner_dim + 1; i < NumDims; ++i) {
341  if (++it[i].count < it[i].size) {
342  block_offset += it[i].block_stride;
343  input_offset += it[i].input_stride;
344  break;
345  }
346  if (i != NumDims - 1) it[i].count = 0;
347  block_offset -= it[i].block_span;
348  input_offset -= it[i].input_span;
349  }
350  }
351 
352  return block_storage.AsTensorMaterializedBlock();
353  }
354 
355  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool vectorized) const {
356  double compute_cost = NumDims * (2 * TensorOpCost::AddCost<Index>() +
357  2 * TensorOpCost::MulCost<Index>() +
358  TensorOpCost::DivCost<Index>());
359  for (int i = 0; i < NumDims; ++i) {
360  if (m_reverse[i]) {
361  compute_cost += 2 * TensorOpCost::AddCost<Index>();
362  }
363  }
364  return m_impl.costPerCoeff(vectorized) +
365  TensorOpCost(0, 0, compute_cost, false /* vectorized */, PacketSize);
366  }
367 
368  EIGEN_DEVICE_FUNC typename Storage::Type data() const { return NULL; }
369 
370 #ifdef EIGEN_USE_SYCL
371  // binding placeholder accessors to a command group handler for SYCL
372  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void bind(cl::sycl::handler &cgh) const {
373  m_impl.bind(cgh);
374  }
375 #endif
376 
377  protected:
378  Dimensions m_dimensions;
379  array<Index, NumDims> m_strides;
380  array<IndexDivisor, NumDims> m_fastStrides;
381  TensorEvaluator<ArgType, Device> m_impl;
382  ReverseDimensions m_reverse;
383  const Device EIGEN_DEVICE_REF m_device;
384 
385  private:
386  struct BlockIteratorState {
387  BlockIteratorState()
388  : size(0),
389  count(0),
390  reverse(false),
391  block_stride(0),
392  block_span(0),
393  input_stride(0),
394  input_span(0) {}
395 
396  Index size;
397  Index count;
398  bool reverse;
399  Index block_stride;
400  Index block_span;
401  Index input_stride;
402  Index input_span;
403  };
404 };
405 
406 // Eval as lvalue
407 
408 template <typename ReverseDimensions, typename ArgType, typename Device>
409 struct TensorEvaluator<TensorReverseOp<ReverseDimensions, ArgType>, Device>
410  : public TensorEvaluator<const TensorReverseOp<ReverseDimensions, ArgType>,
411  Device> {
412  typedef TensorEvaluator<const TensorReverseOp<ReverseDimensions, ArgType>,
413  Device> Base;
414  typedef TensorReverseOp<ReverseDimensions, ArgType> XprType;
415  typedef typename XprType::Index Index;
416  static const int NumDims = internal::array_size<ReverseDimensions>::value;
417  typedef DSizes<Index, NumDims> Dimensions;
418 
419  enum {
420  IsAligned = false,
421  PacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess,
422  BlockAccess = false,
423  PreferBlockAccess = false,
424  Layout = TensorEvaluator<ArgType, Device>::Layout,
425  CoordAccess = false, // to be implemented
426  RawAccess = false
427  };
428  EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
429  : Base(op, device) {}
430 
431  typedef typename XprType::Scalar Scalar;
432  typedef typename XprType::CoeffReturnType CoeffReturnType;
433  typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
434  static const int PacketSize = PacketType<CoeffReturnType, Device>::size;
435 
436  //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
437  typedef internal::TensorBlockNotImplemented TensorBlock;
438  //===--------------------------------------------------------------------===//
439 
440  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
441  const Dimensions& dimensions() const { return this->m_dimensions; }
442 
443  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar& coeffRef(Index index) {
444  return this->m_impl.coeffRef(this->reverseIndex(index));
445  }
446 
447  template <int StoreMode> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
448  void writePacket(Index index, const PacketReturnType& x) {
449  EIGEN_STATIC_ASSERT((PacketSize > 1), YOU_MADE_A_PROGRAMMING_MISTAKE)
450  eigen_assert(index+PacketSize-1 < dimensions().TotalSize());
451 
452  // This code is pilfered from TensorMorphing.h
453  EIGEN_ALIGN_MAX CoeffReturnType values[PacketSize];
454  internal::pstore<CoeffReturnType, PacketReturnType>(values, x);
455  EIGEN_UNROLL_LOOP
456  for (int i = 0; i < PacketSize; ++i) {
457  this->coeffRef(index+i) = values[i];
458  }
459  }
460 };
461 
462 
463 } // end namespace Eigen
464 
465 #endif // EIGEN_CXX11_TENSOR_TENSOR_REVERSE_H
WriteAccessors
Namespace containing all symbols from the Eigen library.
EIGEN_DEFAULT_DENSE_INDEX_TYPE Index