Parallelizer.h

// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2010 Gael Guennebaud <gael.guennebaud@inria.fr>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0
 
#ifndef EIGEN_PARALLELIZER_H
#define EIGEN_PARALLELIZER_H
 
// IWYU pragma: private
#include "../InternalHeaderCheck.h"
 
// 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.
 
#if defined(EIGEN_HAS_OPENMP) && defined(EIGEN_GEMM_THREADPOOL)
#error "EIGEN_HAS_OPENMP and EIGEN_GEMM_THREADPOOL may not both be defined."
#endif
 
namespace Eigen {
 
namespace internal {
inline void manage_multi_threading(Action action, int* v);
}
 
// Public APIs.

EIGEN_DEPRECATED_WITH_REASON("Initialization is no longer needed.") inline void initParallel() {}

inline int nbThreads() {
  int ret;
  internal::manage_multi_threading(GetAction, &ret);
  return ret;
}

inline void setNbThreads(int v) { internal::manage_multi_threading(SetAction, &v); }
 
#ifdef EIGEN_VECTORIZE_SME
#ifndef EIGEN_SME_MIN_TASK_SIZE
#define EIGEN_SME_MIN_TASK_SIZE (double(1 << 20))
#endif
namespace internal {
// Time of one multiply-add on the SME kernel in fp32 multiply-adds: an FP64 outer product covers a quarter of the
// elements of an FP32 one at the same vector length, and a complex multiply-add is four real ones.
template <typename Scalar>
struct sme_madd_cost {
  static constexpr double value = (std::is_same<typename NumTraits<Scalar>::Real, double>::value ? 4.0 : 1.0) *
                                  (NumTraits<Scalar>::IsComplex ? 4.0 : 1.0);
};
// Apple silicon shares one SME unit per core cluster, so more threads than clusters only contend for the units; other
// SME implementations may have one per core. The count comes from macOS's performance-cluster topology and is 0 (no
// cap) everywhere else, where EIGEN_SME_UNITS or setNbSmeUnits() supply it.
inline int detect_sme_units() {
#if defined(EIGEN_SME_UNITS)
  return EIGEN_SME_UNITS;
#elif EIGEN_OS_MAC
  // Performance clusters only (a cluster shares one L2): the work is split evenly over the threads,
  // so a thread on the efficiency cluster's slower unit would set the finishing time.
  int32_t cores = 0, per_l2 = 0;
  size_t sz = sizeof(cores);
  if (sysctlbyname("hw.perflevel0.physicalcpu", &cores, &sz, nullptr, 0) != 0 || cores <= 0) return 0;
  sz = sizeof(per_l2);
  if (sysctlbyname("hw.perflevel0.cpusperl2", &per_l2, &sz, nullptr, 0) != 0 || per_l2 <= 0) return 0;
  return numext::maxi<int32_t>(1, cores / per_l2);
#else
  return 0;
#endif
}
inline void manage_sme_units(Action action, int* v) {
  static int m_units = detect_sme_units();
  if (action == SetAction)
    m_units = *v;
  else
    *v = m_units;
}
}  // namespace internal
#endif

inline int nbSmeUnits() {
#ifdef EIGEN_VECTORIZE_SME
  int ret;
  internal::manage_sme_units(GetAction, &ret);
  return ret;
#else
  return 0;
#endif
}
inline void setNbSmeUnits(int v) {
#ifdef EIGEN_VECTORIZE_SME
  internal::manage_sme_units(SetAction, &v);
#else
  EIGEN_UNUSED_VARIABLE(v);
#endif
}
 
#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 only replaces the stored pointer: work already scheduled on the old
    // ThreadPool is not waited for, and the old pool is not destroyed. Since
    // this returns the new pool, the caller must keep its own pointer to the
    // old one to dispose of it.
    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 {
 
// Implementation.
 
// Granularity of the column split. Aligning to Traits::nr instead would avoid a partial leading
// rhs panel per thread, but nr is 8 on several backends and the coarser split balances worse: on a
// 72-core Neoverse V2, nr-alignment cost 3.5-4.5% at high thread counts (8192x8192 float, 72
// threads: 5132 vs 5315 GFLOP/s). gebp handles a partial panel, so prefer the finer balance.
constexpr int kGemmColGrain = 4;

template <typename Index>
EIGEN_ALWAYS_INLINE void balanced_gemm_range(Index total, Index parts, Index grain, Index part, Index& start,
                                             Index& length) {
  const Index chunks = numext::div_ceil(total, grain);
  const Index base = chunks / parts;
  const Index extra = chunks % parts;
  const Index first_chunk = part * base + numext::mini(part, extra);
  const Index chunk_count = base + (part < extra ? 1 : 0);
  start = numext::mini(first_chunk * grain, total);
  length = numext::mini(chunk_count * grain, total - start);
}
 
#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);
}
 
#else
 
template <typename Index>
struct GemmParallelTaskInfo {
  GemmParallelTaskInfo() {}
  std::atomic<Index> sync{Index(-1)};
  std::atomic<int> users{0};
  Index lhs_start = 0;
  Index lhs_length = 0;
};
 
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<int>(*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);
  }
}
 
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.
  Index size = transpose ? rows : cols;
  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 kMinTaskSize = 50000;  // FIXME: tune this minimum task-size heuristic based on architecture and scalar type.
