11#ifndef EIGEN_TENSOR_TENSOR_EXECUTOR_H
12#define EIGEN_TENSOR_TENSOR_EXECUTOR_H
15#include "./InternalHeaderCheck.h"
37template <
typename Expression,
typename Device,
bool Vectorizable, TiledEvaluation Tiling>
40 typedef typename Expression::Index StorageIndex;
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.");
52 static EIGEN_STRONG_INLINE
void run(
const Expression& expr,
const Device& device = DefaultDevice()) {
54 const bool needs_assign = evaluator.evalSubExprsIfNeeded(
nullptr);
56 const StorageIndex size =
static_cast<StorageIndex
>(array_prod(evaluator.dimensions()));
57 for (StorageIndex i = 0; i < size; ++i) {
58 evaluator.evalScalar(i);
69template <
typename Expression,
typename Device,
typename DoneCallback,
bool Vectorizable, TiledEvaluation Tiling>
75template <
typename Expression>
77 TiledEvaluation::Off> {
79 typedef typename Expression::Index StorageIndex;
81 static EIGEN_STRONG_INLINE
void run(
const Expression& expr,
const DefaultDevice& device = DefaultDevice()) {
83 const bool needs_assign = evaluator.evalSubExprsIfNeeded(
nullptr);
85 const StorageIndex size =
static_cast<StorageIndex
>(array_prod(evaluator.dimensions()));
86 const int PacketSize =
87 unpacket_traits<typename TensorEvaluator<Expression, DefaultDevice>::PacketReturnType>::size;
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);
98 const StorageIndex VectorizedSize = (size / PacketSize) * PacketSize;
99 for (StorageIndex i = UnrolledSize; i < VectorizedSize; i += PacketSize) {
100 evaluator.evalPacket(i);
102 for (StorageIndex i = VectorizedSize; i < size; ++i) {
103 evaluator.evalScalar(i);
114template <
typename Expression,
bool Vectorizable>
116 TiledEvaluation::On> {
118 typedef typename traits<Expression>::Scalar Scalar;
119 typedef std::remove_const_t<Scalar> ScalarNoConst;
122 typedef typename traits<Expression>::Index StorageIndex;
124 static constexpr int NumDims = traits<Expression>::NumDimensions;
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;
130 typedef internal::TensorBlockDescriptor<NumDims, StorageIndex> TensorBlockDesc;
131 typedef internal::TensorBlockScratchAllocator<DefaultDevice> TensorBlockScratch;
133 Evaluator evaluator(expr, device);
136 const bool needs_assign = evaluator.evalSubExprsIfNeeded(
nullptr);
140 const TensorBlockResourceRequirements requirements = evaluator.getResourceRequirements();
142 const TensorBlockMapper block_mapper(
typename TensorBlockDesc::Dimensions(evaluator.dimensions()), requirements);
145 TensorBlockScratch scratch(device);
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);
169#ifdef EIGEN_USE_THREADS
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) {}
177 TensorBlockMapper block_mapper;
179 size_t aligned_blocksize;
184template <
typename Evaluator,
typename TensorBlockMapper,
bool Vectorizable>
185TensorExecutorTilingContext<TensorBlockMapper> GetTensorExecutorTilingContext(
const Evaluator& evaluator) {
187 TensorBlockResourceRequirements requirements = evaluator.getResourceRequirements();
190 double taskSize = TensorCostModel<ThreadPoolDevice>::taskSize(1, requirements.cost_per_coeff);
191 requirements.size =
static_cast<size_t>(1.0 / taskSize);
193 TensorBlockMapper block_mapper(
typename TensorBlockMapper::Dimensions(evaluator.dimensions()), requirements);
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);
200 return {block_mapper, requirements.cost_per_coeff * block_size, aligned_blocksize};
203template <
typename Evaluator,
typename StorageIndex,
bool Vectorizable>
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);
213 static StorageIndex alignBlockSize(StorageIndex size) {
return size; }
216template <
typename Evaluator,
typename StorageIndex>
217struct EvalRange<Evaluator, StorageIndex, true> {
218 static constexpr int PacketSize = unpacket_traits<typename Evaluator::PacketReturnType>::size;
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;
230 for (; i <= last_chunk_offset; i += 4 * PacketSize) {
231 for (StorageIndex j = 0; j < 4; j++) {
232 evaluator.evalPacket(i + j * PacketSize);
235 last_chunk_offset = lastIdx - PacketSize;
236 for (; i <= last_chunk_offset; i += PacketSize) {
237 evaluator.evalPacket(i);
240 for (; i < lastIdx; ++i) {
241 evaluator.evalScalar(i);
245 static StorageIndex alignBlockSize(StorageIndex size) {
247 if (size >= 16 * PacketSize) {
248 return (size + 4 * PacketSize - 1) & ~(4 * PacketSize - 1);
251 return (size + PacketSize - 1) & ~(PacketSize - 1);
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);
268 for (IndexType block_idx = firstBlockIdx; block_idx < lastBlockIdx; ++block_idx) {
269 auto desc = block_mapper.blockDescriptor(block_idx);
270 evaluator.evalBlock(desc, scratch);
276template <
typename Expression,
bool Vectorizable, TiledEvaluation Tiling>
277class TensorExecutor<Expression, ThreadPoolDevice, Vectorizable, Tiling> {
279 typedef typename Expression::Index StorageIndex;
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;
285 Evaluator evaluator(expr, device);
286 const bool needs_assign = evaluator.evalSubExprsIfNeeded(
nullptr);
288 const StorageIndex size =
static_cast<StorageIndex
>(array_prod(evaluator.dimensions()));
290 size, evaluator.costPerCoeff(Vectorizable), EvalRange::alignBlockSize,
291 [&evaluator](StorageIndex firstIdx, StorageIndex lastIdx) { EvalRange::run(&evaluator, firstIdx, lastIdx); });
297template <
typename Expression,
bool Vectorizable>
299 TiledEvaluation::On> {
301 typedef typename traits<Expression>::Index IndexType;
302 typedef typename traits<Expression>::Scalar Scalar;
303 typedef std::remove_const_t<Scalar> ScalarNoConst;
305 static constexpr int NumDims = traits<Expression>::NumDimensions;
307 typedef TensorEvaluator<Expression, ThreadPoolDevice> Evaluator;
308 typedef TensorBlockMapper<NumDims, Evaluator::Layout, IndexType> BlockMapper;
309 typedef TensorExecutorTilingContext<BlockMapper> TilingContext;
311 typedef internal::TensorBlockDescriptor<NumDims, IndexType> TensorBlockDesc;
312 typedef internal::TensorBlockScratchAllocator<ThreadPoolDevice> TensorBlockScratch;
314 static EIGEN_STRONG_INLINE
void run(
const Expression& expr,
const ThreadPoolDevice& device) {
315 Evaluator evaluator(expr, device);
317 const bool needs_assign = evaluator.evalSubExprsIfNeeded(
nullptr);
319 const TilingContext tiling =
320 internal::GetTensorExecutorTilingContext<Evaluator, BlockMapper, Vectorizable>(evaluator);
322 auto eval_block = [&device, &evaluator, &tiling](IndexType firstBlockIdx, IndexType lastBlockIdx) {
323 EvalBlockRange<Evaluator, BlockMapper, IndexType>::run(&evaluator, tiling.block_mapper, device, firstBlockIdx,
328 if (tiling.block_mapper.blockCount() == 1) {
329 TensorBlockScratch scratch(device);
330 TensorBlockDesc desc(0, tiling.block_mapper.blockDimensions());
331 evaluator.evalBlock(desc, scratch);
333 device.parallelFor(tiling.block_mapper.blockCount(), tiling.cost, std::move(eval_block));
340template <
typename Expression,
typename DoneCallback,
bool Vectorizable, TiledEvaluation Tiling>
341class TensorAsyncExecutor<Expression, ThreadPoolDevice, DoneCallback, Vectorizable, Tiling> {
343 typedef typename Expression::Index StorageIndex;
344 typedef TensorEvaluator<Expression, ThreadPoolDevice> Evaluator;
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));
349 const auto on_eval_subexprs = [ctx, &device](
bool need_assign) ->
void {
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; });
363 ctx->evaluator.evalSubExprsIfNeededAsync(
nullptr, on_eval_subexprs);
367 struct TensorAsyncExecutorContext {
368 TensorAsyncExecutorContext(
const Expression& expr,
const ThreadPoolDevice& thread_pool, DoneCallback done)
369 : evaluator(expr, thread_pool), on_done(std::move(done)) {}
371 ~TensorAsyncExecutorContext() {
379 DoneCallback on_done;
383template <
typename Expression,
typename DoneCallback,
bool Vectorizable>
384class TensorAsyncExecutor<Expression, ThreadPoolDevice, DoneCallback, Vectorizable, TiledEvaluation::On> {
386 typedef typename traits<Expression>::Index IndexType;
387 typedef typename traits<Expression>::Scalar Scalar;
388 typedef std::remove_const_t<Scalar> ScalarNoConst;
390 static constexpr int NumDims = traits<Expression>::NumDimensions;
392 typedef TensorEvaluator<Expression, ThreadPoolDevice> Evaluator;
393 typedef TensorBlockMapper<NumDims, Evaluator::Layout, IndexType> BlockMapper;
394 typedef TensorExecutorTilingContext<BlockMapper> TilingContext;
396 typedef internal::TensorBlockDescriptor<NumDims, IndexType> TensorBlockDesc;
397 typedef internal::TensorBlockScratchAllocator<ThreadPoolDevice> TensorBlockScratch;
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));
402 const auto on_eval_subexprs = [ctx](
bool need_assign) ->
void {
408 ctx->tiling = internal::GetTensorExecutorTilingContext<Evaluator, BlockMapper, Vectorizable>(ctx->evaluator);
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);
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);
422 ctx->device.parallelForAsync(ctx->tiling.block_mapper.blockCount(), ctx->tiling.cost, eval_block,
423 [ctx]() { delete ctx; });
427 ctx->evaluator.evalSubExprsIfNeededAsync(
nullptr, on_eval_subexprs);
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)) {}
435 ~TensorAsyncExecutorContext() {
440 const ThreadPoolDevice& device;
442 TilingContext tiling;
445 DoneCallback on_done;
452#if defined(EIGEN_USE_GPU)
454template <
typename Expression,
bool Vectorizable, TiledEvaluation Tiling>
455class TensorExecutor<Expression, GpuDevice, Vectorizable, Tiling> {
457 typedef typename Expression::Index StorageIndex;
458 static void run(
const Expression& expr,
const GpuDevice& device);
461#if defined(EIGEN_GPUCC)
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;
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;
489template <
typename Index>
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) {}
498 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Index operator()(Index index)
const {
499 return can_overflow_ ? saturate_add(index, step_size_) : index + step_size_;
503 const bool can_overflow_;
504 const Index step_size_;
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)) {
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;
526 SafeStep<StorageIndex> safe_vectorized_step(vectorized_size, vectorized_step_size);
528 for (StorageIndex i = firstIdx * PacketSize; i < vectorized_size; i = safe_vectorized_step(i)) {
531 SafeStep<StorageIndex> safe_step(lastIdx, step_size);
532 for (StorageIndex i = saturate_add(vectorized_size, firstIdx); i < lastIdx; i = safe_step(i)) {
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;
543 const bool vectorizable = Evaluator::PacketAccess & Evaluator::IsAligned;
544 EigenMetaKernelEval<Evaluator, StorageIndex, vectorizable>::run(eval, first_index, size, step_size);
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);
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()) /
559 const StorageIndex size =
static_cast<StorageIndex
>(array_prod(evaluator.dimensions()));
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);
564 LAUNCH_GPU_KERNEL((EigenMetaKernel<TensorEvaluator<Expression, GpuDevice>, StorageIndex>), num_blocks, block_size,
565 0, device, evaluator, size);
576template <
typename Evaluator>
577struct ExecExprFunctorKernel {
578 typedef typename Evaluator::Index Index;
581 template <
typename Scratch>
582 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE ExecExprFunctorKernel(
const Scratch, Evaluator evaluator_,
const Index range_)
583 : evaluator(evaluator_), range(range_) {}
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);
591 for (Index i = gId; i < range; i += total_threads) {
592 evaluator.evalScalar(i);
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);
604 gId += vectorizedRange;
605 for (Index i = gId; i < range; i += itemID.get_global_range(0)) {
606 evaluator.evalScalar(i);
611template <
typename Expression,
bool Vectorizable, TiledEvaluation Tiling>
612class TensorExecutor<Expression, Eigen::SyclDevice, Vectorizable, Tiling> {
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);
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);
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),
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