parallel_for.h

Go to the documentation of this file.

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */
 
#pragma once
 
#include <cmath>
#include <limits>
 
#include <dispenso/cpu_set.h>
#include <dispenso/small_buffer_allocator.h>
#include <dispenso/task_set.h>
#include "detail/can_invoke.h"
#include "detail/par_for_stripe.h"
#include "detail/per_thread_info.h"
 
namespace dispenso {
 
#if DISPENSO_HAS_CONCEPTS
template <typename F, typename IntegerT>
concept ParallelForRangeFunc = std::invocable<F, IntegerT, IntegerT>;
 
template <typename F, typename IntegerT>
concept ParallelForIndexFunc = std::invocable<F, IntegerT>;
 
template <typename F, typename StateRef, typename IntegerT>
concept ParallelForStateRangeFunc = std::invocable<F, StateRef, IntegerT, IntegerT>;
 
template <typename F, typename StateRef, typename IntegerT>
concept ParallelForStateIndexFunc = std::invocable<F, StateRef, IntegerT>;
#endif // DISPENSO_HAS_CONCEPTS
 
enum class ParForChunking {
  kStatic,
  kAdaptive,
  kAuto DISPENSO_DEPRECATED("Use ParForChunking::kAdaptive (kAuto will be removed in 2.0).") =
      kAdaptive,
};
 
struct ParForOptions {
  uint32_t maxThreads = std::numeric_limits<int32_t>::max();
  bool wait = true;
 
  ParForChunking defaultChunking = ParForChunking::kStatic;
 
  uint32_t minItemsPerChunk = 1;
 
  uint32_t granularity = 1;
 
  bool reuseExistingState = false;
};
 
template <typename IntegerT = ssize_t>
struct ChunkedRange {
  // We need to utilize 64-bit integers to avoid overflow, e.g. passing -2**30, 2**30 as int32 will
  // result in overflow unless we cast to 64-bit.  Note that if we have a range of e.g. -2**63+1 to
  // 2**63-1, we cannot hold the result in an int64_t.  We could in a uint64_t, but it is quite
  // tricky to make this work.  However, I do not expect ranges larger than can be held in int64_t
  // since people want their computations to finish before the heat death of the sun (slight
  // exaggeration).
  using size_type = std::conditional_t<std::is_signed<IntegerT>::value, int64_t, uint64_t>;
 
  struct Static {};
  struct Auto {};
  static constexpr IntegerT kStatic = std::numeric_limits<IntegerT>::max();
 
  ChunkedRange(IntegerT s, IntegerT e, IntegerT c) : start(s), end(e), chunk(c) {}
  ChunkedRange(IntegerT s, IntegerT e, Static) : ChunkedRange(s, e, kStatic) {}
  ChunkedRange(IntegerT s, IntegerT e, Auto) : ChunkedRange(s, e, 0) {}
 
  bool isStatic() const {
    return chunk == kStatic;
  }
 
  bool isAuto() const {
    return chunk == 0;
  }
 
  bool empty() const {
    return end <= start;
  }
 
  size_type size() const {
    return static_cast<size_type>(end) - start;
  }
 
  template <typename OtherInt>
  std::tuple<size_type, size_type> calcChunkSize(
      OtherInt numLaunched,
      bool oneOnCaller,
      size_type minChunkSize,
      uint32_t granularity = 1,
      size_type maxDynFactor = 16) const {
    size_type workingThreads = static_cast<size_type>(numLaunched) + size_type{oneOnCaller};
    assert(workingThreads > 0);
 
    if (!chunk) {
      size_type dynFactor = std::min<size_type>(maxDynFactor, size() / workingThreads);
      size_type chunkSize;
      do {
        size_type roughChunks = dynFactor * workingThreads;
        chunkSize = (size() + roughChunks - 1) / roughChunks;
        if (granularity > 1) {
          // Round UP to a multiple of granularity (no smaller than granularity).
          chunkSize = ((chunkSize + granularity - 1) / granularity) * granularity;
        }
        --dynFactor;
      } while (chunkSize < minChunkSize);
      return {chunkSize, (size() + chunkSize - 1) / chunkSize};
    } else if (chunk == kStatic) {
      // This should never be called.  The static distribution versions of the parallel_for
      // functions should be invoked instead.
      std::abort();
    }
    return {chunk, (size() + chunk - 1) / chunk};
  }
 
