@smake/eigen 1.0.2 → 1.1.1

package/eigen/Eigen/src/Core/products/GeneralMatrixVector_BLAS.h

@@ -33,104 +33,107 @@
 #ifndef EIGEN_GENERAL_MATRIX_VECTOR_BLAS_H
 #define EIGEN_GENERAL_MATRIX_VECTOR_BLAS_H
 
-namespace Eigen { 
+// IWYU pragma: private
+#include "../InternalHeaderCheck.h"
+
+namespace Eigen {
 
 namespace internal {
 
 /**********************************************************************
-* This file implements general matrix-vector multiplication using BLAS
-* gemv function via partial specialization of
-* general_matrix_vector_product::run(..) method for float, double,
-* std::complex<float> and std::complex<double> types
-**********************************************************************/
+ * This file implements general matrix-vector multiplication using BLAS
+ * gemv function via partial specialization of
+ * general_matrix_vector_product::run(..) method for float, double,
+ * std::complex<float> and std::complex<double> types
+ **********************************************************************/
 
 // gemv specialization
 
-template<typename Index, typename LhsScalar, int StorageOrder, bool ConjugateLhs, typename RhsScalar, bool ConjugateRhs>
+template <typename Index, typename LhsScalar, int StorageOrder, bool ConjugateLhs, typename RhsScalar,
+          bool ConjugateRhs>
 struct general_matrix_vector_product_gemv;
 
-#define EIGEN_BLAS_GEMV_SPECIALIZE(Scalar) \
-template<typename Index, bool ConjugateLhs, bool ConjugateRhs> \
-struct general_matrix_vector_product<Index,Scalar,const_blas_data_mapper<Scalar,Index,ColMajor>,ColMajor,ConjugateLhs,Scalar,const_blas_data_mapper<Scalar,Index,RowMajor>,ConjugateRhs,Specialized> { \
-static void run( \
-  Index rows, Index cols, \
-  const const_blas_data_mapper<Scalar,Index,ColMajor> &lhs, \
-  const const_blas_data_mapper<Scalar,Index,RowMajor> &rhs, \
-  Scalar* res, Index resIncr, Scalar alpha) \
-{ \
-  if (ConjugateLhs) { \
-    general_matrix_vector_product<Index,Scalar,const_blas_data_mapper<Scalar,Index,ColMajor>,ColMajor,ConjugateLhs,Scalar,const_blas_data_mapper<Scalar,Index,RowMajor>,ConjugateRhs,BuiltIn>::run( \
-      rows, cols, lhs, rhs, res, resIncr, alpha); \
-  } else { \
-    general_matrix_vector_product_gemv<Index,Scalar,ColMajor,ConjugateLhs,Scalar,ConjugateRhs>::run( \
-      rows, cols, lhs.data(), lhs.stride(), rhs.data(), rhs.stride(), res, resIncr, alpha); \
-  } \
-} \
-}; \
-template<typename Index, bool ConjugateLhs, bool ConjugateRhs> \
-struct general_matrix_vector_product<Index,Scalar,const_blas_data_mapper<Scalar,Index,RowMajor>,RowMajor,ConjugateLhs,Scalar,const_blas_data_mapper<Scalar,Index,ColMajor>,ConjugateRhs,Specialized> { \
-static void run( \
-  Index rows, Index cols, \
-  const const_blas_data_mapper<Scalar,Index,RowMajor> &lhs, \
-  const const_blas_data_mapper<Scalar,Index,ColMajor> &rhs, \
-  Scalar* res, Index resIncr, Scalar alpha) \
-{ \
-    general_matrix_vector_product_gemv<Index,Scalar,RowMajor,ConjugateLhs,Scalar,ConjugateRhs>::run( \
-      rows, cols, lhs.data(), lhs.stride(), rhs.data(), rhs.stride(), res, resIncr, alpha); \
-} \
-}; \
+#define EIGEN_BLAS_GEMV_SPECIALIZE(Scalar)                                                                       \
+  template <typename Index, bool ConjugateLhs, bool ConjugateRhs>                                                \
+  struct general_matrix_vector_product<Index, Scalar, const_blas_data_mapper<Scalar, Index, ColMajor>, ColMajor, \
+                                       ConjugateLhs, Scalar, const_blas_data_mapper<Scalar, Index, RowMajor>,    \
+                                       ConjugateRhs, Specialized> {                                              \
+    static void run(Index rows, Index cols, const const_blas_data_mapper<Scalar, Index, ColMajor>& lhs,          \
+                    const const_blas_data_mapper<Scalar, Index, RowMajor>& rhs, Scalar* res, Index resIncr,      \
+                    Scalar alpha) {                                                                              \
+      if (ConjugateLhs) {                                                                                        \
+        general_matrix_vector_product<Index, Scalar, const_blas_data_mapper<Scalar, Index, ColMajor>, ColMajor,  \
+                                      ConjugateLhs, Scalar, const_blas_data_mapper<Scalar, Index, RowMajor>,     \
+                                      ConjugateRhs, BuiltIn>::run(rows, cols, lhs, rhs, res, resIncr, alpha);    \
+      } else {                                                                                                   \
+        general_matrix_vector_product_gemv<Index, Scalar, ColMajor, ConjugateLhs, Scalar, ConjugateRhs>::run(    \
+            rows, cols, lhs.data(), lhs.stride(), rhs.data(), rhs.stride(), res, resIncr, alpha);                \
+      }                                                                                                          \
+    }                                                                                                            \
+  };                                                                                                             \
+  template <typename Index, bool ConjugateLhs, bool ConjugateRhs>                                                \
+  struct general_matrix_vector_product<Index, Scalar, const_blas_data_mapper<Scalar, Index, RowMajor>, RowMajor, \
+                                       ConjugateLhs, Scalar, const_blas_data_mapper<Scalar, Index, ColMajor>,    \
+                                       ConjugateRhs, Specialized> {                                              \
+    static void run(Index rows, Index cols, const const_blas_data_mapper<Scalar, Index, RowMajor>& lhs,          \
+                    const const_blas_data_mapper<Scalar, Index, ColMajor>& rhs, Scalar* res, Index resIncr,      \
+                    Scalar alpha) {                                                                              \
+      general_matrix_vector_product_gemv<Index, Scalar, RowMajor, ConjugateLhs, Scalar, ConjugateRhs>::run(      \
+          rows, cols, lhs.data(), lhs.stride(), rhs.data(), rhs.stride(), res, resIncr, alpha);                  \
+    }                                                                                                            \
+  };
 
 EIGEN_BLAS_GEMV_SPECIALIZE(double)
 EIGEN_BLAS_GEMV_SPECIALIZE(float)
 EIGEN_BLAS_GEMV_SPECIALIZE(dcomplex)
 EIGEN_BLAS_GEMV_SPECIALIZE(scomplex)
 
-#define EIGEN_BLAS_GEMV_SPECIALIZATION(EIGTYPE,BLASTYPE,BLASFUNC) \
-template<typename Index, int LhsStorageOrder, bool ConjugateLhs, bool ConjugateRhs> \
-struct general_matrix_vector_product_gemv<Index,EIGTYPE,LhsStorageOrder,ConjugateLhs,EIGTYPE,ConjugateRhs> \
-{ \
-typedef Matrix<EIGTYPE,Dynamic,1,ColMajor> GEMVVector;\
-\
-static void run( \
-  Index rows, Index cols, \
-  const EIGTYPE* lhs, Index lhsStride, \
-  const EIGTYPE* rhs, Index rhsIncr, \
-  EIGTYPE* res, Index resIncr, EIGTYPE alpha) \
-{ \
-  BlasIndex m=convert_index<BlasIndex>(rows), n=convert_index<BlasIndex>(cols), \
-            lda=convert_index<BlasIndex>(lhsStride), incx=convert_index<BlasIndex>(rhsIncr), incy=convert_index<BlasIndex>(resIncr); \
-  const EIGTYPE beta(1); \
-  const EIGTYPE *x_ptr; \
-  char trans=(LhsStorageOrder==ColMajor) ? 'N' : (ConjugateLhs) ? 'C' : 'T'; \
-  if (LhsStorageOrder==RowMajor) { \
-    m = convert_index<BlasIndex>(cols); \
-    n = convert_index<BlasIndex>(rows); \
-  }\
-  GEMVVector x_tmp; \
-  if (ConjugateRhs) { \
-    Map<const GEMVVector, 0, InnerStride<> > map_x(rhs,cols,1,InnerStride<>(incx)); \
-    x_tmp=map_x.conjugate(); \
-    x_ptr=x_tmp.data(); \
-    incx=1; \
-  } else x_ptr=rhs; \
-  BLASFUNC(&trans, &m, &n, (const BLASTYPE*)&numext::real_ref(alpha), (const BLASTYPE*)lhs, &lda, (const BLASTYPE*)x_ptr, &incx, (const BLASTYPE*)&numext::real_ref(beta), (BLASTYPE*)res, &incy); \
-}\
-};
+#define EIGEN_BLAS_GEMV_SPECIALIZATION(EIGTYPE, BLASTYPE, BLASFUNC)                                                 \
+  template <typename Index, int LhsStorageOrder, bool ConjugateLhs, bool ConjugateRhs>                              \
+  struct general_matrix_vector_product_gemv<Index, EIGTYPE, LhsStorageOrder, ConjugateLhs, EIGTYPE, ConjugateRhs> { \
+    typedef Matrix<EIGTYPE, Dynamic, 1, ColMajor> GEMVVector;                                                       \
+                                                                                                                    \
+    static void run(Index rows, Index cols, const EIGTYPE* lhs, Index lhsStride, const EIGTYPE* rhs, Index rhsIncr, \
+                    EIGTYPE* res, Index resIncr, EIGTYPE alpha) {                                                   \
+      if (rows == 0 || cols == 0) return;                                                                           \
+      BlasIndex m = convert_index<BlasIndex>(rows), n = convert_index<BlasIndex>(cols),                             \
+                lda = convert_index<BlasIndex>(lhsStride), incx = convert_index<BlasIndex>(rhsIncr),                \
+                incy = convert_index<BlasIndex>(resIncr);                                                           \
+      const EIGTYPE beta(1);                                                                                        \
+      const EIGTYPE* x_ptr;                                                                                         \
+      char trans = (LhsStorageOrder == ColMajor) ? 'N' : (ConjugateLhs) ? 'C' : 'T';                                \
+      if (LhsStorageOrder == RowMajor) {                                                                            \
+        m = convert_index<BlasIndex>(cols);                                                                         \
+        n = convert_index<BlasIndex>(rows);                                                                         \
+      }                                                                                                             \
+      GEMVVector x_tmp;                                                                                             \
+      if (ConjugateRhs) {                                                                                           \
+        Map<const GEMVVector, 0, InnerStride<> > map_x(rhs, cols, 1, InnerStride<>(incx));                          \
+        x_tmp = map_x.conjugate();                                                                                  \
+        x_ptr = x_tmp.data();                                                                                       \
+        incx = 1;                                                                                                   \
+      } else {                                                                                                      \
+        x_ptr = rhs;                                                                                                \
+      }                                                                                                             \
+      BLASFUNC(&trans, &m, &n, (const BLASTYPE*)&numext::real_ref(alpha), (const BLASTYPE*)lhs, &lda,               \
+               (const BLASTYPE*)x_ptr, &incx, (const BLASTYPE*)&numext::real_ref(beta), (BLASTYPE*)res, &incy);     \
+    }                                                                                                               \
+  };
 
 #ifdef EIGEN_USE_MKL
-EIGEN_BLAS_GEMV_SPECIALIZATION(double,   double, dgemv)
-EIGEN_BLAS_GEMV_SPECIALIZATION(float,    float,  sgemv)
+EIGEN_BLAS_GEMV_SPECIALIZATION(double, double, dgemv)
+EIGEN_BLAS_GEMV_SPECIALIZATION(float, float, sgemv)
 EIGEN_BLAS_GEMV_SPECIALIZATION(dcomplex, MKL_Complex16, zgemv)
-EIGEN_BLAS_GEMV_SPECIALIZATION(scomplex, MKL_Complex8 , cgemv)
+EIGEN_BLAS_GEMV_SPECIALIZATION(scomplex, MKL_Complex8, cgemv)
 #else
-EIGEN_BLAS_GEMV_SPECIALIZATION(double,   double, dgemv_)
-EIGEN_BLAS_GEMV_SPECIALIZATION(float,    float,  sgemv_)
+EIGEN_BLAS_GEMV_SPECIALIZATION(double, double, dgemv_)
+EIGEN_BLAS_GEMV_SPECIALIZATION(float, float, sgemv_)
 EIGEN_BLAS_GEMV_SPECIALIZATION(dcomplex, double, zgemv_)
-EIGEN_BLAS_GEMV_SPECIALIZATION(scomplex, float,  cgemv_)
+EIGEN_BLAS_GEMV_SPECIALIZATION(scomplex, float, cgemv_)
 #endif
 
-} // end namespase internal
+}  // namespace internal
 
-} // end namespace Eigen
+}  // end namespace Eigen
 
