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/detail/can_invoke.h>
#include <dispenso/detail/per_thread_info.h>
#include <dispenso/small_buffer_allocator.h>
#include <dispenso/task_set.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, kAuto };
 
struct ParForOptions {
  uint32_t maxThreads = std::numeric_limits<int32_t>::max();
  bool wait = true;
 
  ParForChunking defaultChunking = ParForChunking::kStatic;
 
  uint32_t minItemsPerChunk = 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) const {
    size_type workingThreads = static_cast<size_type>(numLaunched) + size_type{oneOnCaller};
    assert(workingThreads > 0);
 
    if (!chunk) {
      size_type dynFactor = std::min<size_type>(16, size() / workingThreads);
      size_type chunkSize;
      do {
        size_type roughChunks = dynFactor * workingThreads;
        chunkSize = (size() + roughChunks - 1) / roughChunks;
        --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;
  }
};
 
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 TaskSetT,
    typename IntegerT,
    typename F,
    typename StateContainer,
    typename StateGen>
void parallel_for_staticImpl(
    TaskSetT& taskSet,
    StateContainer& states,
    const StateGen& defaultState,
    const ChunkedRange<IntegerT>& range,
    F&& f,
    ssize_t maxThreads,
    bool wait,
    bool reuseExistingState) {
  using size_type = typename ChunkedRange<IntegerT>::size_type;
 
  size_type numThreads = std::min<size_type>(taskSet.numPoolThreads() + wait, maxThreads);
  // Reduce threads used if they exceed work to be done.
  numThreads = std::min(numThreads, range.size());
 
  detail::initStates(states, defaultState, static_cast<size_t>(numThreads), reuseExistingState);
 
  auto chunking =
      detail::staticChunkSize(static_cast<ssize_t>(range.size()), static_cast<ssize_t>(numThreads));
  IntegerT chunkSize = static_cast<IntegerT>(chunking.ceilChunkSize);
 
  bool perfectlyChunked = static_cast<size_type>(chunking.transitionTaskIndex) == numThreads;
 
  // Number of tasks to schedule (all but the last one if wait is true)
  size_type numToSchedule = wait ? numThreads - 1 : numThreads;
 
  if (numToSchedule > 0) {
    // Precompute range boundaries for the generator
    // First loop: indices [0, transitionTaskIndex) use ceilChunkSize
    // Second loop: indices [transitionTaskIndex, numToSchedule) use ceilChunkSize - 1
    size_type transitionIdx = perfectlyChunked ? numToSchedule : chunking.transitionTaskIndex;
    IntegerT smallChunkSize = static_cast<IntegerT>(chunkSize - !perfectlyChunked);
 
    taskSet.scheduleBulk(
        static_cast<size_t>(numToSchedule),
        [&, chunkSize, smallChunkSize, transitionIdx](size_t idx) {
          // Calculate start position for this chunk
          IntegerT start;
          IntegerT thisChunkSize;
          if (static_cast<size_type>(idx) < transitionIdx) {
            IntegerT i = static_cast<IntegerT>(idx);
            start = static_cast<IntegerT>(range.start + static_cast<IntegerT>(i * chunkSize));
            thisChunkSize = chunkSize;
          } else {
            // After transition, chunks are smaller by 1
            IntegerT ti = static_cast<IntegerT>(transitionIdx);
            IntegerT ri = static_cast<IntegerT>(idx - transitionIdx);
            start = static_cast<IntegerT>(
                range.start + static_cast<IntegerT>(ti * chunkSize) +
                static_cast<IntegerT>(ri * smallChunkSize));
            thisChunkSize = smallChunkSize;
          }
          IntegerT end = static_cast<IntegerT>(start + thisChunkSize);
 
          auto stateIt = states.begin();
          std::advance(stateIt, static_cast<ptrdiff_t>(idx));
 
          return [it = stateIt, start, end, f]() {
            auto recurseInfo = detail::PerPoolPerThreadInfo::parForRecurse();
            f(*it, start, end);
          };
        });
  }
 
  if (wait) {
    // Execute the last chunk on the calling thread
    auto stateIt = states.begin();
    std::advance(stateIt, static_cast<ptrdiff_t>(numThreads - 1));
    // Calculate start of last chunk
    size_type transitionIdx = perfectlyChunked ? numThreads - 1 : chunking.transitionTaskIndex;
    IntegerT smallChunkSize = static_cast<IntegerT>(chunkSize - !perfectlyChunked);
    IntegerT lastStart;
    if (numThreads - 1 < transitionIdx) {
      IntegerT i = static_cast<IntegerT>(numThreads - 1);
      lastStart = static_cast<IntegerT>(range.start + static_cast<IntegerT>(i * chunkSize));
    } else {
      IntegerT ti = static_cast<IntegerT>(transitionIdx);
      IntegerT ri = static_cast<IntegerT>(numThreads - 1 - transitionIdx);
      lastStart = static_cast<IntegerT>(
          range.start + static_cast<IntegerT>(ti * chunkSize) +
          static_cast<IntegerT>(ri * smallChunkSize));
    }
    f(*stateIt, lastStart, range.end);
    taskSet.wait();
  }
}
 
template <typename IntegerT>
struct ChunkSizingResult {
  typename ChunkedRange<IntegerT>::size_type maxThreads;
  bool isStatic;
};
 
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;
 
  // Step 1: never use more threads than the pool actually has
  maxThreads = std::min<size_type>(maxThreads, poolThreads + wait);
 
  if (minItemsPerChunk > 1) {
    // Step 2a: reduce threads so each gets at least minItemsPerChunk items
    size_type maxWorkers = range.size() / minItemsPerChunk;
    if (maxWorkers < maxThreads) {
      maxThreads = maxWorkers;
    }
    // Step 2b: if dynamic chunks would still be too small, use static scheduling
    if (maxThreads > 0 && range.size() / (maxThreads + wait) < minItemsPerChunk && range.isAuto()) {
      isStatic = true;
    }
  } else if (range.size() <= poolThreads + wait) {
    // Step 3: fewer items than threads — static is better (no atomic overhead)
    if (range.isAuto()) {
      isStatic = true;
    } else if (!range.isStatic()) {
      maxThreads = range.size() - wait;
    }
  }
 
  return {maxThreads, isStatic};
}
 