  IntegerT start;
  IntegerT end;
  IntegerT chunk;
};
 
template <typename IntegerA, typename IntegerB>
inline ChunkedRange<std::common_type_t<IntegerA, IntegerB>>
makeChunkedRange(IntegerA start, IntegerB end, ParForChunking chunking = ParForChunking::kStatic) {
  using IntegerT = std::common_type_t<IntegerA, IntegerB>;
  return (chunking == ParForChunking::kStatic)
      ? ChunkedRange<IntegerT>(start, end, typename ChunkedRange<IntegerT>::Static())
      : ChunkedRange<IntegerT>(start, end, typename ChunkedRange<IntegerT>::Auto());
}
 
template <typename IntegerA, typename IntegerB, typename IntegerC>
inline ChunkedRange<std::common_type_t<IntegerA, IntegerB>>
makeChunkedRange(IntegerA start, IntegerB end, IntegerC chunkSize) {
  return ChunkedRange<std::common_type_t<IntegerA, IntegerB>>(start, end, chunkSize);
}
 
namespace detail {
 
struct NoOpIter {
  using difference_type = std::ptrdiff_t;
  using value_type = int;
  using pointer = int*;
  using reference = int&;
  using iterator_category = std::random_access_iterator_tag;
 
  int& operator*() const {
    static DISPENSO_THREAD_LOCAL int dummy = 0;
    return dummy;
  }
  NoOpIter& operator++() {
    return *this;
  }
  NoOpIter operator++(int) {
    return *this;
  }
  NoOpIter& operator--() {
    return *this;
  }
  NoOpIter operator--(int) {
    return *this;
  }
  NoOpIter& operator+=(difference_type) {
    return *this;
  }
  NoOpIter& operator-=(difference_type) {
    return *this;
  }
  NoOpIter operator+(difference_type) const {
    return *this;
  }
  NoOpIter operator-(difference_type) const {
    return *this;
  }
  difference_type operator-(const NoOpIter&) const {
    return 0;
  }
  bool operator==(const NoOpIter&) const {
    return true;
  }
  bool operator!=(const NoOpIter&) const {
    return false;
  }
  bool operator<(const NoOpIter&) const {
    return false;
  }
  int& operator[](difference_type) const {
    static DISPENSO_THREAD_LOCAL int dummy = 0;
    return dummy;
  }
};
 
struct NoOpContainer {
  size_t size() const {
    return 0;
  }
 
  bool empty() const {
    return true;
  }
 
  void clear() {}
 
  NoOpIter begin() {
    return {};
  }
 
  void emplace_back(int) {}
 
