Please, help us to better know about our user community by answering the following short survey: https://forms.gle/wpyrxWi18ox9Z5ae9
TensorScan.h
1 // This file is part of Eigen, a lightweight C++ template library
2 // for linear algebra.
3 //
4 // Copyright (C) 2016 Igor Babuschkin <igor@babuschk.in>
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_SCAN_H
11 #define EIGEN_CXX11_TENSOR_TENSOR_SCAN_H
12 
13 namespace Eigen {
14 
15 namespace internal {
16 
17 template <typename Op, typename XprType>
18 struct traits<TensorScanOp<Op, XprType> >
19  : public traits<XprType> {
20  typedef typename XprType::Scalar Scalar;
21  typedef traits<XprType> XprTraits;
22  typedef typename XprTraits::StorageKind StorageKind;
23  typedef typename XprType::Nested Nested;
24  typedef typename remove_reference<Nested>::type _Nested;
25  static const int NumDimensions = XprTraits::NumDimensions;
26  static const int Layout = XprTraits::Layout;
27  typedef typename XprTraits::PointerType PointerType;
28 };
29 
30 template<typename Op, typename XprType>
31 struct eval<TensorScanOp<Op, XprType>, Eigen::Dense>
32 {
33  typedef const TensorScanOp<Op, XprType>& type;
34 };
35 
36 template<typename Op, typename XprType>
37 struct nested<TensorScanOp<Op, XprType>, 1,
38  typename eval<TensorScanOp<Op, XprType> >::type>
39 {
40  typedef TensorScanOp<Op, XprType> type;
41 };
42 } // end namespace internal
43 
49 template <typename Op, typename XprType>
50 class TensorScanOp
51  : public TensorBase<TensorScanOp<Op, XprType>, ReadOnlyAccessors> {
52 public:
53  typedef typename Eigen::internal::traits<TensorScanOp>::Scalar Scalar;
54  typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
55  typedef typename XprType::CoeffReturnType CoeffReturnType;
56  typedef typename Eigen::internal::nested<TensorScanOp>::type Nested;
57  typedef typename Eigen::internal::traits<TensorScanOp>::StorageKind StorageKind;
58  typedef typename Eigen::internal::traits<TensorScanOp>::Index Index;
59 
60  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorScanOp(
61  const XprType& expr, const Index& axis, bool exclusive = false, const Op& op = Op())
62  : m_expr(expr), m_axis(axis), m_accumulator(op), m_exclusive(exclusive) {}
63 
64  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
65  const Index axis() const { return m_axis; }
66  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
67  const XprType& expression() const { return m_expr; }
68  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
69  const Op accumulator() const { return m_accumulator; }
70  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
71  bool exclusive() const { return m_exclusive; }
72 
73 protected:
74  typename XprType::Nested m_expr;
75  const Index m_axis;
76  const Op m_accumulator;
77  const bool m_exclusive;
78 };
79 
80 
81 namespace internal {
82 
83 template <typename Self>
84 EIGEN_STRONG_INLINE void ReduceScalar(Self& self, Index offset,
85  typename Self::CoeffReturnType* data) {
86  // Compute the scan along the axis, starting at the given offset
87  typename Self::CoeffReturnType accum = self.accumulator().initialize();
88  if (self.stride() == 1) {
89  if (self.exclusive()) {
90  for (Index curr = offset; curr < offset + self.size(); ++curr) {
91  data[curr] = self.accumulator().finalize(accum);
92  self.accumulator().reduce(self.inner().coeff(curr), &accum);
93  }
94  } else {
95  for (Index curr = offset; curr < offset + self.size(); ++curr) {
96  self.accumulator().reduce(self.inner().coeff(curr), &accum);
97  data[curr] = self.accumulator().finalize(accum);
98  }
99  }
100  } else {
101  if (self.exclusive()) {
102  for (Index idx3 = 0; idx3 < self.size(); idx3++) {
103  Index curr = offset + idx3 * self.stride();
104  data[curr] = self.accumulator().finalize(accum);
105  self.accumulator().reduce(self.inner().coeff(curr), &accum);
106  }
107  } else {
108  for (Index idx3 = 0; idx3 < self.size(); idx3++) {
109  Index curr = offset + idx3 * self.stride();
110  self.accumulator().reduce(self.inner().coeff(curr), &accum);
111  data[curr] = self.accumulator().finalize(accum);
112  }
113  }
114  }
115 }
116 
117 template <typename Self>
118 EIGEN_STRONG_INLINE void ReducePacket(Self& self, Index offset,
119  typename Self::CoeffReturnType* data) {
120  using Scalar = typename Self::CoeffReturnType;
121  using Packet = typename Self::PacketReturnType;
122  // Compute the scan along the axis, starting at the calculated offset
123  Packet accum = self.accumulator().template initializePacket<Packet>();
124  if (self.stride() == 1) {
125  if (self.exclusive()) {
126  for (Index curr = offset; curr < offset + self.size(); ++curr) {
127  internal::pstoreu<Scalar, Packet>(data + curr, self.accumulator().finalizePacket(accum));
128  self.accumulator().reducePacket(self.inner().template packet<Unaligned>(curr), &accum);
129  }
130  } else {
131  for (Index curr = offset; curr < offset + self.size(); ++curr) {
132  self.accumulator().reducePacket(self.inner().template packet<Unaligned>(curr), &accum);
133  internal::pstoreu<Scalar, Packet>(data + curr, self.accumulator().finalizePacket(accum));
134  }
135  }
136  } else {
137  if (self.exclusive()) {
138  for (Index idx3 = 0; idx3 < self.size(); idx3++) {
139  const Index curr = offset + idx3 * self.stride();
140  internal::pstoreu<Scalar, Packet>(data + curr, self.accumulator().finalizePacket(accum));
141  self.accumulator().reducePacket(self.inner().template packet<Unaligned>(curr), &accum);
142  }
143  } else {
144  for (Index idx3 = 0; idx3 < self.size(); idx3++) {
145  const Index curr = offset + idx3 * self.stride();
146  self.accumulator().reducePacket(self.inner().template packet<Unaligned>(curr), &accum);
147  internal::pstoreu<Scalar, Packet>(data + curr, self.accumulator().finalizePacket(accum));
148  }
149  }
150  }
151 }
152 
153 template <typename Self, bool Vectorize, bool Parallel>
154 struct ReduceBlock {
155  EIGEN_STRONG_INLINE void operator()(Self& self, Index idx1,
156  typename Self::CoeffReturnType* data) {
157  for (Index idx2 = 0; idx2 < self.stride(); idx2++) {
158  // Calculate the starting offset for the scan
159  Index offset = idx1 + idx2;
160  ReduceScalar(self, offset, data);
161  }
162  }
163 };
164 
165 // Specialization for vectorized reduction.
166 template <typename Self>
167 struct ReduceBlock<Self, /*Vectorize=*/true, /*Parallel=*/false> {
168  EIGEN_STRONG_INLINE void operator()(Self& self, Index idx1,
169  typename Self::CoeffReturnType* data) {
170  using Packet = typename Self::PacketReturnType;
171  const int PacketSize = internal::unpacket_traits<Packet>::size;
172  Index idx2 = 0;
173  for (; idx2 + PacketSize <= self.stride(); idx2 += PacketSize) {
174  // Calculate the starting offset for the packet scan
175  Index offset = idx1 + idx2;
176  ReducePacket(self, offset, data);
177  }
178  for (; idx2 < self.stride(); idx2++) {
179  // Calculate the starting offset for the scan
180  Index offset = idx1 + idx2;
181  ReduceScalar(self, offset, data);
182  }
183  }
184 };
185 
186 // Single-threaded CPU implementation of scan
187 template <typename Self, typename Reducer, typename Device,
188  bool Vectorize =
189  (TensorEvaluator<typename Self::ChildTypeNoConst, Device>::PacketAccess &&
190  internal::reducer_traits<Reducer, Device>::PacketAccess)>
191 struct ScanLauncher {
192  void operator()(Self& self, typename Self::CoeffReturnType* data) {
193  Index total_size = internal::array_prod(self.dimensions());
194 
195  // We fix the index along the scan axis to 0 and perform a
196  // scan per remaining entry. The iteration is split into two nested
197  // loops to avoid an integer division by keeping track of each idx1 and
198  // idx2.
199  for (Index idx1 = 0; idx1 < total_size; idx1 += self.stride() * self.size()) {
200  ReduceBlock<Self, Vectorize, /*Parallel=*/false> block_reducer;
201  block_reducer(self, idx1, data);
202  }
203  }
204 };
205 
206 #ifdef EIGEN_USE_THREADS
207 
208 // Adjust block_size to avoid false sharing of cachelines among
209 // threads. Currently set to twice the cache line size on Intel and ARM
210 // processors.
211 EIGEN_STRONG_INLINE Index AdjustBlockSize(Index item_size, Index block_size) {
212  EIGEN_CONSTEXPR Index kBlockAlignment = 128;
213  const Index items_per_cacheline =
214  numext::maxi<Index>(1, kBlockAlignment / item_size);
215  return items_per_cacheline * divup(block_size, items_per_cacheline);
216 }
217 
218 template <typename Self>
219 struct ReduceBlock<Self, /*Vectorize=*/true, /*Parallel=*/true> {
220  EIGEN_STRONG_INLINE void operator()(Self& self, Index idx1,
221  typename Self::CoeffReturnType* data) {
222  using Scalar = typename Self::CoeffReturnType;
223  using Packet = typename Self::PacketReturnType;
224  const int PacketSize = internal::unpacket_traits<Packet>::size;
225  Index num_scalars = self.stride();
226  Index num_packets = 0;
227  if (self.stride() >= PacketSize) {
228  num_packets = self.stride() / PacketSize;
229  self.device().parallelFor(
230  num_packets,
231  TensorOpCost(PacketSize * self.size(), PacketSize * self.size(),
232  16 * PacketSize * self.size(), true, PacketSize),
233  // Make the shard size large enough that two neighboring threads
234  // won't write to the same cacheline of `data`.
235  [=](Index blk_size) {
236  return AdjustBlockSize(PacketSize * sizeof(Scalar), blk_size);
237  },
238  [&](Index first, Index last) {
239  for (Index packet = first; packet < last; ++packet) {
240  const Index idx2 = packet * PacketSize;
241  ReducePacket(self, idx1 + idx2, data);
242  }
243  });
244  num_scalars -= num_packets * PacketSize;
245  }
246  self.device().parallelFor(
247  num_scalars, TensorOpCost(self.size(), self.size(), 16 * self.size()),
248  // Make the shard size large enough that two neighboring threads
249  // won't write to the same cacheline of `data`.
250  [=](Index blk_size) {
251  return AdjustBlockSize(sizeof(Scalar), blk_size);
252  },
253  [&](Index first, Index last) {
254  for (Index scalar = first; scalar < last; ++scalar) {
255  const Index idx2 = num_packets * PacketSize + scalar;
256  ReduceScalar(self, idx1 + idx2, data);
257  }
258  });
259  }
260 };
261 
262 template <typename Self>
263 struct ReduceBlock<Self, /*Vectorize=*/false, /*Parallel=*/true> {
264  EIGEN_STRONG_INLINE void operator()(Self& self, Index idx1,
265  typename Self::CoeffReturnType* data) {
266  using Scalar = typename Self::CoeffReturnType;
267  self.device().parallelFor(
268  self.stride(), TensorOpCost(self.size(), self.size(), 16 * self.size()),
269  // Make the shard size large enough that two neighboring threads
270  // won't write to the same cacheline of `data`.
271  [=](Index blk_size) {
272  return AdjustBlockSize(sizeof(Scalar), blk_size);
273  },
274  [&](Index first, Index last) {
275  for (Index idx2 = first; idx2 < last; ++idx2) {
276  ReduceScalar(self, idx1 + idx2, data);
277  }
278  });
279  }
280 };
281 
282 // Specialization for multi-threaded execution.
283 template <typename Self, typename Reducer, bool Vectorize>
284 struct ScanLauncher<Self, Reducer, ThreadPoolDevice, Vectorize> {
285  void operator()(Self& self, typename Self::CoeffReturnType* data) {
286  using Scalar = typename Self::CoeffReturnType;
287  using Packet = typename Self::PacketReturnType;
288  const int PacketSize = internal::unpacket_traits<Packet>::size;
289  const Index total_size = internal::array_prod(self.dimensions());
290  const Index inner_block_size = self.stride() * self.size();
291  bool parallelize_by_outer_blocks = (total_size >= (self.stride() * inner_block_size));
292 
293  if ((parallelize_by_outer_blocks && total_size <= 4096) ||
294  (!parallelize_by_outer_blocks && self.stride() < PacketSize)) {
295  ScanLauncher<Self, Reducer, DefaultDevice, Vectorize> launcher;
296  launcher(self, data);
297  return;
298  }
299 
300  if (parallelize_by_outer_blocks) {
301  // Parallelize over outer blocks.
302  const Index num_outer_blocks = total_size / inner_block_size;
303  self.device().parallelFor(
304  num_outer_blocks,
305  TensorOpCost(inner_block_size, inner_block_size,
306  16 * PacketSize * inner_block_size, Vectorize,
307  PacketSize),
308  [=](Index blk_size) {
309  return AdjustBlockSize(inner_block_size * sizeof(Scalar), blk_size);
310  },
311  [&](Index first, Index last) {
312  for (Index idx1 = first; idx1 < last; ++idx1) {
313  ReduceBlock<Self, Vectorize, /*Parallelize=*/false> block_reducer;
314  block_reducer(self, idx1 * inner_block_size, data);
315  }
316  });
317  } else {
318  // Parallelize over inner packets/scalars dimensions when the reduction
319  // axis is not an inner dimension.
320  ReduceBlock<Self, Vectorize, /*Parallelize=*/true> block_reducer;
321  for (Index idx1 = 0; idx1 < total_size;
322  idx1 += self.stride() * self.size()) {
323  block_reducer(self, idx1, data);
324  }
325  }
326  }
327 };
328 #endif // EIGEN_USE_THREADS
329 
330 #if defined(EIGEN_USE_GPU) && (defined(EIGEN_GPUCC))
331 
332 // GPU implementation of scan
333 // TODO(ibab) This placeholder implementation performs multiple scans in
334 // parallel, but it would be better to use a parallel scan algorithm and
335 // optimize memory access.
336 template <typename Self, typename Reducer>
337 __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void ScanKernel(Self self, Index total_size, typename Self::CoeffReturnType* data) {
338  // Compute offset as in the CPU version
339  Index val = threadIdx.x + blockIdx.x * blockDim.x;
340  Index offset = (val / self.stride()) * self.stride() * self.size() + val % self.stride();
341 
342  if (offset + (self.size() - 1) * self.stride() < total_size) {
343  // Compute the scan along the axis, starting at the calculated offset
344  typename Self::CoeffReturnType accum = self.accumulator().initialize();
345  for (Index idx = 0; idx < self.size(); idx++) {
346  Index curr = offset + idx * self.stride();
347  if (self.exclusive()) {
348  data[curr] = self.accumulator().finalize(accum);
349  self.accumulator().reduce(self.inner().coeff(curr), &accum);
350  } else {
351  self.accumulator().reduce(self.inner().coeff(curr), &accum);
352  data[curr] = self.accumulator().finalize(accum);
353  }
354  }
355  }
356  __syncthreads();
357 
358 }
359 
360 template <typename Self, typename Reducer, bool Vectorize>
361 struct ScanLauncher<Self, Reducer, GpuDevice, Vectorize> {
362  void operator()(const Self& self, typename Self::CoeffReturnType* data) {
363  Index total_size = internal::array_prod(self.dimensions());
364  Index num_blocks = (total_size / self.size() + 63) / 64;
365  Index block_size = 64;
366 
367  LAUNCH_GPU_KERNEL((ScanKernel<Self, Reducer>), num_blocks, block_size, 0, self.device(), self, total_size, data);
368  }
369 };
370 #endif // EIGEN_USE_GPU && (EIGEN_GPUCC)
371 
372 } // namespace internal
373 
374 // Eval as rvalue
375 template <typename Op, typename ArgType, typename Device>
376 struct TensorEvaluator<const TensorScanOp<Op, ArgType>, Device> {
377 
378  typedef TensorScanOp<Op, ArgType> XprType;
379  typedef typename XprType::Index Index;
380  typedef const ArgType ChildTypeNoConst;
381  typedef const ArgType ChildType;
382  static const int NumDims = internal::array_size<typename TensorEvaluator<ArgType, Device>::Dimensions>::value;
383  typedef DSizes<Index, NumDims> Dimensions;
384  typedef typename internal::remove_const<typename XprType::Scalar>::type Scalar;
385  typedef typename XprType::CoeffReturnType CoeffReturnType;
386  typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
387  typedef TensorEvaluator<const TensorScanOp<Op, ArgType>, Device> Self;
388  typedef StorageMemory<Scalar, Device> Storage;
389  typedef typename Storage::Type EvaluatorPointerType;
390 
391  enum {
392  IsAligned = false,
393  PacketAccess = (PacketType<CoeffReturnType, Device>::size > 1),
394  BlockAccess = false,
395  PreferBlockAccess = false,
396  Layout = TensorEvaluator<ArgType, Device>::Layout,
397  CoordAccess = false,
398  RawAccess = true
399  };
400 
401  //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
402  typedef internal::TensorBlockNotImplemented TensorBlock;
403  //===--------------------------------------------------------------------===//
404 
405  EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
406  : m_impl(op.expression(), device),
407  m_device(device),
408  m_exclusive(op.exclusive()),
409  m_accumulator(op.accumulator()),
410  m_size(m_impl.dimensions()[op.axis()]),
411  m_stride(1), m_consume_dim(op.axis()),
412  m_output(NULL) {
413 
414  // Accumulating a scalar isn't supported.
415  EIGEN_STATIC_ASSERT((NumDims > 0), YOU_MADE_A_PROGRAMMING_MISTAKE);
416  eigen_assert(op.axis() >= 0 && op.axis() < NumDims);
417 
418  // Compute stride of scan axis
419  const Dimensions& dims = m_impl.dimensions();
420  if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
421  for (int i = 0; i < op.axis(); ++i) {
422  m_stride = m_stride * dims[i];
423  }
424  } else {
425  // dims can only be indexed through unsigned integers,
426  // so let's use an unsigned type to let the compiler knows.
427  // This prevents stupid warnings: ""'*((void*)(& evaluator)+64)[18446744073709551615]' may be used uninitialized in this function"
428  unsigned int axis = internal::convert_index<unsigned int>(op.axis());
429  for (unsigned int i = NumDims - 1; i > axis; --i) {
430  m_stride = m_stride * dims[i];
431  }
432  }
433  }
434 
435  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const {
436  return m_impl.dimensions();
437  }
438 
439  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Index& stride() const {
440  return m_stride;
441  }
442 
443  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Index& consume_dim() const {
444  return m_consume_dim;
445  }
446 
447  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Index& size() const {
448  return m_size;
449  }
450 
451  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Op& accumulator() const {
452  return m_accumulator;
453  }
454 
455  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool exclusive() const {
456  return m_exclusive;
457  }
458 
459  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const TensorEvaluator<ArgType, Device>& inner() const {
460  return m_impl;
461  }
462 
463  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Device& device() const {
464  return m_device;
465  }
466 
467  EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(EvaluatorPointerType data) {
468  m_impl.evalSubExprsIfNeeded(NULL);
469  internal::ScanLauncher<Self, Op, Device> launcher;
470  if (data) {
471  launcher(*this, data);
472  return false;
473  }
474 
475  const Index total_size = internal::array_prod(dimensions());
476  m_output = static_cast<EvaluatorPointerType>(m_device.get((Scalar*) m_device.allocate_temp(total_size * sizeof(Scalar))));
477  launcher(*this, m_output);
478  return true;
479  }
480 
481  template<int LoadMode>
482  EIGEN_DEVICE_FUNC PacketReturnType packet(Index index) const {
483  return internal::ploadt<PacketReturnType, LoadMode>(m_output + index);
484  }
485 
486  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE EvaluatorPointerType data() const
487  {
488  return m_output;
489  }
490 
491  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const
492  {
493  return m_output[index];
494  }
495 
496  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool) const {
497  return TensorOpCost(sizeof(CoeffReturnType), 0, 0);
498  }
499 
500  EIGEN_STRONG_INLINE void cleanup() {
501  if (m_output) {
502  m_device.deallocate_temp(m_output);
503  m_output = NULL;
504  }
505  m_impl.cleanup();
506  }
507 
508 #ifdef EIGEN_USE_SYCL
509  // binding placeholder accessors to a command group handler for SYCL
510  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void bind(cl::sycl::handler &cgh) const {
511  m_impl.bind(cgh);
512  m_output.bind(cgh);
513  }
514 #endif
515 protected:
516  TensorEvaluator<ArgType, Device> m_impl;
517  const Device EIGEN_DEVICE_REF m_device;
518  const bool m_exclusive;
519  Op m_accumulator;
520  const Index m_size;
521  Index m_stride;
522  Index m_consume_dim;
523  EvaluatorPointerType m_output;
524 };
525 
526 } // end namespace Eigen
527 
528 #endif // EIGEN_CXX11_TENSOR_TENSOR_SCAN_H
Namespace containing all symbols from the Eigen library.
static const symbolic::SymbolExpr< internal::symbolic_last_tag > last
EIGEN_DEFAULT_DENSE_INDEX_TYPE Index