Eigen-Contrib  5.0.1
 
Loading...
Searching...
No Matches
TensorExecutor.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_EXECUTOR_H
12#define EIGEN_TENSOR_TENSOR_EXECUTOR_H
13
14// IWYU pragma: private
15#include "./InternalHeaderCheck.h"
16
17namespace Eigen {
18
19namespace internal {
20
37template <typename Expression, typename Device, bool Vectorizable, TiledEvaluation Tiling>
39 public:
40 typedef typename Expression::Index StorageIndex;
41
42 // Including `contrib/Eigen/Tensor` in different translation units
43 // with/without `EIGEN_USE_THREADS` or `EIGEN_USE_GPU` is a potential ODR
44 // violation. If this template is instantiated with a non-default device, it
45 // means that this header file was included without defining
46 // `EIGEN_USE_THREADS`, `EIGEN_USE_GPU` or `EIGEN_USE_SYCL`.
47 static_assert(std::is_same<Device, DefaultDevice>::value,
48 "Default executor instantiated with non-default device. "
49 "You must #define EIGEN_USE_THREADS, EIGEN_USE_GPU or "
50 "EIGEN_USE_SYCL before including Eigen headers.");
51
52 static EIGEN_STRONG_INLINE void run(const Expression& expr, const Device& device = DefaultDevice()) {
53 TensorEvaluator<Expression, Device> evaluator(expr, device);
54 const bool needs_assign = evaluator.evalSubExprsIfNeeded(nullptr);
55 if (needs_assign) {
56 const StorageIndex size = static_cast<StorageIndex>(array_prod(evaluator.dimensions()));
57 for (StorageIndex i = 0; i < size; ++i) {
58 evaluator.evalScalar(i);
59 }
60 }
61 evaluator.cleanup();
62 }
63};
64
69template <typename Expression, typename Device, typename DoneCallback, bool Vectorizable, TiledEvaluation Tiling>
71
75template <typename Expression>
76class TensorExecutor<Expression, DefaultDevice, /*Vectorizable=*/true,
77 /*Tiling=*/TiledEvaluation::Off> {
78 public:
79 typedef typename Expression::Index StorageIndex;
80
81 static EIGEN_STRONG_INLINE void run(const Expression& expr, const DefaultDevice& device = DefaultDevice()) {
82 TensorEvaluator<Expression, DefaultDevice> evaluator(expr, device);
83 const bool needs_assign = evaluator.evalSubExprsIfNeeded(nullptr);
84 if (needs_assign) {
85 const StorageIndex size = static_cast<StorageIndex>(array_prod(evaluator.dimensions()));
86 const int PacketSize =
87 unpacket_traits<typename TensorEvaluator<Expression, DefaultDevice>::PacketReturnType>::size;
88
89 // Give compiler a strong possibility to unroll the loop. But don't insist
90 // on unrolling, because if the function is expensive compiler should not
91 // unroll the loop at the expense of inlining.
92 const StorageIndex UnrolledSize = (size / (4 * PacketSize)) * 4 * PacketSize;
93 for (StorageIndex i = 0; i < UnrolledSize; i += 4 * PacketSize) {
94 for (StorageIndex j = 0; j < 4; j++) {
95 evaluator.evalPacket(i + j * PacketSize);
96 }
97 }
98 const StorageIndex VectorizedSize = (size / PacketSize) * PacketSize;
99 for (StorageIndex i = UnrolledSize; i < VectorizedSize; i += PacketSize) {
100 evaluator.evalPacket(i);
101 }
102 for (StorageIndex i = VectorizedSize; i < size; ++i) {
103 evaluator.evalScalar(i);
104 }
105 }
106 evaluator.cleanup();
107 }
108};
109
114template <typename Expression, bool Vectorizable>
115class TensorExecutor<Expression, DefaultDevice, Vectorizable,
116 /*Tiling=*/TiledEvaluation::On> {
117 public:
118 typedef typename traits<Expression>::Scalar Scalar;
119 typedef std::remove_const_t<Scalar> ScalarNoConst;
120
122 typedef typename traits<Expression>::Index StorageIndex;
123
124 static constexpr int NumDims = traits<Expression>::NumDimensions;
125
126 EIGEN_DEVICE_FUNC static EIGEN_STRONG_INLINE void run(const Expression& expr,
127 const DefaultDevice& device = DefaultDevice()) {
128 typedef TensorBlockMapper<NumDims, Evaluator::Layout, StorageIndex> TensorBlockMapper;
129
130 typedef internal::TensorBlockDescriptor<NumDims, StorageIndex> TensorBlockDesc;
131 typedef internal::TensorBlockScratchAllocator<DefaultDevice> TensorBlockScratch;
132
133 Evaluator evaluator(expr, device);
134
135 // TODO(ezhulenev): Do not use tiling for small tensors?
136 const bool needs_assign = evaluator.evalSubExprsIfNeeded(nullptr);
137
138 if (needs_assign) {
139 // Query expression tree for desired block size/shape.
140 const TensorBlockResourceRequirements requirements = evaluator.getResourceRequirements();
141
142 const TensorBlockMapper block_mapper(typename TensorBlockDesc::Dimensions(evaluator.dimensions()), requirements);
143
144 // Share scratch memory allocator between all blocks.
145 TensorBlockScratch scratch(device);
146
147 const StorageIndex total_block_count = block_mapper.blockCount();
148 for (StorageIndex i = 0; i < total_block_count; ++i) {
149 TensorBlockDesc desc = block_mapper.blockDescriptor(i);
150 evaluator.evalBlock(desc, scratch);
151 scratch.reset();
152 }
153 }
154 evaluator.cleanup();
155 }
156};
157
169#ifdef EIGEN_USE_THREADS
170
171template <typename TensorBlockMapper>
172struct TensorExecutorTilingContext {
173 TensorExecutorTilingContext() = default;
174 TensorExecutorTilingContext(const TensorBlockMapper& b_mapper, const TensorOpCost& b_cost, size_t b_aligned_size)
175 : block_mapper(b_mapper), cost(b_cost), aligned_blocksize(b_aligned_size) {}
176
177 TensorBlockMapper block_mapper; // navigate through blocks
178 TensorOpCost cost; // cost of computing a single block
179 size_t aligned_blocksize; // block size after memory alignment
180};
181
182// Computes block evaluation parameters, and allocates temporary memory buffer
183// for blocks. See TensorExecutor/TensorAsyncExecutor (Tiling=On) below.
184template <typename Evaluator, typename TensorBlockMapper, bool Vectorizable>
185TensorExecutorTilingContext<TensorBlockMapper> GetTensorExecutorTilingContext(const Evaluator& evaluator) {
186 // Query expression tree for desired block size/shape.
187 TensorBlockResourceRequirements requirements = evaluator.getResourceRequirements();
188
189 // Update target block size based on cost model.
190 double taskSize = TensorCostModel<ThreadPoolDevice>::taskSize(1, requirements.cost_per_coeff);
191 requirements.size = static_cast<size_t>(1.0 / taskSize);
192
193 TensorBlockMapper block_mapper(typename TensorBlockMapper::Dimensions(evaluator.dimensions()), requirements);
194
195 size_t block_size = block_mapper.blockTotalSize();
196 const size_t align = numext::maxi(EIGEN_MAX_ALIGN_BYTES, 1);
197 const size_t aligned_blocksize =
198 align * numext::div_ceil<size_t>(block_size * sizeof(typename Evaluator::Scalar), align);
199
200 return {block_mapper, requirements.cost_per_coeff * block_size, aligned_blocksize};
201}
202
203template <typename Evaluator, typename StorageIndex, bool Vectorizable>
204struct EvalRange {
205 static void run(Evaluator* evaluator_in, const StorageIndex firstIdx, const StorageIndex lastIdx) {
206 Evaluator evaluator = *evaluator_in;
207 eigen_assert(lastIdx >= firstIdx);
208 for (StorageIndex i = firstIdx; i < lastIdx; ++i) {
209 evaluator.evalScalar(i);
210 }
211 }
212
213 static StorageIndex alignBlockSize(StorageIndex size) { return size; }
214};
215
216template <typename Evaluator, typename StorageIndex>
217struct EvalRange<Evaluator, StorageIndex, /*Vectorizable*/ true> {
218 static constexpr int PacketSize = unpacket_traits<typename Evaluator::PacketReturnType>::size;
219
220 static void run(Evaluator* evaluator_in, const StorageIndex firstIdx, const StorageIndex lastIdx) {
221 Evaluator evaluator = *evaluator_in;
222 eigen_assert(lastIdx >= firstIdx);
223 StorageIndex i = firstIdx;
224 if (lastIdx - firstIdx >= PacketSize) {
225 eigen_assert(firstIdx % PacketSize == 0);
226 StorageIndex last_chunk_offset = lastIdx - 4 * PacketSize;
227 // Give compiler a strong possibility to unroll the loop. But don't insist
228 // on unrolling, because if the function is expensive compiler should not
229 // unroll the loop at the expense of inlining.
230 for (; i <= last_chunk_offset; i += 4 * PacketSize) {
231 for (StorageIndex j = 0; j < 4; j++) {
232 evaluator.evalPacket(i + j * PacketSize);
233 }
234 }
235 last_chunk_offset = lastIdx - PacketSize;
236 for (; i <= last_chunk_offset; i += PacketSize) {
237 evaluator.evalPacket(i);
238 }
239 }
240 for (; i < lastIdx; ++i) {
241 evaluator.evalScalar(i);
242 }
243 }
244
245 static StorageIndex alignBlockSize(StorageIndex size) {
246 // Align block size to packet size and account for unrolling in run above.
247 if (size >= 16 * PacketSize) {
248 return (size + 4 * PacketSize - 1) & ~(4 * PacketSize - 1);
249 }
250 // Aligning to 4 * PacketSize would increase block size by more than 25%.
251 return (size + PacketSize - 1) & ~(PacketSize - 1);
252 }
253};
254
255// Evaluates a range of blocks for the tiled executors below. Each task copies
256// the evaluator so that concurrent tasks never share per-instance state (e.g.
257// stateful functors reached through coeff()), matching EvalRange for the
258// non-tiled path.
259template <typename Evaluator, typename BlockMapper, typename IndexType>
260struct EvalBlockRange {
261 static void run(Evaluator* evaluator_in, const BlockMapper& block_mapper, const ThreadPoolDevice& device,
262 IndexType firstBlockIdx, IndexType lastBlockIdx) {
263 typedef TensorBlockScratchAllocator<ThreadPoolDevice> TensorBlockScratch;
264 Evaluator evaluator = *evaluator_in;
265 eigen_assert(lastBlockIdx >= firstBlockIdx);
266 TensorBlockScratch scratch(device);
267
268 for (IndexType block_idx = firstBlockIdx; block_idx < lastBlockIdx; ++block_idx) {
269 auto desc = block_mapper.blockDescriptor(block_idx);
270 evaluator.evalBlock(desc, scratch);
271 scratch.reset();
272 }
273 }
274};
275
276template <typename Expression, bool Vectorizable, TiledEvaluation Tiling>
277class TensorExecutor<Expression, ThreadPoolDevice, Vectorizable, Tiling> {
278 public:
279 typedef typename Expression::Index StorageIndex;
280
281 static EIGEN_STRONG_INLINE void run(const Expression& expr, const ThreadPoolDevice& device) {
282 typedef TensorEvaluator<Expression, ThreadPoolDevice> Evaluator;
283 typedef EvalRange<Evaluator, StorageIndex, Vectorizable> EvalRange;
284
285 Evaluator evaluator(expr, device);
286 const bool needs_assign = evaluator.evalSubExprsIfNeeded(nullptr);
287 if (needs_assign) {
288 const StorageIndex size = static_cast<StorageIndex>(array_prod(evaluator.dimensions()));
289 device.parallelFor(
290 size, evaluator.costPerCoeff(Vectorizable), EvalRange::alignBlockSize,
291 [&evaluator](StorageIndex firstIdx, StorageIndex lastIdx) { EvalRange::run(&evaluator, firstIdx, lastIdx); });
292 }
293 evaluator.cleanup();
294 }
295};
296
297template <typename Expression, bool Vectorizable>
298class TensorExecutor<Expression, ThreadPoolDevice, Vectorizable,
299 /*Tiling=*/TiledEvaluation::On> {
300 public:
301 typedef typename traits<Expression>::Index IndexType;
302 typedef typename traits<Expression>::Scalar Scalar;
303 typedef std::remove_const_t<Scalar> ScalarNoConst;
304
305 static constexpr int NumDims = traits<Expression>::NumDimensions;
306
307 typedef TensorEvaluator<Expression, ThreadPoolDevice> Evaluator;
308 typedef TensorBlockMapper<NumDims, Evaluator::Layout, IndexType> BlockMapper;
309 typedef TensorExecutorTilingContext<BlockMapper> TilingContext;
310
311 typedef internal::TensorBlockDescriptor<NumDims, IndexType> TensorBlockDesc;
312 typedef internal::TensorBlockScratchAllocator<ThreadPoolDevice> TensorBlockScratch;
313
314 static EIGEN_STRONG_INLINE void run(const Expression& expr, const ThreadPoolDevice& device) {
315 Evaluator evaluator(expr, device);
316
317 const bool needs_assign = evaluator.evalSubExprsIfNeeded(nullptr);
318 if (needs_assign) {
319 const TilingContext tiling =
320 internal::GetTensorExecutorTilingContext<Evaluator, BlockMapper, Vectorizable>(evaluator);
321
322 auto eval_block = [&device, &evaluator, &tiling](IndexType firstBlockIdx, IndexType lastBlockIdx) {
323 EvalBlockRange<Evaluator, BlockMapper, IndexType>::run(&evaluator, tiling.block_mapper, device, firstBlockIdx,
324 lastBlockIdx);
325 };
326
327 // Evaluate small expressions directly as a single block.
328 if (tiling.block_mapper.blockCount() == 1) {
329 TensorBlockScratch scratch(device);
330 TensorBlockDesc desc(0, tiling.block_mapper.blockDimensions());
331 evaluator.evalBlock(desc, scratch);
332 } else {
333 device.parallelFor(tiling.block_mapper.blockCount(), tiling.cost, std::move(eval_block));
334 }
335 }
336 evaluator.cleanup();
337 }
338};
339
340template <typename Expression, typename DoneCallback, bool Vectorizable, TiledEvaluation Tiling>
341class TensorAsyncExecutor<Expression, ThreadPoolDevice, DoneCallback, Vectorizable, Tiling> {
342 public:
343 typedef typename Expression::Index StorageIndex;
344 typedef TensorEvaluator<Expression, ThreadPoolDevice> Evaluator;
345
346 static EIGEN_STRONG_INLINE void runAsync(const Expression& expr, const ThreadPoolDevice& device, DoneCallback done) {
347 TensorAsyncExecutorContext* const ctx = new TensorAsyncExecutorContext(expr, device, std::move(done));
348
349 const auto on_eval_subexprs = [ctx, &device](bool need_assign) -> void {
350 if (!need_assign) {
351 delete ctx;
352 return;
353 }
354
355 typedef EvalRange<Evaluator, StorageIndex, Vectorizable> EvalRange;
356 const StorageIndex size = static_cast<StorageIndex>(array_prod(ctx->evaluator.dimensions()));
357 device.parallelForAsync(
358 size, ctx->evaluator.costPerCoeff(Vectorizable), EvalRange::alignBlockSize,
359 [ctx](StorageIndex firstIdx, StorageIndex lastIdx) { EvalRange::run(&ctx->evaluator, firstIdx, lastIdx); },
360 [ctx]() { delete ctx; });
361 };
362
363 ctx->evaluator.evalSubExprsIfNeededAsync(nullptr, on_eval_subexprs);
364 }
365
366 private:
367 struct TensorAsyncExecutorContext {
368 TensorAsyncExecutorContext(const Expression& expr, const ThreadPoolDevice& thread_pool, DoneCallback done)
369 : evaluator(expr, thread_pool), on_done(std::move(done)) {}
370
371 ~TensorAsyncExecutorContext() {
372 evaluator.cleanup();
373 on_done();
374 }
375
376 Evaluator evaluator;
377
378 private:
379 DoneCallback on_done;
380 };
381};
382
383template <typename Expression, typename DoneCallback, bool Vectorizable>
384class TensorAsyncExecutor<Expression, ThreadPoolDevice, DoneCallback, Vectorizable, /*Tileable*/ TiledEvaluation::On> {
385 public:
386 typedef typename traits<Expression>::Index IndexType;
387 typedef typename traits<Expression>::Scalar Scalar;
388 typedef std::remove_const_t<Scalar> ScalarNoConst;
389
390 static constexpr int NumDims = traits<Expression>::NumDimensions;
391
392 typedef TensorEvaluator<Expression, ThreadPoolDevice> Evaluator;
393 typedef TensorBlockMapper<NumDims, Evaluator::Layout, IndexType> BlockMapper;
394 typedef TensorExecutorTilingContext<BlockMapper> TilingContext;
395
396 typedef internal::TensorBlockDescriptor<NumDims, IndexType> TensorBlockDesc;
397 typedef internal::TensorBlockScratchAllocator<ThreadPoolDevice> TensorBlockScratch;
398
399 static EIGEN_STRONG_INLINE void runAsync(const Expression& expr, const ThreadPoolDevice& device, DoneCallback done) {
400 TensorAsyncExecutorContext* const ctx = new TensorAsyncExecutorContext(expr, device, std::move(done));
401
402 const auto on_eval_subexprs = [ctx](bool need_assign) -> void {
403 if (!need_assign) {
404 delete ctx;
405 return;
406 }
407
408 ctx->tiling = internal::GetTensorExecutorTilingContext<Evaluator, BlockMapper, Vectorizable>(ctx->evaluator);
409
410 auto eval_block = [ctx](IndexType firstBlockIdx, IndexType lastBlockIdx) {
411 EvalBlockRange<Evaluator, BlockMapper, IndexType>::run(&ctx->evaluator, ctx->tiling.block_mapper, ctx->device,
412 firstBlockIdx, lastBlockIdx);
413 };
414
415 // Evaluate small expressions directly as a single block.
416 if (ctx->tiling.block_mapper.blockCount() == 1) {
417 TensorBlockScratch scratch(ctx->device);
418 TensorBlockDesc desc(0, ctx->tiling.block_mapper.blockDimensions());
419 ctx->evaluator.evalBlock(desc, scratch);
420 delete ctx;
421 } else {
422 ctx->device.parallelForAsync(ctx->tiling.block_mapper.blockCount(), ctx->tiling.cost, eval_block,
423 [ctx]() { delete ctx; });
424 }
425 };
426
427 ctx->evaluator.evalSubExprsIfNeededAsync(nullptr, on_eval_subexprs);
428 }
429
430 private:
431 struct TensorAsyncExecutorContext {
432 TensorAsyncExecutorContext(const Expression& expr, const ThreadPoolDevice& thread_pool, DoneCallback done)
433 : device(thread_pool), evaluator(expr, thread_pool), on_done(std::move(done)) {}
434
435 ~TensorAsyncExecutorContext() {
436 evaluator.cleanup();
437 on_done();
438 }
439
440 const ThreadPoolDevice& device;
441 Evaluator evaluator;
442 TilingContext tiling;
443
444 private:
445 DoneCallback on_done;
446 };
447};
448
449#endif // EIGEN_USE_THREADS
450
451// GPU: the evaluation of the expression is offloaded to a GPU.
452#if defined(EIGEN_USE_GPU)
453
454template <typename Expression, bool Vectorizable, TiledEvaluation Tiling>
455class TensorExecutor<Expression, GpuDevice, Vectorizable, Tiling> {
456 public:
457 typedef typename Expression::Index StorageIndex;
458 static void run(const Expression& expr, const GpuDevice& device);
459};
460
461#if defined(EIGEN_GPUCC)
462// Returns 1 if lhs + rhs would overflow, -1 if it would underflow, otherwise 0.
463template <typename Index>
464EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE int sum_will_overflow(Index lhs, Index rhs) {
465 const Index highest = NumTraits<Index>::highest();
466 const Index lowest = NumTraits<Index>::lowest();
467 if (lhs > 0 && rhs > 0) {
468 return lhs > highest - rhs ? 1 : 0;
469 } else if (lhs < 0 && rhs < 0) {
470 return lhs < lowest - rhs ? -1 : 0;
471 } else {
472 return 0;
473 }
474}
475
476// Returns lhs + rhs, saturating to the highest/lowest representable value on
477// overflow/underflow respectively.
478template <typename Index>
479EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Index saturate_add(Index lhs, Index rhs) {
480 const Index highest = NumTraits<Index>::highest();
481 const Index lowest = NumTraits<Index>::lowest();
482 int overflow = sum_will_overflow(lhs, rhs);
483 return overflow == 1 ? highest : overflow == -1 ? lowest : lhs + rhs;
484}
485
486// A functor that adds step_size to a given index, saturating to avoid
487// overflow/underflow. If overflow/underflow is not possible, regular addition
488// is used (for efficiency).
489template <typename Index>
490struct SafeStep {
491 // lastIdx is one past the end of the possible indexes.
492 // step_size is the value that will be added to the given index when the
493 // functor is called.
494 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE SafeStep(Index lastIdx, Index step_size)
495 : can_overflow_(sum_will_overflow(lastIdx, step_size)), step_size_(step_size) {}
496
497 // Adds step_size to index, saturating on overflow (if overflow is possible).
498 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Index operator()(Index index) const {
499 return can_overflow_ ? saturate_add(index, step_size_) : index + step_size_;
500 }
501
502 private:
503 const bool can_overflow_;
504 const Index step_size_;
505};
506
507template <typename Evaluator, typename StorageIndex, bool Vectorizable>
508struct EigenMetaKernelEval {
509 static EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE void run(Evaluator& eval, StorageIndex firstIdx, StorageIndex lastIdx,
510 StorageIndex step_size) {
511 SafeStep<StorageIndex> safe_step(lastIdx, step_size);
512 for (StorageIndex i = firstIdx; i < lastIdx; i = safe_step(i)) {
513 eval.evalScalar(i);
514 }
515 }
516};
517
518template <typename Evaluator, typename StorageIndex>
519struct EigenMetaKernelEval<Evaluator, StorageIndex, true> {
520 static EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE void run(Evaluator& eval, StorageIndex firstIdx, StorageIndex lastIdx,
521 StorageIndex step_size) {
522 const StorageIndex PacketSize = unpacket_traits<typename Evaluator::PacketReturnType>::size;
523 const StorageIndex vectorized_size = (lastIdx / PacketSize) * PacketSize;
524 const StorageIndex vectorized_step_size = step_size * PacketSize;
525
526 SafeStep<StorageIndex> safe_vectorized_step(vectorized_size, vectorized_step_size);
527 // Use the vector path
528 for (StorageIndex i = firstIdx * PacketSize; i < vectorized_size; i = safe_vectorized_step(i)) {
529 eval.evalPacket(i);
530 }
531 SafeStep<StorageIndex> safe_step(lastIdx, step_size);
532 for (StorageIndex i = saturate_add(vectorized_size, firstIdx); i < lastIdx; i = safe_step(i)) {
533 eval.evalScalar(i);
534 }
535 }
536};
537
538template <typename Evaluator, typename StorageIndex>
539__global__ void __launch_bounds__(1024) EigenMetaKernel(Evaluator eval, StorageIndex size) {
540 const StorageIndex first_index = blockIdx.x * blockDim.x + threadIdx.x;
541 const StorageIndex step_size = blockDim.x * gridDim.x;
542
543 const bool vectorizable = Evaluator::PacketAccess & Evaluator::IsAligned;
544 EigenMetaKernelEval<Evaluator, StorageIndex, vectorizable>::run(eval, first_index, size, step_size);
545}
546
547/*static*/
548template <typename Expression, bool Vectorizable, TiledEvaluation Tiling>
549EIGEN_STRONG_INLINE void TensorExecutor<Expression, GpuDevice, Vectorizable, Tiling>::run(const Expression& expr,
550 const GpuDevice& device) {
551 TensorEvaluator<Expression, GpuDevice> evaluator(expr, device);
552 const bool needs_assign = evaluator.evalSubExprsIfNeeded(nullptr);
553 if (needs_assign) {
554 const int block_size = device.maxGpuThreadsPerBlock();
555 const int max_blocks = static_cast<int>(
556 numext::mini<int64_t>(device.getNumGpuMultiProcessors() * device.maxGpuThreadsPerMultiProcessor(),
557 NumTraits<StorageIndex>::highest()) /
558 block_size);
559 const StorageIndex size = static_cast<StorageIndex>(array_prod(evaluator.dimensions()));
560 // Create at least one block to ensure we don't crash with tensors of size 0.
561 const int num_blocks = numext::maxi<int>(
562 numext::mini<int>(max_blocks, static_cast<int>(numext::div_ceil<StorageIndex>(size, block_size))), 1);
563
564 LAUNCH_GPU_KERNEL((EigenMetaKernel<TensorEvaluator<Expression, GpuDevice>, StorageIndex>), num_blocks, block_size,
565 0, device, evaluator, size);
566 }
567 evaluator.cleanup();
568}
569
570#endif // EIGEN_GPUCC
571#endif // EIGEN_USE_GPU
572
573// SYCL Executor policy
574#ifdef EIGEN_USE_SYCL
575
576template <typename Evaluator>
577struct ExecExprFunctorKernel {
578 typedef typename Evaluator::Index Index;
579 Evaluator evaluator;
580 const Index range;
581 template <typename Scratch>
582 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE ExecExprFunctorKernel(const Scratch, Evaluator evaluator_, const Index range_)
583 : evaluator(evaluator_), range(range_) {}
584
585 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE void operator()(cl::sycl::nd_item<1> itemID) const { compute(itemID); }
586 template <bool is_vec = Evaluator::PacketAccess>
587 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE std::enable_if_t<!is_vec> compute(const cl::sycl::nd_item<1>& itemID) const {
588 Index gId = static_cast<Index>(itemID.get_global_linear_id());
589 Index total_threads = itemID.get_global_range(0);
590
591 for (Index i = gId; i < range; i += total_threads) {
592 evaluator.evalScalar(i);
593 }
594 }
595 template <bool is_vec = Evaluator::PacketAccess>
596 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE std::enable_if_t<is_vec> compute(const cl::sycl::nd_item<1>& itemID) const {
597 const Index vectorizedRange = (range / Evaluator::PacketSize) * Evaluator::PacketSize;
598 Index gId = static_cast<Index>(itemID.get_global_linear_id());
599 const Index step = Evaluator::PacketSize * itemID.get_global_range(0);
600 const Index start = Evaluator::PacketSize * gId;
601 for (Index i = start; i < vectorizedRange; i += step) {
602 evaluator.evalPacket(i);
603 }
604 gId += vectorizedRange;
605 for (Index i = gId; i < range; i += itemID.get_global_range(0)) {
606 evaluator.evalScalar(i);
607 }
608 }
609};
610
611template <typename Expression, bool Vectorizable, TiledEvaluation Tiling>
612class TensorExecutor<Expression, Eigen::SyclDevice, Vectorizable, Tiling> {
613 public:
614 typedef typename Expression::Index Index;
615 static EIGEN_STRONG_INLINE void run(const Expression& expr, const Eigen::SyclDevice& dev) {
616 typedef Eigen::TensorEvaluator<Expression, Eigen::SyclDevice> Evaluator;
617 Evaluator evaluator(expr, dev);
618 const bool needs_assign = evaluator.evalSubExprsIfNeeded(nullptr);
619 if (needs_assign) {
620 Index range, GRange, tileSize;
621 Index total_size = ::Eigen::internal::array_prod(evaluator.dimensions());
622 total_size = (total_size == 0) ? 1 : total_size;
623 const int PacketSize = Eigen::PacketType<typename Evaluator::CoeffReturnType, Eigen::SyclDevice>::size;
624 Index vectorizable_threads = static_cast<Index>(total_size / PacketSize);
625 dev.parallel_for_setup(vectorizable_threads, tileSize, range, GRange);
626 range = total_size;
627
628 dev.template nullary_kernel_launcher<typename Evaluator::CoeffReturnType, ExecExprFunctorKernel<Evaluator> >(
629 evaluator, cl::sycl::nd_range<1>(cl::sycl::range<1>(GRange), cl::sycl::range<1>(tileSize)), Index(1),
630 range)
631 .wait();
632 }
633 evaluator.cleanup();
634 }
635};
636
637#endif
638
639} // end namespace internal
640
641} // end namespace Eigen
642
643#endif // EIGEN_TENSOR_TENSOR_EXECUTOR_H
Definition TensorExecutor.h:70
The tensor executor class.
Definition TensorExecutor.h:38
Namespace containing all symbols from the Eigen library.
The tensor evaluator class.
Definition TensorEvaluator.h:47