-#endif // EIGEN_GENERAL_MATRIX_VECTOR_BLAS_H
+#endif  // EIGEN_GENERAL_MATRIX_VECTOR_BLAS_H

package/eigen/Eigen/src/Core/products/Parallelizer.h

@@ -10,157 +10,273 @@
 #ifndef EIGEN_PARALLELIZER_H
 #define EIGEN_PARALLELIZER_H
 
-namespace Eigen {
+// IWYU pragma: private
+#include "../InternalHeaderCheck.h"
 
-namespace internal {
+// Note that in the following, there are 3 different uses of the concept
+// "number of threads":
+//  1. Max number of threads used by OpenMP or ThreadPool.
+//     * For OpenMP this is typically the value set by the OMP_NUM_THREADS
+//       environment variable, or by a call to omp_set_num_threads() prior to
+//       calling Eigen.
+//     * For ThreadPool, this is the number of threads in the ThreadPool.
+//  2. Max number of threads currently allowed to be used by parallel Eigen
+//     operations. This is set by setNbThreads(), and cannot exceed the value
+//     in 1.
+//  3. The actual number of threads used for a given parallel Eigen operation.
+//     This is typically computed on the fly using a cost model and cannot exceed
+//     the value in 2.
+//     * For OpenMP, this is typically the number of threads specified in individual
+//       "omp parallel" pragmas associated with an Eigen operation.
+//     * For ThreadPool, it is the number of concurrent tasks scheduled in the
+//       threadpool for a given Eigen operation. Notice that since the threadpool
+//       uses task stealing, there is no way to limit the number of concurrently
+//       executing tasks to below the number in 1. except by limiting the total
+//       number of tasks in flight.
 
-/** \internal */
-inline void manage_multi_threading(Action action, int* v)
-{
-  static int m_maxThreads = -1;
-  EIGEN_UNUSED_VARIABLE(m_maxThreads);
+#if defined(EIGEN_HAS_OPENMP) && defined(EIGEN_GEMM_THREADPOOL)
+#error "EIGEN_HAS_OPENMP and EIGEN_GEMM_THREADPOOL may not both be defined."
+#endif
 
-  if(action==SetAction)
-  {
-    eigen_internal_assert(v!=0);
-    m_maxThreads = *v;
-  }
-  else if(action==GetAction)
-  {
-    eigen_internal_assert(v!=0);
-    #ifdef EIGEN_HAS_OPENMP
-    if(m_maxThreads>0)
-      *v = m_maxThreads;
-    else
-      *v = omp_get_max_threads();
-    #else
-    *v = 1;
-    #endif
-  }
-  else
-  {
-    eigen_internal_assert(false);
-  }
-}
+namespace Eigen {
 
+namespace internal {
+inline void manage_multi_threading(Action action, int* v);
 }
 
+// Public APIs.
+
 /** Must be call first when calling Eigen from multiple threads */
-inline void initParallel()
-{
-  int nbt;
-  internal::manage_multi_threading(GetAction, &nbt);
-  std::ptrdiff_t l1, l2, l3;
-  internal::manage_caching_sizes(GetAction, &l1, &l2, &l3);
-}
+EIGEN_DEPRECATED_WITH_REASON("Initialization is no longer needed.") inline void initParallel() {}
 
 /** \returns the max number of threads reserved for Eigen
-  * \sa setNbThreads */
-inline int nbThreads()
-{
+ * \sa setNbThreads */
+inline int nbThreads() {
   int ret;
   internal::manage_multi_threading(GetAction, &ret);
   return ret;
 }
 
 /** Sets the max number of threads reserved for Eigen
-  * \sa nbThreads */
-inline void setNbThreads(int v)
-{
-  internal::manage_multi_threading(SetAction, &v);
+ * \sa nbThreads */
+inline void setNbThreads(int v) { internal::manage_multi_threading(SetAction, &v); }
+
+#ifdef EIGEN_GEMM_THREADPOOL
+// Sets the ThreadPool used by Eigen parallel Gemm.
+//
+// NOTICE: This function has a known race condition with
+// parallelize_gemm below, and should not be called while
+// an instance of that function is running.
+//
+// TODO(rmlarsen): Make the device API available instead of
+// storing a local static pointer variable to avoid this issue.
+inline ThreadPool* setGemmThreadPool(ThreadPool* new_pool) {
+  static ThreadPool* pool = nullptr;
+  if (new_pool != nullptr) {
+    // This will wait for work in all threads in *pool to finish,
+    // then destroy the old ThreadPool, and then replace it with new_pool.
+    pool = new_pool;
+    // Reset the number of threads to the number of threads on the new pool.
+    setNbThreads(pool->NumThreads());
+  }
+  return pool;
 }
 
+// Gets the ThreadPool used by Eigen parallel Gemm.
+inline ThreadPool* getGemmThreadPool() { return setGemmThreadPool(nullptr); }
+#endif
+
 namespace internal {
 
-template<typename Index> struct GemmParallelInfo
-{
-  GemmParallelInfo() : sync(-1), users(0), lhs_start(0), lhs_length(0) {}
+// Implementation.
+
+#if defined(EIGEN_USE_BLAS) || (!defined(EIGEN_HAS_OPENMP) && !defined(EIGEN_GEMM_THREADPOOL))
+
+inline void manage_multi_threading(Action action, int* v) {
+  if (action == SetAction) {
+    eigen_internal_assert(v != nullptr);
+  } else if (action == GetAction) {
+    eigen_internal_assert(v != nullptr);
+    *v = 1;
+  } else {
+    eigen_internal_assert(false);
+  }
+}
+template <typename Index>
+struct GemmParallelInfo {};
+template <bool Condition, typename Functor, typename Index>
+EIGEN_STRONG_INLINE void parallelize_gemm(const Functor& func, Index rows, Index cols, Index /*unused*/,
+                                          bool /*unused*/) {
+  func(0, rows, 0, cols);
+}
 
-  Index volatile sync;
-  int volatile users;
+#else
 
+template <typename Index>
+struct GemmParallelTaskInfo {
+  GemmParallelTaskInfo() : sync(-1), users(0), lhs_start(0), lhs_length(0) {}
+  std::atomic<Index> sync;
+  std::atomic<int> users;
   Index lhs_start;
   Index lhs_length;
 };
 
-template<bool Condition, typename Functor, typename Index>
-void parallelize_gemm(const Functor& func, Index rows, Index cols, Index depth, bool transpose)
-{
-  // TODO when EIGEN_USE_BLAS is defined,
-  // we should still enable OMP for other scalar types
-#if !(defined (EIGEN_HAS_OPENMP)) || defined (EIGEN_USE_BLAS)
-  // FIXME the transpose variable is only needed to properly split
-  // the matrix product when multithreading is enabled. This is a temporary
-  // fix to support row-major destination matrices. This whole
-  // parallelizer mechanism has to be redisigned anyway.
-  EIGEN_UNUSED_VARIABLE(depth);
-  EIGEN_UNUSED_VARIABLE(transpose);
-  func(0,rows, 0,cols);
+template <typename Index>
+struct GemmParallelInfo {
+  const int logical_thread_id;
+  const int num_threads;
+  GemmParallelTaskInfo<Index>* task_info;
+
+  GemmParallelInfo(int logical_thread_id_, int num_threads_, GemmParallelTaskInfo<Index>* task_info_)
+      : logical_thread_id(logical_thread_id_), num_threads(num_threads_), task_info(task_info_) {}
+};
+
+inline void manage_multi_threading(Action action, int* v) {
+  static int m_maxThreads = -1;
+  if (action == SetAction) {
+    eigen_internal_assert(v != nullptr);
+#if defined(EIGEN_HAS_OPENMP)
+    // Calling action == SetAction and *v = 0 means
+    // restoring m_maxThreads to the maximum number of threads specified
+    // for OpenMP.
+    eigen_internal_assert(*v >= 0);
+    int omp_threads = omp_get_max_threads();
+    m_maxThreads = (*v == 0 ? omp_threads : std::min(*v, omp_threads));
+#elif defined(EIGEN_GEMM_THREADPOOL)
+    // Calling action == SetAction and *v = 0 means
+    // restoring m_maxThreads to the number of threads in the ThreadPool,
+    // which defaults to 1 if no pool was provided.
+    eigen_internal_assert(*v >= 0);
+    ThreadPool* pool = getGemmThreadPool();
+    int pool_threads = pool != nullptr ? pool->NumThreads() : 1;
+    m_maxThreads = (*v == 0 ? pool_threads : numext::mini(pool_threads, *v));
+#endif
+  } else if (action == GetAction) {
+    eigen_internal_assert(v != nullptr);
+#if defined(EIGEN_HAS_OPENMP)
+    if (m_maxThreads > 0)
+      *v = m_maxThreads;
+    else
+      *v = omp_get_max_threads();
 #else
+    *v = m_maxThreads;
+#endif
+  } else {
+    eigen_internal_assert(false);
+  }
+}
 
-  // Dynamically check whether we should enable or disable OpenMP.
+template <bool Condition, typename Functor, typename Index>
+EIGEN_STRONG_INLINE void parallelize_gemm(const Functor& func, Index rows, Index cols, Index depth, bool transpose) {
+  // Dynamically check whether we should even try to execute in parallel.
   // The conditions are:
   // - the max number of threads we can create is greater than 1
   // - we are not already in a parallel code
   // - the sizes are large enough
 
   // compute the maximal number of threads from the size of the product:
-  // This first heuristic takes into account that the product kernel is fully optimized when working with nr columns at once.
+  // This first heuristic takes into account that the product kernel is fully optimized when working with nr columns at
+  // once.
   Index size = transpose ? rows : cols;
-  Index pb_max_threads = std::max<Index>(1,size / Functor::Traits::nr);
+  Index pb_max_threads = std::max<Index>(1, size / Functor::Traits::nr);
 
   // compute the maximal number of threads from the total amount of work:
-  double work = static_cast<double>(rows) * static_cast<double>(cols) *
-      static_cast<double>(depth);
+  double work = static_cast<double>(rows) * static_cast<double>(cols) * static_cast<double>(depth);
   double kMinTaskSize = 50000;  // FIXME improve this heuristic.
-  pb_max_threads = std::max<Index>(1, std::min<Index>(pb_max_threads, work / kMinTaskSize));
+  pb_max_threads = std::max<Index>(1, std::min<Index>(pb_max_threads, static_cast<Index>(work / kMinTaskSize)));
 
   // compute the number of threads we are going to use
-  Index threads = std::min<Index>(nbThreads(), pb_max_threads);
+  int threads = std::min<int>(nbThreads(), static_cast<int>(pb_max_threads));
 
-  // if multi-threading is explicitely disabled, not useful, or if we already are in a parallel session,
-  // then abort multi-threading
-  // FIXME omp_get_num_threads()>1 only works for openmp, what if the user does not use openmp?
-  if((!Condition) || (threads==1) || (omp_get_num_threads()>1))
-    return func(0,rows, 0,cols);
+  // if multi-threading is explicitly disabled, not useful, or if we already are
+  // inside a parallel session, then abort multi-threading
+  bool dont_parallelize = (!Condition) || (threads <= 1);
+#if defined(EIGEN_HAS_OPENMP)
+  // don't parallelize if we are executing in a parallel context already.
+  dont_parallelize |= omp_get_num_threads() > 1;
+#elif defined(EIGEN_GEMM_THREADPOOL)
+  // don't parallelize if we have a trivial threadpool or the current thread id
+  // is != -1, indicating that we are already executing on a thread inside the pool.
+  // In other words, we do not allow nested parallelism, since this would lead to
+  // deadlocks due to the workstealing nature of the threadpool.
+  ThreadPool* pool = getGemmThreadPool();
+  dont_parallelize |= (pool == nullptr || pool->CurrentThreadId() != -1);
+#endif
+  if (dont_parallelize) return func(0, rows, 0, cols);
 
-  Eigen::initParallel();
   func.initParallelSession(threads);
 
-  if(transpose)
-    std::swap(rows,cols);
+  if (transpose) std::swap(rows, cols);
 
-  ei_declare_aligned_stack_constructed_variable(GemmParallelInfo<Index>,info,threads,0);
+  ei_declare_aligned_stack_constructed_variable(GemmParallelTaskInfo<Index>, task_info, threads, 0);
 
-  #pragma omp parallel num_threads(threads)
+#if defined(EIGEN_HAS_OPENMP)
+#pragma omp parallel num_threads(threads)
   {
     Index i = omp_get_thread_num();
-    // Note that the actual number of threads might be lower than the number of request ones.
+    // Note that the actual number of threads might be lower than the number of
+    // requested ones
     Index actual_threads = omp_get_num_threads();
+    GemmParallelInfo<Index> info(static_cast<int>(i), static_cast<int>(actual_threads), task_info);
+
+    Index blockCols = (cols / actual_threads) & ~Index(0x3);
+    Index blockRows = (rows / actual_threads);
+    blockRows = (blockRows / Functor::Traits::mr) * Functor::Traits::mr;
+
+    Index r0 = i * blockRows;
+    Index actualBlockRows = (i + 1 == actual_threads) ? rows - r0 : blockRows;
 
+    Index c0 = i * blockCols;
+    Index actualBlockCols = (i + 1 == actual_threads) ? cols - c0 : blockCols;
+
+    info.task_info[i].lhs_start = r0;
+    info.task_info[i].lhs_length = actualBlockRows;
+
+    if (transpose)
+      func(c0, actualBlockCols, 0, rows, &info);
+    else
+      func(0, rows, c0, actualBlockCols, &info);
+  }
+
+#elif defined(EIGEN_GEMM_THREADPOOL)
+  Barrier barrier(threads);
+  auto task = [=, &func, &barrier, &task_info](int i) {
+    Index actual_threads = threads;
+    GemmParallelInfo<Index> info(i, static_cast<int>(actual_threads), task_info);
     Index blockCols = (cols / actual_threads) & ~Index(0x3);
     Index blockRows = (rows / actual_threads);
-    blockRows = (blockRows/Functor::Traits::mr)*Functor::Traits::mr;
+    blockRows = (blockRows / Functor::Traits::mr) * Functor::Traits::mr;
 
-    Index r0 = i*blockRows;
-    Index actualBlockRows = (i+1==actual_threads) ? rows-r0 : blockRows;
+    Index r0 = i * blockRows;
+    Index actualBlockRows = (i + 1 == actual_threads) ? rows - r0 : blockRows;
 
-    Index c0 = i*blockCols;
-    Index actualBlockCols = (i+1==actual_threads) ? cols-c0 : blockCols;
+    Index c0 = i * blockCols;
+    Index actualBlockCols = (i + 1 == actual_threads) ? cols - c0 : blockCols;
 
-    info[i].lhs_start = r0;
-    info[i].lhs_length = actualBlockRows;
+    info.task_info[i].lhs_start = r0;
+    info.task_info[i].lhs_length = actualBlockRows;
 
-    if(transpose)
-      func(c0, actualBlockCols, 0, rows, info);
+    if (transpose)
+      func(c0, actualBlockCols, 0, rows, &info);
     else
-      func(0, rows, c0, actualBlockCols, info);
+      func(0, rows, c0, actualBlockCols, &info);
+
+    barrier.Notify();
+  };
+  // Notice that we do not schedule more than "threads" tasks, which allows us to
+  // limit number of running threads, even if the threadpool itself was constructed
+  // with a larger number of threads.
+  for (int i = 0; i < threads - 1; ++i) {
+    pool->Schedule([=, task = std::move(task)] { task(i); });
   }
+  task(threads - 1);
+  barrier.Wait();
 #endif
 }
 
-} // end namespace internal
+#endif
 
-} // end namespace Eigen
+}  // end namespace internal
+}  // end namespace Eigen
 
-#endif // EIGEN_PARALLELIZER_H
+#endif  // EIGEN_PARALLELIZER_H