template <
    typename TaskSetT,
    typename IntegerT,
    typename F,
    typename StateContainer,
    typename IndexRef,
    typename ExitAction>
void parallel_for_dynamicImpl(
    TaskSetT& taskSet,
    StateContainer& states,
    IntegerT start,
    IntegerT end,
    F&& f,
    size_t numToLaunch,
    typename ChunkedRange<IntegerT>::size_type chunkSize,
    typename ChunkedRange<IntegerT>::size_type numChunks,
    IndexRef& index,
    ExitAction exitAction,
    bool wait) {
  auto worker = [start, end, &index, f, chunkSize, numChunks, exitAction](auto& s) {
    auto recurseInfo = detail::PerPoolPerThreadInfo::parForRecurse();
    while (true) {
      auto cur = index.fetch_add(1, std::memory_order_relaxed);
      if (cur >= numChunks) {
        exitAction(cur);
        break;
      }
      auto sidx = static_cast<IntegerT>(start + cur * chunkSize);
      if (cur + 1 == numChunks) {
        f(s, sidx, end);
      } else {
        f(s, sidx, static_cast<IntegerT>(sidx + chunkSize));
      }
    }
  };
 
  taskSet.scheduleBulk(static_cast<size_t>(numToLaunch), [&states, &worker](size_t i) {
    auto it = states.begin();
    std::advance(it, static_cast<ptrdiff_t>(i));
    return [&s = *it, worker]() { worker(s); };
  });
 
  if (wait) {
    auto it = states.begin();
    std::advance(it, static_cast<ptrdiff_t>(numToLaunch));
    worker(*it);
    taskSet.wait();
  }
}
 
} // namespace detail
 
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;
 
  uint32_t minItemsPerChunk = std::max<uint32_t>(1, options.minItemsPerChunk);
  size_type maxThreads = std::max<int32_t>(options.maxThreads, 1);
  bool isStatic = range.isStatic();
 
  const size_type N = taskSet.numPoolThreads();
  if (N == 0 || !options.maxThreads || range.size() <= minItemsPerChunk ||
      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(range, maxThreads, isStatic, minItemsPerChunk, N, options.wait);
  maxThreads = chunkSizing.maxThreads;
  isStatic = chunkSizing.isStatic;
 
  // If adjustment reduced threads below 2, run inline — not worth parallelizing.
  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,
        range,
        std::forward<F>(f),
        static_cast<ssize_t>(maxThreads),
        options.wait,
        options.reuseExistingState);
    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);
 
  if (numToLaunch == 1 && !options.wait) {
    taskSet.schedule(
        [&s = states.front(), range, f = std::move(f)]() { f(s, range.start, range.end); });
    return;
  }
 
  auto chunkInfo = range.calcChunkSize(numToLaunch, options.wait, minItemsPerChunk);
  auto chunkSize = std::get<0>(chunkInfo);
  auto numChunks = std::get<1>(chunkInfo);
 
  if (options.wait) {
    alignas(kCacheLineSize) std::atomic<decltype(numChunks)> index(0);
    detail::parallel_for_dynamicImpl(
        taskSet,
        states,
        range.start,
        range.end,
        std::forward<F>(f),
        static_cast<size_t>(numToLaunch),
        chunkSize,
        numChunks,
        index,
        [](auto) {},
        options.wait);
  } else {
    using SizeType = decltype(numChunks);
    struct ChunkIndex {
      std::atomic<SizeType> index;
    };
    static_assert(sizeof(ChunkIndex) <= kCacheLineSize, "ChunkIndex must fit in one cache line");
    char* mem = allocSmallBuffer<kCacheLineSize>();
    auto* ci = new (mem) ChunkIndex{{0}};
    SizeType lastExit = numChunks + static_cast<SizeType>(numToLaunch) - 1;
    detail::parallel_for_dynamicImpl(
        taskSet,
        states,
        range.start,
        range.end,
        std::forward<F>(f),
        static_cast<size_t>(numToLaunch),
        chunkSize,
        numChunks,
        ci->index,
        [ci, lastExit](auto cur) {
          if (cur == lastExit) {
            deallocSmallBuffer<kCacheLineSize>(ci);
          }
        },
        options.wait);
  }
}
 
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);
}
 
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>
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);
}
 
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>
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);
}
 
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>
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);
}
 
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>
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);
}
 
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>
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);
}
 
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>
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