  int& front() {
    static int i;
    return i;
  }
};
 
struct NoOpStateGen {
  int operator()() const {
    return 0;
  }
};
 
// Round size DOWN to a multiple of granularity. Granularity must be >= 1.
template <typename IntegerT>
inline IntegerT roundDownToGranularity(IntegerT size, uint32_t granularity) {
  if (granularity <= 1) {
    return size;
  }
  using U = typename std::make_unsigned<IntegerT>::type;
  return static_cast<IntegerT>((static_cast<U>(size) / granularity) * granularity);
}
 
// Round size UP to a multiple of granularity. Granularity must be >= 1.
template <typename IntegerT>
inline IntegerT roundUpToGranularity(IntegerT size, uint32_t granularity) {
  if (granularity <= 1) {
    return size;
  }
  using U = typename std::make_unsigned<IntegerT>::type;
  return static_cast<IntegerT>(
      ((static_cast<U>(size) + granularity - 1) / granularity) * granularity);
}
 
// Initialize states container with enough entries for the given thread count.
// Respects reuseExistingState: when true, only adds entries if the container
// doesn't already have enough.
template <typename StateContainer, typename StateGen>
void initStates(
    StateContainer& states,
    const StateGen& defaultState,
    size_t numNeeded,
    bool reuseExistingState) {
  if (!reuseExistingState) {
    states.clear();
  }
  for (size_t i = states.size(); i < numNeeded; ++i) {
    states.emplace_back(defaultState());
  }
}
 
template <typename IntegerT>
struct ChunkSizingResult {
  typename ChunkedRange<IntegerT>::size_type maxThreads;
  bool isStatic;
};
 
// Adjust the thread count and static/dynamic scheduling decision based on work size.
template <typename IntegerT>
ChunkSizingResult<IntegerT> adjustChunkSizing(
    const ChunkedRange<IntegerT>& range,
    typename ChunkedRange<IntegerT>::size_type maxThreads,
    bool isStatic,
    uint32_t minItemsPerChunk,
    typename ChunkedRange<IntegerT>::size_type poolThreads,
    bool wait) {
  using size_type = typename ChunkedRange<IntegerT>::size_type;
 
  maxThreads = std::min<size_type>(maxThreads, poolThreads + 1);
 
  if (minItemsPerChunk > 1) {
    size_type maxWorkers = range.size() / minItemsPerChunk;
    if (maxWorkers < maxThreads) {
      maxThreads = maxWorkers;
    }
    if (maxThreads > 0 && range.size() / (maxThreads + wait) < minItemsPerChunk && range.isAuto()) {
      isStatic = true;
    }
  } else if (range.size() <= poolThreads + wait) {
    if (range.isAuto()) {
      isStatic = true;
    } else if (!range.isStatic()) {
      maxThreads = range.size() - wait;
    }
  }
 
  return {maxThreads, isStatic};
}
 
// Compute effective granularity and trimmed range end.
template <typename IntegerT>
struct GranularityInfo {
  uint32_t granularity;
  IntegerT trimmedEnd;
  bool hasTail;
};
 
template <typename IntegerT>
GranularityInfo<IntegerT> computeGranularity(
    const ChunkedRange<IntegerT>& range,
    uint32_t requested) {
  using size_type = typename ChunkedRange<IntegerT>::size_type;
  uint32_t granularity = (range.chunk == 0 || range.chunk == ChunkedRange<IntegerT>::kStatic)
      ? std::max<uint32_t>(1, requested)
      : 1;
 
  IntegerT trimmedEnd = range.end;
  bool hasTail = false;
  if (granularity > 1) {
    size_type rem = range.size() % granularity;
    if (rem > 0) {
      trimmedEnd = static_cast<IntegerT>(range.end - static_cast<IntegerT>(rem));
      hasTail = true;
    }
  }
  return {granularity, trimmedEnd, hasTail};
}
 
// Adaptive (stripe-based) wait dispatch for the top-level parallel_for.
template <typename TaskSetT, typename IntegerT, typename F, typename StateContainer>
void parallel_for_adaptiveWaitDispatch(
    TaskSetT& taskSet,
    StateContainer& states,
    const ChunkedRange<IntegerT>& parRange,
    F&& f,
    size_t numToLaunch,
    uint32_t minItemsPerChunk,
    uint32_t granularity) {
  using size_type = typename ChunkedRange<IntegerT>::size_type;
  size_type numStripeWorkers = static_cast<size_type>(numToLaunch) + 1;
  auto adaptiveChunkInfo =
      parRange.calcChunkSize(numToLaunch, true, minItemsPerChunk, granularity, /*maxDynFactor=*/64);
  auto adaptiveChunkSize = std::get<0>(adaptiveChunkInfo);
 
  detail::StripeState<IntegerT> stripeState;
  detail::initStripeState(
      stripeState,
      parRange.start,
      parRange.end,
      static_cast<uint32_t>(numStripeWorkers),
      static_cast<IntegerT>(adaptiveChunkSize),
      granularity);
  auto stateBegin = states.begin();
  auto worker = [&stripeState, &f](auto& userState, uint32_t myIdx) {
    auto recurseInfo = detail::PerPoolPerThreadInfo::parForRecurse();
    detail::runStripeWorker(stripeState, myIdx, userState, f);
  };
  if (numToLaunch > 0) {
    taskSet.scheduleBulk(numToLaunch, [stateBegin, worker](size_t idx) {
      auto stateIt = stateBegin;
      std::advance(stateIt, static_cast<ptrdiff_t>(idx));
      uint32_t myIdx = static_cast<uint32_t>(idx);
      return [&userState = *stateIt, myIdx, worker]() { worker(userState, myIdx); };
    });
  }
  auto callerIt = states.begin();
  std::advance(callerIt, static_cast<ptrdiff_t>(numToLaunch));
  worker(*callerIt, static_cast<uint32_t>(numToLaunch));
  taskSet.wait();
}
 
} // namespace detail
 
} // namespace dispenso
 
