define XTENSOR_TBB_THRESHOLD

use tbb threshold in assignment
masmirnoff · Nov 9, 2021 · 11851f9 · 11851f9
1 parent 26e3cc5
commit 11851f9
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 3 deletions.
diff --git a/include/xtensor/xassign.hpp b/include/xtensor/xassign.hpp
@@ -641,10 +641,20 @@ namespace xt
         }
 
 #if defined(XTENSOR_USE_TBB)
-        tbb::parallel_for(align_begin, align_end, simd_size, [&e1, &e2](size_t i)
+        if (size >= XTENSOR_TBB_THRESHOLD)
         {
-            e1.template store_simd<lhs_align_mode>(i, e2.template load_simd<rhs_align_mode, value_type>(i));
-        });
+            tbb::parallel_for(align_begin, align_end, simd_size, [&e1, &e2](size_t i)
+            {
+                e1.template store_simd<lhs_align_mode>(i, e2.template load_simd<rhs_align_mode, value_type>(i));
+            });
+        }
+        else
+        {
+            for (size_type i = align_begin; i < align_end; i += simd_size)
+            {
+                e1.template store_simd<lhs_align_mode>(i, e2.template load_simd<rhs_align_mode, value_type>(i));
+            }
+        }
 #elif defined(XTENSOR_USE_OPENMP)
         if (size >= XTENSOR_OPENMP_TRESHOLD)
         {

diff --git a/include/xtensor/xtensor_config.hpp b/include/xtensor/xtensor_config.hpp
@@ -108,6 +108,10 @@
 #define XTENSOR_OPENMP_TRESHOLD 0
 #endif
 
+#ifndef XTENSOR_TBB_THRESHOLD
+#define XTENSOR_TBB_THRESHOLD 0
+#endif
+
 #ifndef XTENSOR_SELECT_ALIGN
 #define XTENSOR_SELECT_ALIGN(T) (XTENSOR_DEFAULT_ALIGNMENT != 0 ? XTENSOR_DEFAULT_ALIGNMENT : alignof(T))
 #endif