Please, help us to better know about our user community by answering the following short survey: https://forms.gle/wpyrxWi18ox9Z5ae9
 
Loading...
Searching...
No Matches
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
13namespace Eigen {
14
21namespace internal {
22template<typename ReverseDimensions, typename XprType>
23struct 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
37template<typename ReverseDimensions, typename XprType>
38struct eval<TensorReverseOp<ReverseDimensions, XprType>, Eigen::Dense>
39{
40 typedef const TensorReverseOp<ReverseDimensions, XprType>& type;
41};
42
43template<typename ReverseDimensions, typename XprType>
44struct 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
52template<typename ReverseDimensions, typename XprType>
53class 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
86template<typename ReverseDimensions, typename ArgType, typename Device>
87struct 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
408template <typename ReverseDimensions, typename ArgType, typename Device>
409struct 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