// Implementation detail headers — included after all shared types are defined.
// These are not standalone headers; they depend on types defined above.
#include "detail/par_for_dynamic.h"
#include "detail/par_for_static.h"
 
namespace dispenso {
 
template <
    typename TaskSetT,
    typename IntegerT,
    typename F,
    typename StateContainer,
    typename StateGen>
void parallel_for(
    TaskSetT& taskSet,
    StateContainer& states,
    const StateGen& defaultState,
    const ChunkedRange<IntegerT>& range,
    F&& f,
    ParForOptions options = {}) {
  if (range.empty()) {
    if (options.wait) {
      taskSet.wait();
    }
    return;
  }
 
  using size_type = typename ChunkedRange<IntegerT>::size_type;
 
  auto granInfo = detail::computeGranularity(range, options.granularity);
  uint32_t granularity = granInfo.granularity;
  IntegerT trimmedEnd = granInfo.trimmedEnd;
  bool hasTail = granInfo.hasTail;
 
  auto runTail = [&]() {
    if (hasTail) {
      f(*states.begin(), trimmedEnd, range.end);
    }
  };
 
  uint32_t minItemsPerChunk = std::max<uint32_t>(1, options.minItemsPerChunk);
  size_type maxThreads = std::max<int32_t>(options.maxThreads, 1);
  bool isStatic = range.isStatic();
 
  ChunkedRange<IntegerT> parRange = range;
  parRange.end = trimmedEnd;
 
  const size_type N = taskSet.numPoolThreads();
 
  // If the parallel portion is empty (entire range is a sub-granularity tail),
  // or there's no pool / recursion, run everything inline.
  if (parRange.empty() || N == 0 ||
      detail::PerPoolPerThreadInfo::isParForRecursive(&taskSet.pool())) {
    detail::initStates(states, defaultState, 1, options.reuseExistingState);
    f(*states.begin(), range.start, range.end);
    if (options.wait) {
      taskSet.wait();
    }
    return;
  }
 
  auto chunkSizing =
      detail::adjustChunkSizing(parRange, maxThreads, isStatic, minItemsPerChunk, N, options.wait);
  maxThreads = chunkSizing.maxThreads;
  isStatic = chunkSizing.isStatic;
 
  if (maxThreads < 2) {
    detail::initStates(states, defaultState, 1, options.reuseExistingState);
    f(*states.begin(), range.start, range.end);
    if (options.wait) {
      taskSet.wait();
    }
    return;
  }
 
  if (isStatic) {
    detail::parallel_for_staticImpl(
        taskSet,
        states,
        defaultState,
        parRange,
        std::forward<F>(f),
        static_cast<ssize_t>(maxThreads),
        options.wait,
        options.reuseExistingState,
        granularity);
    runTail();
    return;
  }
 
  const size_type numToLaunch = std::min<size_type>(maxThreads - options.wait, N);
 
  detail::initStates(
      states,
      defaultState,
      static_cast<size_t>(numToLaunch + options.wait),
      options.reuseExistingState);
 
  bool useAdaptive = range.chunk == 0;
 
  auto chunkInfo = parRange.calcChunkSize(numToLaunch, options.wait, minItemsPerChunk, granularity);
  auto chunkSize = std::get<0>(chunkInfo);
  auto numChunks = std::get<1>(chunkInfo);
 
  if (useAdaptive && options.wait) {
    detail::parallel_for_adaptiveWaitDispatch(
        taskSet,
        states,
        parRange,
        std::forward<F>(f),
        static_cast<size_t>(numToLaunch),
        minItemsPerChunk,
        granularity);
    runTail();
    return;
  }
 
  if (options.wait) {
    alignas(kCacheLineSize) std::atomic<decltype(numChunks)> index(0);
    detail::parallel_for_dynamicImpl(
        taskSet,
        states,
        parRange.start,
        parRange.end,
        std::forward<F>(f),
        static_cast<size_t>(numToLaunch),
        chunkSize,
        numChunks,
        index,
        [](auto) {},
        options.wait);
    runTail();
  } else {
    detail::parallel_for_dynamicNoWaitDispatch(
        taskSet,
        states,
        parRange,
        std::forward<F>(f),
        static_cast<size_t>(numToLaunch),
        chunkSize,
        numChunks,
        range.end,
        hasTail);
  }
}
 
template <typename TaskSetT, typename IntegerT, typename F>
DISPENSO_REQUIRES(ParallelForRangeFunc<F, IntegerT>)
void parallel_for(
    TaskSetT& taskSet,
    const ChunkedRange<IntegerT>& range,
    F&& f,
    ParForOptions options = {}) {
  detail::NoOpContainer container;
  parallel_for(
      taskSet,
      container,
      detail::NoOpStateGen(),
      range,
      [f = std::move(f)](int /*noop*/, auto i, auto j) { f(i, j); },
      options);
}
 
template <typename IntegerT, typename F>
DISPENSO_REQUIRES(ParallelForRangeFunc<F, IntegerT>)
void parallel_for(const ChunkedRange<IntegerT>& range, F&& f, ParForOptions options = {}) {
  TaskSet taskSet(globalThreadPool());
  options.wait = true;
  parallel_for(taskSet, range, std::forward<F>(f), options);
}
 
template <typename F, typename IntegerT, typename StateContainer, typename StateGen>
void parallel_for(
    StateContainer& states,
    const StateGen& defaultState,
    const ChunkedRange<IntegerT>& range,
    F&& f,
    ParForOptions options = {}) {
  TaskSet taskSet(globalThreadPool());
  options.wait = true;
  parallel_for(taskSet, states, defaultState, range, std::forward<F>(f), options);
}
 
#if DISPENSO_HAS_CONCEPTS
template <typename TaskSetT, std::integral IntegerA, std::integral IntegerB, typename F>
  requires std::invocable<F, IntegerA>
#else
template <
    typename TaskSetT,
    typename IntegerA,
    typename IntegerB,
    typename F,
    std::enable_if_t<std::is_integral<IntegerA>::value, bool> = true,
    std::enable_if_t<std::is_integral<IntegerB>::value, bool> = true,
    std::enable_if_t<detail::CanInvoke<F(IntegerA)>::value, bool> = true>
#endif
void parallel_for(
    TaskSetT& taskSet,
    IntegerA start,
    IntegerB end,
    F&& f,
    ParForOptions options = {}) {
  using IntegerT = std::common_type_t<IntegerA, IntegerB>;
 
  auto range = makeChunkedRange(start, end, options.defaultChunking);
  parallel_for(
      taskSet,
      range,
      [f = std::move(f)](IntegerT s, IntegerT e) {
        for (IntegerT i = s; i < e; ++i) {
          f(i);
        }
      },
      options);
}
 
#if DISPENSO_HAS_CONCEPTS
template <typename TaskSetT, std::integral IntegerA, std::integral IntegerB, typename F>
  requires std::invocable<F, IntegerA, IntegerB>
#else
template <
    typename TaskSetT,
    typename IntegerA,
    typename IntegerB,
    typename F,
    std::enable_if_t<std::is_integral<IntegerA>::value, bool> = true,
    std::enable_if_t<std::is_integral<IntegerB>::value, bool> = true,
    std::enable_if_t<detail::CanInvoke<F(IntegerA, IntegerB)>::value, bool> = true>
#endif
void parallel_for(
    TaskSetT& taskSet,
    IntegerA start,
    IntegerB end,
    F&& f,
    ParForOptions options = {}) {
  auto range = makeChunkedRange(start, end, options.defaultChunking);
  parallel_for(taskSet, range, std::forward<F>(f), options);
}
 
#if DISPENSO_HAS_CONCEPTS
template <std::integral IntegerA, std::integral IntegerB, typename F>
#else
template <
    typename IntegerA,
    typename IntegerB,
    typename F,
    std::enable_if_t<std::is_integral<IntegerA>::value, bool> = true,
    std::enable_if_t<std::is_integral<IntegerB>::value, bool> = true>
#endif
void parallel_for(IntegerA start, IntegerB end, F&& f, ParForOptions options = {}) {
  TaskSet taskSet(globalThreadPool());
  options.wait = true;
  parallel_for(taskSet, start, end, std::forward<F>(f), options);
}
 
#if DISPENSO_HAS_CONCEPTS
template <
    typename TaskSetT,
    std::integral IntegerA,
    std::integral IntegerB,
    typename F,
    typename StateContainer,
    typename StateGen>
  requires std::invocable<F, typename StateContainer::reference, IntegerA>
#else
template <
    typename TaskSetT,
    typename IntegerA,
    typename IntegerB,
    typename F,
    typename StateContainer,
    typename StateGen,
    std::enable_if_t<std::is_integral<IntegerA>::value, bool> = true,
    std::enable_if_t<std::is_integral<IntegerB>::value, bool> = true,
    std::enable_if_t<
        detail::CanInvoke<F(typename StateContainer::reference, IntegerA)>::value,
        bool> = true>
#endif
void parallel_for(
    TaskSetT& taskSet,
    StateContainer& states,
    const StateGen& defaultState,
    IntegerA start,
    IntegerB end,
    F&& f,
    ParForOptions options = {}) {
  using IntegerT = std::common_type_t<IntegerA, IntegerB>;
  auto range = makeChunkedRange(start, end, options.defaultChunking);
  parallel_for(
      taskSet,
      states,
      defaultState,
      range,
      [f = std::move(f)](auto& state, IntegerT s, IntegerT e) {
        for (IntegerT i = s; i < e; ++i) {
          f(state, i);
        }
      },
      options);
}
 
#if DISPENSO_HAS_CONCEPTS
template <
    typename TaskSetT,
    std::integral IntegerA,
    std::integral IntegerB,
    typename F,
    typename StateContainer,
    typename StateGen>
  requires std::invocable<F, typename StateContainer::reference, IntegerA, IntegerB>
#else
template <
    typename TaskSetT,
    typename IntegerA,
    typename IntegerB,
    typename F,
    typename StateContainer,
    typename StateGen,
    std::enable_if_t<std::is_integral<IntegerA>::value, bool> = true,
    std::enable_if_t<std::is_integral<IntegerB>::value, bool> = true,
    std::enable_if_t<
        detail::CanInvoke<F(typename StateContainer::reference, IntegerA, IntegerB)>::value,
        bool> = true>
#endif
void parallel_for(
    TaskSetT& taskSet,
    StateContainer& states,
    const StateGen& defaultState,
    IntegerA start,
    IntegerB end,
    F&& f,
    ParForOptions options = {}) {
  auto range = makeChunkedRange(start, end, options.defaultChunking);
  parallel_for(taskSet, states, defaultState, range, std::forward<F>(f), options);
}
 
#if DISPENSO_HAS_CONCEPTS
template <
    std::integral IntegerA,
    std::integral IntegerB,
    typename F,
    typename StateContainer,
    typename StateGen>
#else
template <
    typename IntegerA,
    typename IntegerB,
    typename F,
    typename StateContainer,
    typename StateGen,
    std::enable_if_t<std::is_integral<IntegerA>::value, bool> = true,
    std::enable_if_t<std::is_integral<IntegerB>::value, bool> = true>
#endif
void parallel_for(
    StateContainer& states,
    const StateGen& defaultState,
    IntegerA start,
    IntegerB end,
    F&& f,
    ParForOptions options = {}) {
  TaskSet taskSet(globalThreadPool());
  options.wait = true;
  parallel_for(taskSet, states, defaultState, start, end, std::forward<F>(f), options);
}
 
} // namespace dispenso