From 33e174613987cfc6c83576dc0fe8086c7a5d1b1f Mon Sep 17 00:00:00 2001
From: Eugene Zhulenev <ezhulenev@google.com>
Date: Wed, 9 Oct 2019 12:45:31 -0700
Subject: Block evaluation for TensorChipping + fixed bugs in TensorPadding and
 TensorSlicing

---
 unsupported/Eigen/CXX11/src/Tensor/TensorExecutor.h | 13 +++++++++++++
 1 file changed, 13 insertions(+)

(limited to 'unsupported/Eigen/CXX11/src/Tensor/TensorExecutor.h')

diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorExecutor.h b/unsupported/Eigen/CXX11/src/Tensor/TensorExecutor.h
index 97ac96db1..6ad6327a6 100644
--- a/unsupported/Eigen/CXX11/src/Tensor/TensorExecutor.h
+++ b/unsupported/Eigen/CXX11/src/Tensor/TensorExecutor.h
@@ -521,6 +521,19 @@ class TensorExecutor<Expression, ThreadPoolDevice, Vectorizable,
   static EIGEN_STRONG_INLINE void run(const Expression& expr,
                                       const ThreadPoolDevice& device) {
     Evaluator evaluator(expr, device);
+    Index total_size = array_prod(evaluator.dimensions());
+    Index cache_size = device.firstLevelCacheSize() / sizeof(Scalar);
+
+    // TODO(ezuhulenev): For small expressions cost of block mapping and
+    // resource requirements gathering dominates the cost of expression
+    // evaluatiuon.
+    if (total_size < cache_size &&
+        !ExpressionHasTensorBroadcastingOp<Expression>::value) {
+      internal::TensorExecutor<Expression, ThreadPoolDevice, Vectorizable,
+                               /*Tiling=*/TiledEvaluation::Off>::run(expr, device);
+      evaluator.cleanup();
+      return;
+    }
 
     const bool needs_assign = evaluator.evalSubExprsIfNeeded(nullptr);
     if (needs_assign) {
-- 
cgit v1.2.3