Eigen-Contrib  5.0.1
 
Loading...
Searching...
No Matches
TensorMorphing.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// SPDX-License-Identifier: MPL-2.0
10
11#ifndef EIGEN_TENSOR_TENSOR_MORPHING_H
12#define EIGEN_TENSOR_TENSOR_MORPHING_H
13
14// IWYU pragma: private
15#include "./InternalHeaderCheck.h"
16
17namespace Eigen {
18
19namespace internal {
20template <typename NewDimensions, typename XprType>
21struct traits<TensorReshapingOp<NewDimensions, XprType>> : public traits<XprType> {
22 typedef typename XprType::Scalar Scalar;
23 typedef traits<XprType> XprTraits;
24 typedef typename XprTraits::StorageKind StorageKind;
25 typedef typename XprTraits::Index Index;
26 static constexpr int NumDimensions = array_size<NewDimensions>::value;
27 static constexpr int Layout = XprTraits::Layout;
28 typedef typename XprTraits::PointerType PointerType;
29};
30
31template <typename NewDimensions, typename XprType>
32struct eval<TensorReshapingOp<NewDimensions, XprType>, Eigen::Dense> {
33 typedef const TensorReshapingOp<NewDimensions, XprType> EIGEN_DEVICE_REF type;
34};
35
36} // end namespace internal
37
43template <typename NewDimensions, typename XprType>
44class TensorReshapingOp : public TensorBase<TensorReshapingOp<NewDimensions, XprType>, WriteAccessors> {
45 public:
47 typedef typename Eigen::internal::traits<TensorReshapingOp>::Scalar Scalar;
48 typedef std::remove_const_t<typename XprType::CoeffReturnType> CoeffReturnType;
49 typedef typename Eigen::internal::ref_selector<TensorReshapingOp>::type Nested;
50 typedef typename Eigen::internal::traits<TensorReshapingOp>::StorageKind StorageKind;
51 typedef typename Eigen::internal::traits<TensorReshapingOp>::Index Index;
52
53 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorReshapingOp(const XprType& expr, const NewDimensions& dims)
54 : m_xpr(expr), m_dims(dims) {}
55
56 EIGEN_DEVICE_FUNC const NewDimensions& dimensions() const { return m_dims; }
57
58 EIGEN_DEVICE_FUNC const internal::remove_all_t<typename XprType::Nested>& expression() const { return m_xpr; }
59
60 EIGEN_INHERIT_ASSIGNMENT_OPERATORS(TensorReshapingOp)
61
62 protected:
63 typename XprType::Nested m_xpr;
64 const NewDimensions m_dims;
65};
66
67// Eval as rvalue
68template <typename NewDimensions, typename ArgType, typename Device>
69struct TensorEvaluator<const TensorReshapingOp<NewDimensions, ArgType>, Device> {
71 typedef NewDimensions Dimensions;
72
73 typedef typename XprType::Index Index;
74 typedef typename XprType::Scalar Scalar;
75 typedef typename XprType::CoeffReturnType CoeffReturnType;
76 typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
77 typedef StorageMemory<CoeffReturnType, Device> Storage;
78 typedef typename Storage::Type EvaluatorPointerType;
79 typedef StorageMemory<std::remove_const_t<CoeffReturnType>, Device> ConstCastStorage;
80
81 static constexpr int NumOutputDims = internal::array_size<Dimensions>::value;
82 static constexpr int NumInputDims =
83 internal::array_size<typename TensorEvaluator<ArgType, Device>::Dimensions>::value;
84
85 enum ReshapingKind {
86 // We do not use layout information to determine reshaping kind.
87 // Depending on the layout `N` can be inner or outer dimension.
88 OneByN = 0, // expr.reshape(1, N)
89 NByOne = 1, // expr.reshape(N, 1)
90 Runtime = 2 // Reshape dimensions are dynamic (specified at runtime).
91 };
92
93 // clang-format off
94 static const ReshapingKind kind =
95 (NumOutputDims == 2 && internal::index_statically_eq<NewDimensions>(/*index=*/0, /*value=*/1)) ? OneByN
96 : (NumOutputDims == 2 && internal::index_statically_eq<NewDimensions>(/*index=*/1, /*value=*/1)) ? NByOne
97 : Runtime;
98 // clang-format on
99
100 static constexpr int Layout = TensorEvaluator<ArgType, Device>::Layout;
101 enum {
102 IsAligned = TensorEvaluator<ArgType, Device>::IsAligned,
103 PacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess,
104 // For trivial reshapes with raw access to underlying data we will provide
105 // zero overhead block access.
106 // TODO(ezhulenev): Consider adding block access without raw access?
107 BlockAccess = TensorEvaluator<ArgType, Device>::RawAccess && NumInputDims > 0 && NumOutputDims > 0,
108 PreferBlockAccess = false,
109 CoordAccess = false, // to be implemented
110 RawAccess = TensorEvaluator<ArgType, Device>::RawAccess
111 };
112
113 typedef std::remove_const_t<Scalar> ScalarNoConst;
114
115 //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
116 typedef internal::TensorBlockDescriptor<NumOutputDims, Index> TensorBlockDesc;
117 typedef internal::TensorBlockScratchAllocator<Device> TensorBlockScratch;
118
119 typedef typename internal::TensorMaterializedBlock<ScalarNoConst, NumOutputDims, Layout, Index> TensorBlock;
120 //===--------------------------------------------------------------------===//
121
122 EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
123 : m_impl(op.expression(), device), m_dimensions(op.dimensions()) {
124 // The total size of the reshaped tensor must be equal to the total size
125 // of the input tensor.
126 eigen_assert(internal::array_prod(m_impl.dimensions()) == internal::array_prod(op.dimensions()));
127 }
128
129 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_dimensions; }
130
131#ifdef EIGEN_USE_THREADS
132 template <typename EvalSubExprsCallback>
133 EIGEN_STRONG_INLINE void evalSubExprsIfNeededAsync(EvaluatorPointerType data, EvalSubExprsCallback done) {
134 m_impl.evalSubExprsIfNeededAsync(data, std::move(done));
135 }
136#endif
137
138 EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(EvaluatorPointerType data) { return m_impl.evalSubExprsIfNeeded(data); }
139 EIGEN_STRONG_INLINE void cleanup() { m_impl.cleanup(); }
140
141 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const { return m_impl.coeff(index); }
142
143 template <int LoadMode>
144 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packet(Index index) const {
145 return m_impl.template packet<LoadMode>(index);
146 }
147
148 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool vectorized) const {
149 return m_impl.costPerCoeff(vectorized);
150 }
151
152 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE internal::TensorBlockResourceRequirements getResourceRequirements() const {
153 return internal::TensorBlockResourceRequirements::any();
154 }
155
156 struct BlockIteratorState {
157 Index stride;
158 Index span;
159 Index size;
160 Index count;
161 };
162
163 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBlock block(TensorBlockDesc& desc, TensorBlockScratch& scratch,
164 bool /*root_of_expr_ast*/ = false) const {
165 eigen_assert(m_impl.data() != nullptr);
166 eigen_assert((kind == Runtime) || (kind == OneByN && desc.dimensions()[0] == 1) ||
167 (kind == NByOne && desc.dimensions()[1] == 1));
168
169 if (kind == OneByN || kind == NByOne) {
170 // We can guarantee at compile time that block is just a contiguous slice
171 // of the underlying expression memory buffer.
172 return TensorBlock(internal::TensorBlockKind::kView, m_impl.data() + desc.offset(), desc.dimensions());
173 } else {
174 // This will do additional runtime checks, and in the end it might be also
175 // a view, or it might be a block materialized in the temporary buffer.
176 return TensorBlock::materialize(m_impl.data(), m_dimensions, desc, scratch);
177 }
178 }
179
180 EIGEN_DEVICE_FUNC typename Storage::Type data() const { return constCast(m_impl.data()); }
181
182 EIGEN_DEVICE_FUNC const TensorEvaluator<ArgType, Device>& impl() const { return m_impl; }
183
184 protected:
185 TensorEvaluator<ArgType, Device> m_impl;
186 NewDimensions m_dimensions;
187};
188
189// Eval as lvalue
190template <typename NewDimensions, typename ArgType, typename Device>
191struct TensorEvaluator<TensorReshapingOp<NewDimensions, ArgType>, Device>
192 : public TensorEvaluator<const TensorReshapingOp<NewDimensions, ArgType>, Device>
193
194{
195 typedef TensorEvaluator<const TensorReshapingOp<NewDimensions, ArgType>, Device> Base;
196 typedef TensorReshapingOp<NewDimensions, ArgType> XprType;
197 typedef NewDimensions Dimensions;
198
199 static constexpr int Layout = TensorEvaluator<ArgType, Device>::Layout;
200 enum {
201 IsAligned = TensorEvaluator<ArgType, Device>::IsAligned,
202 PacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess,
203 BlockAccess = TensorEvaluator<ArgType, Device>::RawAccess,
204 PreferBlockAccess = false,
205 CoordAccess = false, // to be implemented
206 RawAccess = TensorEvaluator<ArgType, Device>::RawAccess
207 };
208
209 EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device) : Base(op, device) {}
210
211 typedef typename XprType::Index Index;
212 typedef typename XprType::Scalar Scalar;
213 typedef typename XprType::CoeffReturnType CoeffReturnType;
214 typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
215
216 //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
217 typedef internal::TensorBlockDescriptor<TensorEvaluator::NumOutputDims, Index> TensorBlockDesc;
218 //===--------------------------------------------------------------------===//
219
220 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType& coeffRef(Index index) const {
221 return this->m_impl.coeffRef(index);
222 }
223
224 template <int StoreMode>
225 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void writePacket(Index index, const PacketReturnType& x) const {
226 this->m_impl.template writePacket<StoreMode>(index, x);
227 }
228
229 template <typename TensorBlock>
230 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void writeBlock(const TensorBlockDesc& desc, const TensorBlock& block) {
231 eigen_assert(this->m_impl.data() != nullptr);
232
233 typedef typename TensorBlock::XprType TensorBlockExpr;
234 typedef internal::TensorBlockAssignment<Scalar, TensorEvaluator::NumOutputDims, TensorBlockExpr, Index>
235 TensorBlockAssign;
236
237 TensorBlockAssign::Run(TensorBlockAssign::target(desc.dimensions(), internal::strides<Layout>(this->dimensions()),
238 this->m_impl.data(), desc.offset()),
239 block.expr());
240 }
241};
242
250namespace internal {
251template <typename StartIndices, typename Sizes, typename XprType>
252struct traits<TensorSlicingOp<StartIndices, Sizes, XprType>> : public traits<XprType> {
253 typedef typename XprType::Scalar Scalar;
254 typedef traits<XprType> XprTraits;
255 typedef typename XprTraits::StorageKind StorageKind;
256 typedef typename XprTraits::Index Index;
257 static constexpr int NumDimensions = array_size<StartIndices>::value;
258 static constexpr int Layout = XprTraits::Layout;
259 typedef typename XprTraits::PointerType PointerType;
260};
261
262template <typename StartIndices, typename Sizes, typename XprType>
263struct eval<TensorSlicingOp<StartIndices, Sizes, XprType>, Eigen::Dense> {
264 typedef const TensorSlicingOp<StartIndices, Sizes, XprType> EIGEN_DEVICE_REF type;
265};
266
267} // end namespace internal
268
269template <typename StartIndices, typename Sizes, typename XprType>
270class TensorSlicingOp : public TensorBase<TensorSlicingOp<StartIndices, Sizes, XprType>> {
271 public:
272 typedef TensorBase<TensorSlicingOp<StartIndices, Sizes, XprType>> Base;
273 typedef typename Eigen::internal::traits<TensorSlicingOp>::Scalar Scalar;
274 typedef typename XprType::CoeffReturnType CoeffReturnType;
275 typedef typename Eigen::internal::ref_selector<TensorSlicingOp>::type Nested;
276 typedef typename Eigen::internal::traits<TensorSlicingOp>::StorageKind StorageKind;
277 typedef typename Eigen::internal::traits<TensorSlicingOp>::Index Index;
278
279 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorSlicingOp(const XprType& expr, const StartIndices& indices,
280 const Sizes& sizes)
281 : m_xpr(expr), m_indices(indices), m_sizes(sizes) {}
282
283 EIGEN_DEVICE_FUNC const StartIndices& startIndices() const { return m_indices; }
284 EIGEN_DEVICE_FUNC const Sizes& sizes() const { return m_sizes; }
285
286 EIGEN_DEVICE_FUNC const internal::remove_all_t<typename XprType::Nested>& expression() const { return m_xpr; }
287
288 EIGEN_INHERIT_ASSIGNMENT_OPERATORS(TensorSlicingOp)
289
290 protected:
291 typename XprType::Nested m_xpr;
292 const StartIndices m_indices;
293 const Sizes m_sizes;
294};
295
296namespace internal {
297
298// FIXME: Figure out the exact threshold.
299template <typename Index, typename Device, bool BlockAccess>
300struct MemcpyTriggerForSlicing {
301 EIGEN_DEVICE_FUNC MemcpyTriggerForSlicing(const Device& device)
302 : threshold_(static_cast<Index>(2 * device.numThreads())) {}
303 EIGEN_DEVICE_FUNC bool operator()(Index total, Index contiguous) const {
304 const bool prefer_block_evaluation = BlockAccess && total > 32 * 1024;
305 return !prefer_block_evaluation && contiguous > threshold_;
306 }
307
308 private:
309 Index threshold_;
310};
311
312// It is very expensive to start the memcpy kernel on GPU: we therefore only
313// use it for large copies.
314#ifdef EIGEN_USE_GPU
315template <typename Index, bool BlockAccess>
316struct MemcpyTriggerForSlicing<Index, GpuDevice, BlockAccess> {
317 EIGEN_DEVICE_FUNC MemcpyTriggerForSlicing(const GpuDevice&) {}
318 EIGEN_DEVICE_FUNC bool operator()(Index, Index contiguous) const { return contiguous > 4 * 1024 * 1024; }
319};
320#endif
321
322// It is very expensive to start the memcpy kernel on GPU: we therefore only
323// use it for large copies.
324#ifdef EIGEN_USE_SYCL
325template <typename Index, bool BlockAccess>
326struct MemcpyTriggerForSlicing<Index, Eigen::SyclDevice, BlockAccess> {
327 EIGEN_DEVICE_FUNC MemcpyTriggerForSlicing(const SyclDevice&) {}
328 EIGEN_DEVICE_FUNC bool operator()(Index, Index contiguous) const { return contiguous > 4 * 1024 * 1024; }
329};
330#endif
331
332} // namespace internal
333
334// Eval as rvalue
335template <typename StartIndices, typename Sizes, typename ArgType, typename Device>
336struct TensorEvaluator<const TensorSlicingOp<StartIndices, Sizes, ArgType>, Device> {
337 typedef TensorSlicingOp<StartIndices, Sizes, ArgType> XprType;
338 static constexpr int NumDims = internal::array_size<Sizes>::value;
339
340 typedef typename XprType::Index Index;
341 typedef typename XprType::Scalar Scalar;
342 typedef typename XprType::CoeffReturnType CoeffReturnType;
343 typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
344 typedef Sizes Dimensions;
345 typedef StorageMemory<CoeffReturnType, Device> Storage;
346 typedef StorageMemory<std::remove_const_t<CoeffReturnType>, Device> ConstCastStorage;
347 typedef typename Storage::Type EvaluatorPointerType;
348
349 static constexpr int Layout = TensorEvaluator<ArgType, Device>::Layout;
350 enum {
351 // Alignment can't be guaranteed at compile time since it depends on the
352 // slice offsets and sizes.
353 IsAligned = false,
354 PacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess,
355 BlockAccess = TensorEvaluator<ArgType, Device>::BlockAccess,
356 PreferBlockAccess = true,
357 CoordAccess = false,
358 RawAccess = false
359 };
360
361 typedef std::remove_const_t<Scalar> ScalarNoConst;
362
363 //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
364 typedef internal::TensorBlockDescriptor<NumDims, Index> TensorBlockDesc;
365 typedef internal::TensorBlockScratchAllocator<Device> TensorBlockScratch;
366
367 // Tensor slicing does not change the block type.
368 typedef typename TensorEvaluator<const ArgType, Device>::TensorBlock TensorBlock;
369 //===--------------------------------------------------------------------===//
370
371 EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
372 : m_impl(op.expression(), device), m_device(device), m_dimensions(op.sizes()), m_offsets(op.startIndices()) {
373 m_is_identity = true;
374 for (int i = 0; i < internal::array_size<Dimensions>::value; ++i) {
375 eigen_assert(m_impl.dimensions()[i] >= op.sizes()[i] + op.startIndices()[i]);
376 if (m_impl.dimensions()[i] != op.sizes()[i] || op.startIndices()[i] != 0) {
377 m_is_identity = false;
378 }
379 }
380
381 initStrides(m_impl.dimensions(), op.sizes());
382 }
383
384 template <int ND = NumDims>
385 EIGEN_STRONG_INLINE std::enable_if_t<ND == 0, void> initStrides(
386 const typename TensorEvaluator<ArgType, Device>::Dimensions& /*input_dims*/, const Sizes& /*output_dims*/) {}
387
388 template <int ND = NumDims>
389 EIGEN_STRONG_INLINE std::enable_if_t<(ND > 0), void> initStrides(
390 const typename TensorEvaluator<ArgType, Device>::Dimensions& input_dims, const Sizes& output_dims) {
391 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
392 m_inputStrides[0] = 1;
393 for (int i = 1; i < ND; ++i) {
394 m_inputStrides[i] = m_inputStrides[i - 1] * input_dims[i - 1];
395 }
396
397 // Don't initialize m_fastOutputStrides[0] since it won't ever be accessed.
398 m_outputStrides[0] = 1;
399 for (int i = 1; i < ND; ++i) {
400 m_outputStrides[i] = m_outputStrides[i - 1] * output_dims[i - 1];
401 m_fastOutputStrides[i] = internal::TensorIntDivisor<Index>(m_outputStrides[i] > 0 ? m_outputStrides[i] : 1);
402 }
403 } else {
404 m_inputStrides[ND - 1] = 1;
405 for (int i = ND - 2; i >= 0; --i) {
406 m_inputStrides[i] = m_inputStrides[i + 1] * input_dims[i + 1];
407 }
408
409 // Don't initialize m_fastOutputStrides[NumDims-1] since it won't ever be accessed.
410 m_outputStrides[ND - 1] = 1;
411 for (int i = ND - 2; i >= 0; --i) {
412 m_outputStrides[i] = m_outputStrides[i + 1] * output_dims[i + 1];
413 m_fastOutputStrides[i] = internal::TensorIntDivisor<Index>(m_outputStrides[i] > 0 ? m_outputStrides[i] : 1);
414 }
415 }
416 }
417
418 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_dimensions; }
419
420 EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(EvaluatorPointerType data) {
421 m_impl.evalSubExprsIfNeeded(nullptr);
422 EIGEN_IF_CONSTEXPR (!NumTraits<std::remove_const_t<Scalar>>::RequireInitialization) {
423 if (data && m_impl.data()) {
424 Index contiguous_values = 1;
425 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
426 for (int i = 0; i < NumDims; ++i) {
427 contiguous_values *= dimensions()[i];
428 if (dimensions()[i] != m_impl.dimensions()[i]) {
429 break;
430 }
431 }
432 } else {
433 for (int i = NumDims - 1; i >= 0; --i) {
434 contiguous_values *= dimensions()[i];
435 if (dimensions()[i] != m_impl.dimensions()[i]) {
436 break;
437 }
438 }
439 }
440 // Use memcpy if it's going to be faster than using the regular evaluation.
441 const internal::MemcpyTriggerForSlicing<Index, Device, BlockAccess> trigger(m_device);
442 if (trigger(internal::array_prod(dimensions()), contiguous_values)) {
443 EvaluatorPointerType src = (EvaluatorPointerType)m_impl.data();
444 for (Index i = 0; i < internal::array_prod(dimensions()); i += contiguous_values) {
445 Index offset = srcCoeff(i);
446 m_device.memcpy((void*)(m_device.get(data + i)), m_device.get(src + offset),
447 contiguous_values * sizeof(Scalar));
448 }
449 return false;
450 }
451 }
452 }
453 return true;
454 }
455
456#ifdef EIGEN_USE_THREADS
457 template <typename EvalSubExprsCallback>
458 EIGEN_STRONG_INLINE void evalSubExprsIfNeededAsync(EvaluatorPointerType /*data*/, EvalSubExprsCallback done) {
459 m_impl.evalSubExprsIfNeededAsync(nullptr, [done](bool) { done(true); });
460 }
461#endif // EIGEN_USE_THREADS
462
463 EIGEN_STRONG_INLINE void cleanup() { m_impl.cleanup(); }
464
465 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const {
466 if (m_is_identity) {
467 return m_impl.coeff(index);
468 } else {
469 return m_impl.coeff(srcCoeff(index));
470 }
471 }
472
473 template <int LoadMode>
474 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packet(Index index) const {
475 const int packetSize = PacketType<CoeffReturnType, Device>::size;
476 EIGEN_STATIC_ASSERT((packetSize > 1), YOU_MADE_A_PROGRAMMING_MISTAKE)
477 eigen_assert(index + packetSize - 1 < internal::array_prod(dimensions()));
478
479 if (m_is_identity) {
480 return m_impl.template packet<LoadMode>(index);
481 }
482
483 Index inputIndices[] = {0, 0};
484 Index indices[] = {index, index + packetSize - 1};
485 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
486 EIGEN_UNROLL_LOOP
487 for (int i = NumDims - 1; i > 0; --i) {
488 const Index idx0 = indices[0] / m_fastOutputStrides[i];
489 const Index idx1 = indices[1] / m_fastOutputStrides[i];
490 inputIndices[0] += (idx0 + m_offsets[i]) * m_inputStrides[i];
491 inputIndices[1] += (idx1 + m_offsets[i]) * m_inputStrides[i];
492 indices[0] -= idx0 * m_outputStrides[i];
493 indices[1] -= idx1 * m_outputStrides[i];
494 }
495 inputIndices[0] += (indices[0] + m_offsets[0]);
496 inputIndices[1] += (indices[1] + m_offsets[0]);
497 } else {
498 EIGEN_UNROLL_LOOP
499 for (int i = 0; i < NumDims - 1; ++i) {
500 const Index idx0 = indices[0] / m_fastOutputStrides[i];
501 const Index idx1 = indices[1] / m_fastOutputStrides[i];
502 inputIndices[0] += (idx0 + m_offsets[i]) * m_inputStrides[i];
503 inputIndices[1] += (idx1 + m_offsets[i]) * m_inputStrides[i];
504 indices[0] -= idx0 * m_outputStrides[i];
505 indices[1] -= idx1 * m_outputStrides[i];
506 }
507 inputIndices[0] += (indices[0] + m_offsets[NumDims - 1]);
508 inputIndices[1] += (indices[1] + m_offsets[NumDims - 1]);
509 }
510 if (inputIndices[1] - inputIndices[0] == packetSize - 1) {
511 PacketReturnType rslt = m_impl.template packet<Unaligned>(inputIndices[0]);
512 return rslt;
513 } else {
514 EIGEN_ALIGN_TO_BOUNDARY(internal::unpacket_traits<PacketReturnType>::alignment)
515 std::remove_const_t<CoeffReturnType> values[packetSize];
516 values[0] = m_impl.coeff(inputIndices[0]);
517 values[packetSize - 1] = m_impl.coeff(inputIndices[1]);
518 EIGEN_UNROLL_LOOP
519 for (int i = 1; i < packetSize - 1; ++i) {
520 values[i] = coeff(index + i);
521 }
522 PacketReturnType rslt = internal::pload<PacketReturnType>(values);
523 return rslt;
524 }
525 }
526
527 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool vectorized) const {
528 return m_impl.costPerCoeff(vectorized) + TensorOpCost(0, 0, m_is_identity ? 1 : NumDims);
529 }
530
531 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE internal::TensorBlockResourceRequirements getResourceRequirements() const {
532 const size_t target_size = m_device.lastLevelCacheSize();
533 return internal::TensorBlockResourceRequirements::merge(
534 internal::TensorBlockResourceRequirements::skewed<Scalar>(target_size), m_impl.getResourceRequirements());
535 }
536
537 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBlock block(TensorBlockDesc& desc, TensorBlockScratch& scratch,
538 bool /*root_of_expr_ast*/ = false) const {
539 TensorBlockDesc arg_desc = desc.WithOffset(srcCoeff(desc.offset()));
540 TensorBlock block = m_impl.block(arg_desc, scratch);
541 if (!arg_desc.HasDestinationBuffer()) desc.DropDestinationBuffer();
542 return block;
543 }
544
545 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE typename Storage::Type data() const {
546 typename Storage::Type result = constCast(m_impl.data());
547 if (result) {
548 Index offset = 0;
549 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
550 for (int i = 0; i < NumDims; ++i) {
551 if (m_dimensions[i] != m_impl.dimensions()[i]) {
552 offset += m_offsets[i] * m_inputStrides[i];
553 for (int j = i + 1; j < NumDims; ++j) {
554 if (m_dimensions[j] > 1) {
555 return nullptr;
556 }
557 offset += m_offsets[j] * m_inputStrides[j];
558 }
559 break;
560 }
561 }
562 } else {
563 for (int i = NumDims - 1; i >= 0; --i) {
564 if (m_dimensions[i] != m_impl.dimensions()[i]) {
565 offset += m_offsets[i] * m_inputStrides[i];
566 for (int j = i - 1; j >= 0; --j) {
567 if (m_dimensions[j] > 1) {
568 return nullptr;
569 }
570 offset += m_offsets[j] * m_inputStrides[j];
571 }
572 break;
573 }
574 }
575 }
576 return result + offset;
577 }
578 return nullptr;
579 }
580
581 protected:
582 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index srcCoeff(Index index) const { return srcCoeffImpl(index); }
583
584 template <int ND = NumDims>
585 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::enable_if_t<ND == 0, Index> srcCoeffImpl(Index /*index*/) const {
586 return 0;
587 }
588
589 template <int ND = NumDims>
590 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::enable_if_t<(ND > 0), Index> srcCoeffImpl(Index index) const {
591 Index inputIndex = 0;
592 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
593 EIGEN_UNROLL_LOOP
594 for (int i = ND - 1; i > 0; --i) {
595 const Index idx = index / m_fastOutputStrides[i];
596 inputIndex += (idx + m_offsets[i]) * m_inputStrides[i];
597 index -= idx * m_outputStrides[i];
598 }
599 inputIndex += (index + m_offsets[0]);
600 } else {
601 EIGEN_UNROLL_LOOP
602 for (int i = 0; i < ND - 1; ++i) {
603 const Index idx = index / m_fastOutputStrides[i];
604 inputIndex += (idx + m_offsets[i]) * m_inputStrides[i];
605 index -= idx * m_outputStrides[i];
606 }
607 inputIndex += (index + m_offsets[ND - 1]);
608 }
609 return inputIndex;
610 }
611
612 array<Index, NumDims> m_outputStrides;
613 array<internal::TensorIntDivisor<Index>, NumDims> m_fastOutputStrides;
614 array<Index, NumDims> m_inputStrides;
615 TensorEvaluator<ArgType, Device> m_impl;
616 const Device EIGEN_DEVICE_REF m_device;
617 Dimensions m_dimensions;
618 bool m_is_identity;
619 const StartIndices m_offsets;
620};
621
622// Eval as lvalue
623template <typename StartIndices, typename Sizes, typename ArgType, typename Device>
624struct TensorEvaluator<TensorSlicingOp<StartIndices, Sizes, ArgType>, Device>
625 : public TensorEvaluator<const TensorSlicingOp<StartIndices, Sizes, ArgType>, Device> {
626 typedef TensorEvaluator<const TensorSlicingOp<StartIndices, Sizes, ArgType>, Device> Base;
627 typedef TensorSlicingOp<StartIndices, Sizes, ArgType> XprType;
628 static constexpr int NumDims = internal::array_size<Sizes>::value;
629
630 typedef typename XprType::Index Index;
631 typedef typename XprType::Scalar Scalar;
632 typedef typename XprType::CoeffReturnType CoeffReturnType;
633 typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
634 typedef Sizes Dimensions;
635
636 static constexpr int Layout = TensorEvaluator<ArgType, Device>::Layout;
637 enum {
638 IsAligned = false,
639 PacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess,
640 BlockAccess = TensorEvaluator<ArgType, Device>::BlockAccess,
641 PreferBlockAccess = true,
642 CoordAccess = false,
643 RawAccess = (NumDims == 1) & TensorEvaluator<ArgType, Device>::RawAccess
644 };
645
646 typedef std::remove_const_t<Scalar> ScalarNoConst;
647
648 //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
649 typedef internal::TensorBlockDescriptor<NumDims, Index> TensorBlockDesc;
650 typedef internal::TensorBlockScratchAllocator<Device> TensorBlockScratch;
651 //===--------------------------------------------------------------------===//
652
653 EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device) : Base(op, device) {}
654
655 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType& coeffRef(Index index) const {
656 if (this->m_is_identity) {
657 return this->m_impl.coeffRef(index);
658 } else {
659 return this->m_impl.coeffRef(this->srcCoeff(index));
660 }
661 }
662
663 template <int StoreMode>
664 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void writePacket(Index index, const PacketReturnType& x) const {
665 if (this->m_is_identity) {
666 this->m_impl.template writePacket<StoreMode>(index, x);
667 return;
668 }
669
670 const int packetSize = PacketType<CoeffReturnType, Device>::size;
671 Index inputIndices[] = {0, 0};
672 Index indices[] = {index, index + packetSize - 1};
673 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
674 EIGEN_UNROLL_LOOP
675 for (int i = NumDims - 1; i > 0; --i) {
676 const Index idx0 = indices[0] / this->m_fastOutputStrides[i];
677 const Index idx1 = indices[1] / this->m_fastOutputStrides[i];
678 inputIndices[0] += (idx0 + this->m_offsets[i]) * this->m_inputStrides[i];
679 inputIndices[1] += (idx1 + this->m_offsets[i]) * this->m_inputStrides[i];
680 indices[0] -= idx0 * this->m_outputStrides[i];
681 indices[1] -= idx1 * this->m_outputStrides[i];
682 }
683 inputIndices[0] += (indices[0] + this->m_offsets[0]);
684 inputIndices[1] += (indices[1] + this->m_offsets[0]);
685 } else {
686 EIGEN_UNROLL_LOOP
687 for (int i = 0; i < NumDims - 1; ++i) {
688 const Index idx0 = indices[0] / this->m_fastOutputStrides[i];
689 const Index idx1 = indices[1] / this->m_fastOutputStrides[i];
690 inputIndices[0] += (idx0 + this->m_offsets[i]) * this->m_inputStrides[i];
691 inputIndices[1] += (idx1 + this->m_offsets[i]) * this->m_inputStrides[i];
692 indices[0] -= idx0 * this->m_outputStrides[i];
693 indices[1] -= idx1 * this->m_outputStrides[i];
694 }
695 inputIndices[0] += (indices[0] + this->m_offsets[NumDims - 1]);
696 inputIndices[1] += (indices[1] + this->m_offsets[NumDims - 1]);
697 }
698 if (inputIndices[1] - inputIndices[0] == packetSize - 1) {
699 this->m_impl.template writePacket<StoreMode>(inputIndices[0], x);
700 } else {
701 EIGEN_ALIGN_TO_BOUNDARY(internal::unpacket_traits<PacketReturnType>::alignment)
702 CoeffReturnType values[packetSize];
703 internal::pstore<CoeffReturnType, PacketReturnType>(values, x);
704 this->m_impl.coeffRef(inputIndices[0]) = values[0];
705 this->m_impl.coeffRef(inputIndices[1]) = values[packetSize - 1];
706 EIGEN_UNROLL_LOOP
707 for (int i = 1; i < packetSize - 1; ++i) {
708 this->coeffRef(index + i) = values[i];
709 }
710 }
711 }
712
713 template <typename TensorBlock>
714 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void writeBlock(const TensorBlockDesc& desc, const TensorBlock& block) {
715 TensorBlockDesc arg_desc = desc.WithOffset(this->srcCoeff(desc.offset()));
716 this->m_impl.writeBlock(arg_desc, block);
717 }
718};
719
720namespace internal {
721template <typename StartIndices, typename StopIndices, typename Strides, typename XprType>
722struct traits<TensorStridingSlicingOp<StartIndices, StopIndices, Strides, XprType>> : public traits<XprType> {
723 typedef typename XprType::Scalar Scalar;
724 typedef traits<XprType> XprTraits;
725 typedef typename XprTraits::StorageKind StorageKind;
726 typedef typename XprTraits::Index Index;
727 static constexpr int NumDimensions = array_size<StartIndices>::value;
728 static constexpr int Layout = XprTraits::Layout;
729 typedef typename XprTraits::PointerType PointerType;
730};
731
732template <typename StartIndices, typename StopIndices, typename Strides, typename XprType>
733struct eval<TensorStridingSlicingOp<StartIndices, StopIndices, Strides, XprType>, Eigen::Dense> {
734 typedef const TensorStridingSlicingOp<StartIndices, StopIndices, Strides, XprType> EIGEN_DEVICE_REF type;
735};
736
737} // end namespace internal
738
739template <typename StartIndices, typename StopIndices, typename Strides, typename XprType>
740class TensorStridingSlicingOp
741 : public TensorBase<TensorStridingSlicingOp<StartIndices, StopIndices, Strides, XprType>> {
742 public:
743 typedef TensorBase<TensorStridingSlicingOp<StartIndices, StopIndices, Strides, XprType>> Base;
744 typedef typename internal::traits<TensorStridingSlicingOp>::Scalar Scalar;
745 typedef typename XprType::CoeffReturnType CoeffReturnType;
746 typedef typename internal::ref_selector<TensorStridingSlicingOp>::type Nested;
747 typedef typename internal::traits<TensorStridingSlicingOp>::StorageKind StorageKind;
748 typedef typename internal::traits<TensorStridingSlicingOp>::Index Index;
749
750 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorStridingSlicingOp(const XprType& expr, const StartIndices& startIndices,
751 const StopIndices& stopIndices, const Strides& strides)
752 : m_xpr(expr), m_startIndices(startIndices), m_stopIndices(stopIndices), m_strides(strides) {}
753
754 EIGEN_DEVICE_FUNC const StartIndices& startIndices() const { return m_startIndices; }
755 EIGEN_DEVICE_FUNC const StartIndices& stopIndices() const { return m_stopIndices; }
756 EIGEN_DEVICE_FUNC const StartIndices& strides() const { return m_strides; }
757
758 EIGEN_DEVICE_FUNC const internal::remove_all_t<typename XprType::Nested>& expression() const { return m_xpr; }
759
760 EIGEN_INHERIT_ASSIGNMENT_OPERATORS(TensorStridingSlicingOp)
761
762 protected:
763 typename XprType::Nested m_xpr;
764 const StartIndices m_startIndices;
765 const StopIndices m_stopIndices;
766 const Strides m_strides;
767};
768
769// Eval as rvalue
770template <typename StartIndices, typename StopIndices, typename Strides, typename ArgType, typename Device>
771struct TensorEvaluator<const TensorStridingSlicingOp<StartIndices, StopIndices, Strides, ArgType>, Device> {
772 typedef TensorStridingSlicingOp<StartIndices, StopIndices, Strides, ArgType> XprType;
773 static constexpr int NumDims = internal::array_size<Strides>::value;
774 typedef typename XprType::Index Index;
775 typedef typename XprType::Scalar Scalar;
776 typedef typename XprType::CoeffReturnType CoeffReturnType;
777 typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
778 static constexpr int PacketSize = PacketType<CoeffReturnType, Device>::size;
779 typedef StorageMemory<CoeffReturnType, Device> Storage;
780 typedef typename Storage::Type EvaluatorPointerType;
781 typedef Strides Dimensions;
782
783 static constexpr int Layout = TensorEvaluator<ArgType, Device>::Layout;
784 enum {
785 // Alignment can't be guaranteed at compile time since it depends on the
786 // slice offsets and sizes.
787 IsAligned = false,
788 PacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess,
789 // Blocks are materialized from the nested evaluator's raw buffer with one
790 // (possibly strided or reversed) line copy per inner run.
791 BlockAccess = TensorEvaluator<ArgType, Device>::RawAccess,
792 PreferBlockAccess = true,
793 RawAccess = false
794 };
795
796 typedef std::remove_const_t<Scalar> ScalarNoConst;
797
798 //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
799 typedef internal::TensorBlockDescriptor<NumDims, Index> TensorBlockDesc;
800 typedef internal::TensorBlockScratchAllocator<Device> TensorBlockScratch;
801
802 typedef typename internal::TensorMaterializedBlock<ScalarNoConst, NumDims, Layout, Index> TensorBlock;
803 //===--------------------------------------------------------------------===//
804
805 EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
806 : m_impl(op.expression(), device), m_device(device), m_strides(op.strides()) {
807 // Handle degenerate intervals by gracefully clamping and allowing m_dimensions to be zero
808 DSizes<Index, NumDims> startIndicesClamped, stopIndicesClamped;
809 for (ptrdiff_t i = 0; i < internal::array_size<Dimensions>::value; ++i) {
810 eigen_assert(m_strides[i] != 0 && "0 stride is invalid");
811 if (m_strides[i] > 0) {
812 startIndicesClamped[i] = clamp(op.startIndices()[i], 0, m_impl.dimensions()[i]);
813 stopIndicesClamped[i] = clamp(op.stopIndices()[i], 0, m_impl.dimensions()[i]);
814 } else {
815 /* implies m_strides[i] < 0 by assert */
816 startIndicesClamped[i] = clamp(op.startIndices()[i], -1, m_impl.dimensions()[i] - 1);
817 stopIndicesClamped[i] = clamp(op.stopIndices()[i], -1, m_impl.dimensions()[i] - 1);
818 }
819 m_startIndices[i] = startIndicesClamped[i];
820 }
821
822 typedef typename TensorEvaluator<ArgType, Device>::Dimensions InputDimensions;
823 const InputDimensions& input_dims = m_impl.dimensions();
824
825 // compute output tensor shape
826 m_is_identity = true;
827 for (int i = 0; i < NumDims; i++) {
828 Index interval = stopIndicesClamped[i] - startIndicesClamped[i];
829 if (interval == 0 || ((interval < 0) != (m_strides[i] < 0))) {
830 m_dimensions[i] = 0;
831 } else {
832 m_dimensions[i] = (interval / m_strides[i]) + (interval % m_strides[i] != 0 ? 1 : 0);
833 eigen_assert(m_dimensions[i] >= 0);
834 }
835 if (m_strides[i] != 1 || interval != m_impl.dimensions()[i]) {
836 m_is_identity = false;
837 }
838 }
839
840 Strides output_dims = m_dimensions;
841
842 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
843 m_inputStrides[0] = m_strides[0];
844 m_offsets[0] = startIndicesClamped[0];
845 Index previousDimProduct = 1;
846 for (int i = 1; i < NumDims; ++i) {
847 previousDimProduct *= input_dims[i - 1];
848 m_inputStrides[i] = previousDimProduct * m_strides[i];
849 m_offsets[i] = startIndicesClamped[i] * previousDimProduct;
850 }
851
852 // Don't initialize m_fastOutputStrides[0] since it won't ever be accessed.
853 m_outputStrides[0] = 1;
854 for (int i = 1; i < NumDims; ++i) {
855 m_outputStrides[i] = m_outputStrides[i - 1] * output_dims[i - 1];
856 m_fastOutputStrides[i] = internal::TensorIntDivisor<Index>(m_outputStrides[i] > 0 ? m_outputStrides[i] : 1);
857 }
858 } else {
859 m_inputStrides[NumDims - 1] = m_strides[NumDims - 1];
860 m_offsets[NumDims - 1] = startIndicesClamped[NumDims - 1];
861 Index previousDimProduct = 1;
862 for (int i = NumDims - 2; i >= 0; --i) {
863 previousDimProduct *= input_dims[i + 1];
864 m_inputStrides[i] = previousDimProduct * m_strides[i];
865 m_offsets[i] = startIndicesClamped[i] * previousDimProduct;
866 }
867
868 m_outputStrides[NumDims - 1] = 1;
869 for (int i = NumDims - 2; i >= 0; --i) {
870 m_outputStrides[i] = m_outputStrides[i + 1] * output_dims[i + 1];
871 m_fastOutputStrides[i] = internal::TensorIntDivisor<Index>(m_outputStrides[i] > 0 ? m_outputStrides[i] : 1);
872 }
873 }
874 }
875
876 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_dimensions; }
877
878 EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(EvaluatorPointerType) {
879 m_impl.evalSubExprsIfNeeded(nullptr);
880 return true;
881 }
882
883 EIGEN_STRONG_INLINE void cleanup() { m_impl.cleanup(); }
884
885 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const {
886 if (m_is_identity) {
887 return m_impl.coeff(index);
888 } else {
889 return m_impl.coeff(srcCoeff(index));
890 }
891 }
892
893 template <int LoadMode>
894 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packet(Index index) const {
895 EIGEN_STATIC_ASSERT((PacketSize > 1), YOU_MADE_A_PROGRAMMING_MISTAKE)
896 eigen_assert(index + PacketSize - 1 < internal::array_prod(dimensions()));
897
898 if (m_is_identity) {
899 return m_impl.template packet<LoadMode>(index);
900 }
901
902 // A packet that stays inside one inner-most slice is an arithmetic
903 // progression in the input: compute the index mapping once per packet and
904 // step by the inner stride, instead of once per coefficient. With inner
905 // stride +/-1 the progression is a contiguous run served by a single
906 // packet load. The two-endpoint distance test used by TensorSlicingOp and
907 // TensorStridingOp is unsound here: with a negative inner stride the
908 // interior of a packet can be permuted even though the endpoints are
909 // PacketSize-1 apart.
910 constexpr int inner_dim = (static_cast<int>(Layout) == static_cast<int>(ColMajor)) ? 0 : NumDims - 1;
911 const Index inner_size = m_dimensions[inner_dim];
912 Index inner_pos;
913 const Index base = srcCoeffInner(index, inner_pos);
914 EIGEN_ALIGN_TO_BOUNDARY(internal::unpacket_traits<PacketReturnType>::alignment)
915 std::remove_const_t<CoeffReturnType> values[PacketSize];
916 if (inner_pos + PacketSize <= inner_size) {
917 const Index inner_stride = m_inputStrides[inner_dim];
918 if (inner_stride == 1) {
919 return m_impl.template packet<Unaligned>(base);
920 } else EIGEN_IF_CONSTEXPR (!std::is_same<internal::remove_all_t<Device>, SyclDevice>::value) {
921 // SYCL packets do not implement preverse; use the scalar gather below.
922 if (inner_stride == -1) {
923 return internal::preverse(m_impl.template packet<Unaligned>(base - (PacketSize - 1)));
924 }
925 }
926 EIGEN_UNROLL_LOOP
927 for (int i = 0; i < PacketSize; ++i) {
928 values[i] = m_impl.coeff(base + i * inner_stride);
929 }
930 return internal::pload<PacketReturnType>(values);
931 }
932
933 // The packet crosses an inner-slice boundary: assemble it scalar by
934 // scalar.
935 EIGEN_UNROLL_LOOP
936 for (int i = 0; i < PacketSize; ++i) {
937 values[i] = coeff(index + i);
938 }
939 return internal::pload<PacketReturnType>(values);
940 }
941
942 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool vectorized) const {
943 constexpr int inner_dim = (static_cast<int>(Layout) == static_cast<int>(ColMajor)) ? 0 : NumDims - 1;
944 const bool packets_stay_in_inner =
945 m_is_identity || (m_dimensions[inner_dim] >= PacketSize && m_dimensions[inner_dim] % PacketSize == 0);
946 const Index inner_stride = m_inputStrides[inner_dim];
947 const bool packetizes_arg =
948 m_is_identity || (packets_stay_in_inner &&
949 (inner_stride == 1 ||
950 (inner_stride == -1 && !std::is_same<internal::remove_all_t<Device>, SyclDevice>::value)));
951 return m_impl.costPerCoeff(vectorized && packetizes_arg) +
952 TensorOpCost(0, 0, m_is_identity ? 1 : NumDims, vectorized && packets_stay_in_inner, PacketSize);
953 }
954
955 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE internal::TensorBlockResourceRequirements getResourceRequirements() const {
956 // Blocks materialize into scratch and are then re-read by the consumer, so
957 // the round trip must stay cache-resident: L1-sized blocks (as in
958 // TensorShuffling) make the extra pass nearly free, while LLC-sized blocks
959 // cost a measured ~40% on cheap coefficient-wise consumers of a
960 // larger-than-cache slice.
961 const size_t target_size = m_device.firstLevelCacheSize();
962 constexpr int inner_dim = (static_cast<int>(Layout) == static_cast<int>(ColMajor)) ? 0 : NumDims - 1;
963 // A non-unit (or negative) inner stride turns the per-line copies into
964 // gathers or reversed reads the default cost model does not see in bytes
965 // loaded and stored; surface that like TensorReverse does.
966 const double gather_cost = m_inputStrides[inner_dim] == 1 ? 0.0 : 24.0;
967 return internal::TensorBlockResourceRequirements::skewed<Scalar>(target_size).addCostPerCoeff({0, 0, gather_cost});
968 }
969
970 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBlock block(TensorBlockDesc& desc, TensorBlockScratch& scratch,
971 bool root_of_expr_ast = false) const {
972 eigen_assert(m_impl.data() != nullptr);
973
974 if (m_is_identity) {
975 return TensorBlock::materialize(m_impl.data(), m_impl.dimensions(), desc, scratch);
976 }
977
978 typedef internal::TensorBlockIO<ScalarNoConst, Index, NumDims, Layout> TensorBlockIO;
979 typedef typename TensorBlockIO::Dst TensorBlockIODst;
980 typedef typename TensorBlockIO::Src TensorBlockIOSrc;
981
982 const typename TensorBlock::Storage block_storage =
983 TensorBlock::prepareStorage(desc, scratch, /*allow_strided_storage=*/root_of_expr_ast);
984
985 // Per output dimension, one step moves the input by the (possibly
986 // negative) dilated stride; the block's first element sits at the mapped
987 // offset of the block descriptor's origin.
988 typename TensorBlockIO::Dimensions input_strides(m_inputStrides);
989 TensorBlockIOSrc src(input_strides, m_impl.data(), srcCoeff(desc.offset()));
990
991 TensorBlockIODst dst(block_storage.dimensions(), block_storage.strides(), block_storage.data());
992
993 TensorBlockIO::Copy(dst, src);
994 return block_storage.AsTensorMaterializedBlock();
995 }
996
997 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE typename Storage::Type data() const { return nullptr; }
998
999 protected:
1000 // Computes the input index of output index `index` and, as a by-product of
1001 // the same fast-divisor walk, the output's inner-dimension coordinate. The
1002 // packet paths use the latter to test whether a whole packet stays inside
1003 // one inner-most slice without spending an extra division on it.
1004 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index srcCoeffInner(Index index, Index& inner_pos) const {
1005 Index inputIndex = 0;
1006 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
1007 EIGEN_UNROLL_LOOP
1008 for (int i = NumDims - 1; i > 0; --i) {
1009 const Index idx = index / m_fastOutputStrides[i];
1010 inputIndex += idx * m_inputStrides[i] + m_offsets[i];
1011 index -= idx * m_outputStrides[i];
1012 }
1013 inputIndex += index * m_inputStrides[0] + m_offsets[0];
1014 } else {
1015 EIGEN_UNROLL_LOOP
1016 for (int i = 0; i < NumDims - 1; ++i) {
1017 const Index idx = index / m_fastOutputStrides[i];
1018 inputIndex += idx * m_inputStrides[i] + m_offsets[i];
1019 index -= idx * m_outputStrides[i];
1020 }
1021 inputIndex += index * m_inputStrides[NumDims - 1] + m_offsets[NumDims - 1];
1022 }
1023 inner_pos = index;
1024 return inputIndex;
1025 }
1026
1027 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index srcCoeff(Index index) const {
1028 Index inner_pos;
1029 return srcCoeffInner(index, inner_pos);
1030 }
1031
1032 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index clamp(Index value, Index min, Index max) {
1033#ifndef SYCL_DEVICE_ONLY
1034 return numext::maxi(min, numext::mini(max, value));
1035#else
1036 return cl::sycl::clamp(value, min, max);
1037#endif
1038 }
1039
1040 array<Index, NumDims> m_outputStrides;
1041 array<internal::TensorIntDivisor<Index>, NumDims> m_fastOutputStrides;
1042 array<Index, NumDims> m_inputStrides;
1043 bool m_is_identity;
1044 TensorEvaluator<ArgType, Device> m_impl;
1045 const Device EIGEN_DEVICE_REF m_device;
1046 DSizes<Index, NumDims> m_startIndices; // clamped startIndices
1047 DSizes<Index, NumDims> m_dimensions;
1048 DSizes<Index, NumDims> m_offsets; // offset in a flattened shape
1049 const Strides m_strides;
1050};
1051
1052// Eval as lvalue
1053template <typename StartIndices, typename StopIndices, typename Strides, typename ArgType, typename Device>
1054struct TensorEvaluator<TensorStridingSlicingOp<StartIndices, StopIndices, Strides, ArgType>, Device>
1055 : public TensorEvaluator<const TensorStridingSlicingOp<StartIndices, StopIndices, Strides, ArgType>, Device> {
1056 typedef TensorEvaluator<const TensorStridingSlicingOp<StartIndices, StopIndices, Strides, ArgType>, Device> Base;
1057 typedef TensorStridingSlicingOp<StartIndices, StopIndices, Strides, ArgType> XprType;
1058 static constexpr int NumDims = internal::array_size<Strides>::value;
1059 static constexpr int Layout = TensorEvaluator<ArgType, Device>::Layout;
1060
1061 enum {
1062 IsAligned = false,
1063 PacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess,
1064 // Blocks are scattered into the nested evaluator's raw buffer with one
1065 // (possibly strided or reversed) line copy per inner run.
1066 BlockAccess = TensorEvaluator<ArgType, Device>::RawAccess,
1067 PreferBlockAccess = true,
1068 CoordAccess = TensorEvaluator<ArgType, Device>::CoordAccess,
1069 RawAccess = false
1070 };
1071
1072 typedef typename XprType::Index Index;
1073 typedef typename XprType::Scalar Scalar;
1074 typedef typename XprType::CoeffReturnType CoeffReturnType;
1075 typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
1076 typedef Strides Dimensions;
1077 typedef std::remove_const_t<Scalar> ScalarNoConst;
1078
1079 //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
1080 typedef internal::TensorBlockDescriptor<NumDims, Index> TensorBlockDesc;
1081 //===--------------------------------------------------------------------===//
1082
1083 EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device) : Base(op, device) {}
1084
1085 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType& coeffRef(Index index) const {
1086 if (this->m_is_identity) {
1087 return this->m_impl.coeffRef(index);
1088 } else {
1089 return this->m_impl.coeffRef(this->srcCoeff(index));
1090 }
1091 }
1092
1093 template <int StoreMode>
1094 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void writePacket(Index index, const PacketReturnType& x) const {
1095 EIGEN_STATIC_ASSERT((Base::PacketSize > 1), YOU_MADE_A_PROGRAMMING_MISTAKE)
1096 eigen_assert(index + Base::PacketSize - 1 < internal::array_prod(this->dimensions()));
1097
1098 if (this->m_is_identity) {
1099 this->m_impl.template writePacket<StoreMode>(index, x);
1100 return;
1101 }
1102
1103 // Mirrors packet() in the rvalue evaluator above: within one inner-most
1104 // slice the packet is an arithmetic progression in the input, so the
1105 // index mapping is computed once per packet; inner stride +/-1 becomes a
1106 // single (possibly reversed) packet store.
1107 constexpr int inner_dim = (static_cast<int>(Layout) == static_cast<int>(ColMajor)) ? 0 : NumDims - 1;
1108 const Index inner_size = this->m_dimensions[inner_dim];
1109 Index inner_pos;
1110 const Index base = this->srcCoeffInner(index, inner_pos);
1111 EIGEN_ALIGN_TO_BOUNDARY(internal::unpacket_traits<PacketReturnType>::alignment)
1112 CoeffReturnType values[Base::PacketSize];
1113 if (inner_pos + Base::PacketSize <= inner_size) {
1114 const Index inner_stride = this->m_inputStrides[inner_dim];
1115 if (inner_stride == 1) {
1116 this->m_impl.template writePacket<Unaligned>(base, x);
1117 return;
1118 } else EIGEN_IF_CONSTEXPR (!std::is_same<internal::remove_all_t<Device>, SyclDevice>::value) {
1119 // SYCL packets do not implement preverse; use the scalar scatter below.
1120 if (inner_stride == -1) {
1121 this->m_impl.template writePacket<Unaligned>(base - (Base::PacketSize - 1), internal::preverse(x));
1122 return;
1123 }
1124 }
1125 internal::pstore<CoeffReturnType, PacketReturnType>(values, x);
1126 EIGEN_UNROLL_LOOP
1127 for (int i = 0; i < Base::PacketSize; ++i) {
1128 this->m_impl.coeffRef(base + i * inner_stride) = values[i];
1129 }
1130 return;
1131 }
1132
1133 internal::pstore<CoeffReturnType, PacketReturnType>(values, x);
1134 EIGEN_UNROLL_LOOP
1135 for (int i = 0; i < Base::PacketSize; ++i) {
1136 this->coeffRef(index + i) = values[i];
1137 }
1138 }
1139
1140 template <typename TensorBlock>
1141 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void writeBlock(const TensorBlockDesc& desc, const TensorBlock& block) {
1142 eigen_assert(this->m_impl.data() != nullptr);
1143
1144 typedef internal::TensorBlockIO<ScalarNoConst, Index, NumDims, Layout> TensorBlockIO;
1145 typedef typename TensorBlockIO::Dst TensorBlockIODst;
1146 typedef typename TensorBlockIO::Src TensorBlockIOSrc;
1147 typedef internal::TensorBlockAssignment<ScalarNoConst, NumDims, typename TensorBlock::XprType, Index>
1148 TensorBlockAssignment;
1149
1150 const Scalar* block_buffer = block.data();
1151
1152 // Unlike TensorShuffling this op preserves dimension order, so with a unit
1153 // inner stride a block expression can be assigned straight into the
1154 // destination through the dilated outer strides, skipping the
1155 // materialize-into-temp pass below. TensorBlockAssignment requires exactly
1156 // that unit inner stride.
1157 constexpr int inner_dim = (static_cast<int>(Layout) == static_cast<int>(ColMajor)) ? 0 : NumDims - 1;
1158 if (block_buffer == nullptr && this->m_inputStrides[inner_dim] == 1) {
1159 const typename TensorBlockIO::Dimensions output_strides(this->m_inputStrides);
1160 TensorBlockAssignment::Run(TensorBlockAssignment::target(desc.dimensions(), output_strides, this->m_impl.data(),
1161 this->srcCoeff(desc.offset())),
1162 block.expr());
1163 return;
1164 }
1165
1166 const typename TensorBlockIO::Dimensions block_strides = internal::strides<Layout>(desc.dimensions());
1167
1168 // TODO(ezhulenev): TensorBlockIO should be able to read from any Eigen
1169 // expression with coefficient and packet access as `src`.
1170 void* mem = nullptr;
1171 if (block_buffer == nullptr) {
1172 mem = this->m_device.allocate(desc.size() * sizeof(Scalar));
1173 ScalarNoConst* buf = static_cast<ScalarNoConst*>(mem);
1174
1175 TensorBlockAssignment::Run(TensorBlockAssignment::target(desc.dimensions(), block_strides, buf), block.expr());
1176
1177 block_buffer = buf;
1178 }
1179
1180 TensorBlockIOSrc src(block_strides, block_buffer);
1181
1182 // Write through the (possibly negative) dilated strides.
1183 typename TensorBlockIO::Dimensions output_strides(this->m_inputStrides);
1184 TensorBlockIODst dst(desc.dimensions(), output_strides, this->m_impl.data(), this->srcCoeff(desc.offset()));
1185
1186 TensorBlockIO::Copy(dst, src);
1187
1188 // Deallocate temporary buffer used for the block materialization.
1189 if (mem != nullptr) this->m_device.deallocate(mem);
1190 }
1191};
1192
1193} // end namespace Eigen
1194
1195#endif // EIGEN_TENSOR_TENSOR_MORPHING_H
The tensor base class.
Definition TensorForwardDeclarations.h:69
Tensor reshaping class.
Definition TensorMorphing.h:44
WriteAccessors
Namespace containing all symbols from the Eigen library.
The tensor evaluator class.
Definition TensorEvaluator.h:47