#ifdef EIGEN_VECTORIZE_SME
  // With a known SME unit count (one unit per core cluster), the SME kernel does far more work per unit of time than
  // the generic one, so a thread needs more of it to repay the handoff; elsewhere the size is not measured and the
  // generic one applies. With no preallocated buffers, each thread runs a disjoint part of the result (below).
  constexpr bool kSme =
      sme_has_gebp_kernel<typename Functor::Traits::LhsScalar, typename Functor::Traits::RhsScalar>::value;
  const bool sme_units_known = kSme && nbSmeUnits() > 0;
  const bool sme_disjoint = sme_units_known && func.ownsNoBuffers();
  if (sme_units_known)
    kMinTaskSize = EIGEN_SME_MIN_TASK_SIZE / sme_madd_cost<typename Functor::Traits::LhsScalar>::value;
  if (sme_disjoint) {
    const Index kernel_rows = transpose ? cols : rows, kernel_cols = transpose ? rows : cols;
    pb_max_threads =
        std::max<Index>(1, (std::max)(kernel_rows / Functor::Traits::mr, kernel_cols / Functor::Traits::nr));
  }
#endif
  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
  int threads = std::min<int>(nbThreads(), static_cast<int>(pb_max_threads));
#ifdef EIGEN_VECTORIZE_SME
  // The SME kernels share one unit per cluster; more threads than units only contend.
  EIGEN_IF_CONSTEXPR (kSme) {
    const int units = nbSmeUnits();
    if (units > 0) threads = std::min<int>(threads, units);
  }
#endif
 
  // 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);
 
#ifdef EIGEN_VECTORIZE_SME
  // A known unit count stands for one SME unit per core cluster, each with its own L2 (see nbSmeUnits): the
  // cooperative session below would read the other cluster's packed LHS and wait for it at every depth step, so each
  // thread runs a disjoint part of the result instead.
  EIGEN_IF_CONSTEXPR (kSme) {
    if (sme_disjoint) {
      // On the kernel's problem (the transposed one for a row-major result), a row split repacks the whole RHS per
      // thread and a column split the whole LHS (a copy, or nothing when read in place): split the rows only when
      // they are twice the columns and give 8 LHS panels per part, since a narrower result already runs close to peak.
      const Index kernel_rows = transpose ? cols : rows, kernel_cols = transpose ? rows : cols;
      const Index row_parts = kernel_rows / (8 * Functor::Traits::mr);
      const bool split_kernel_rows = kernel_rows >= 2 * kernel_cols && row_parts >= 2;
      const bool split_rows = transpose ? !split_kernel_rows : split_kernel_rows;
      if (split_kernel_rows) threads = static_cast<int>((std::min)(Index(threads), row_parts));
      // The result's rows are the kernel's columns for a row-major result. No part may be empty.
      const Index row_grain = transpose ? Functor::Traits::nr : Functor::Traits::mr;
      const Index col_grain = transpose ? Functor::Traits::mr : Functor::Traits::nr;
      const Index chunks = split_rows ? (rows + row_grain - 1) / row_grain : (cols + col_grain - 1) / col_grain;
      threads = static_cast<int>((std::min)(Index(threads), chunks));
      if (threads <= 1) return func(0, rows, 0, cols);
      auto part = [&func, rows, cols, threads, split_rows, row_grain, col_grain](int i) {
        Index start, length;
        if (split_rows) {
          balanced_gemm_range<Index>(rows, threads, row_grain, i, start, length);
          if (length > 0) func(start, length, 0, cols);
        } else {
          balanced_gemm_range<Index>(cols, threads, col_grain, i, start, length);
          if (length > 0) func(0, rows, start, length);
        }
      };
#if defined(EIGEN_HAS_OPENMP)
#pragma omp parallel for num_threads(threads) schedule(static, 1)
      for (int i = 0; i < threads; ++i) part(i);
#elif defined(EIGEN_GEMM_THREADPOOL)
      // The parts take about as long as each other, so the caller spins for the others (tens of microseconds, less than
      // a blocking wake-up costs a short product) before it blocks on the barrier, which it always passes, also
      // before an exception from its own part leaves this frame.
      std::atomic<int> pending(threads - 1);
      Barrier done(threads - 1);
      for (int i = 0; i < threads - 1; ++i)
        pool->Schedule([&part, &pending, &done, i] {
          part(i);
          pending.fetch_sub(1, std::memory_order_release);
          done.Notify();
        });
      const auto wait = [&pending, &done] {
        for (int spin = 0; spin < 65536 && pending.load(std::memory_order_acquire) != 0; ++spin) {
        }
        done.Wait();
      };
      EIGEN_TRY { part(threads - 1); }
      EIGEN_CATCH(...) {
        wait();
        EIGEN_THROW;
      }
      wait();
#endif
      return;
    }
  }
#endif
 
  func.initParallelSession(threads);
 
  if (transpose) std::swap(rows, cols);
 
  ei_declare_aligned_stack_constructed_variable(GemmParallelTaskInfo<Index>, task_info, threads, 0);
 
#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
    // requested ones
    Index actual_threads = omp_get_num_threads();
    GemmParallelInfo<Index> info(static_cast<int>(i), static_cast<int>(actual_threads), task_info);
 
    Index r0, actualBlockRows;
    balanced_gemm_range<Index>(rows, actual_threads, Index(Functor::Traits::mr), i, r0, actualBlockRows);
 
    Index c0, actualBlockCols;
    balanced_gemm_range<Index>(cols, actual_threads, Index(kGemmColGrain), i, c0, actualBlockCols);
 
    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 r0, actualBlockRows;
    balanced_gemm_range<Index>(rows, actual_threads, Index(Functor::Traits::mr), i, r0, actualBlockRows);
 
    Index c0, actualBlockCols;
    balanced_gemm_range<Index>(cols, actual_threads, Index(kGemmColGrain), i, c0, actualBlockCols);
 
    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);
 
    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
}
 
#endif
 
}  // end namespace internal
}  // end namespace Eigen
 
#endif  // EIGEN_PARALLELIZER_H