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 <memory>
 
#include <dispenso/detail/can_invoke.h>
#include <dispenso/detail/per_thread_info.h>
#include <dispenso/task_set.h>
 
namespace dispenso {
 
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 {
  int& operator*() const {
    static int i = 0;
    return i;
  }
  NoOpIter& operator++() {
    return *this;
  }
  NoOpIter operator++(int) {
    return *this;
  }
};
 
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 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());
 
  if (!reuseExistingState) {
    states.clear();
  }
 
  size_t numToEmplace = states.size() < static_cast<size_t>(numThreads)
      ? static_cast<size_t>(numThreads) - states.size()
      : 0;
 
  for (; numToEmplace--;) {
    states.emplace_back(defaultState());
  }
 
  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;
 
  // (!perfectlyChunked) ? chunking.transitionTaskIndex : numThreads - 1;
  size_type firstLoopLen = chunking.transitionTaskIndex - perfectlyChunked;
 
  auto stateIt = states.begin();
  IntegerT start = range.start;
  size_type t;
  for (t = 0; t < firstLoopLen; ++t) {
    IntegerT next = static_cast<IntegerT>(start + chunkSize);
    taskSet.schedule([it = stateIt++, start, next, f]() {
      auto recurseInfo = detail::PerPoolPerThreadInfo::parForRecurse();
      f(*it, start, next);
    });
    start = next;
  }
 
  // Reduce the remaining chunk sizes by 1.
  chunkSize = static_cast<IntegerT>(chunkSize - !perfectlyChunked);
  // Finish submitting all but the last item.
  for (; t < numThreads - 1; ++t) {
    IntegerT next = static_cast<IntegerT>(start + chunkSize);
    taskSet.schedule([it = stateIt++, start, next, f]() {
      auto recurseInfo = detail::PerPoolPerThreadInfo::parForRecurse();
      f(*it, start, next);
    });
    start = next;
  }
 
  if (wait) {
    f(*stateIt, start, range.end);
    taskSet.wait();
  } else {
    taskSet.schedule(
        [stateIt, start, end = range.end, f]() {
          auto recurseInfo = detail::PerPoolPerThreadInfo::parForRecurse();
          f(*stateIt, start, end);
        },
        ForceQueuingTag());
  }
}
 
} // 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;
 
  // Ensure minItemsPerChunk is sane
  uint32_t minItemsPerChunk = std::max<uint32_t>(1, options.minItemsPerChunk);
 
  // 0 indicates serial execution per API spec
  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())) {
    if (!options.reuseExistingState) {
      states.clear();
    }
    if (states.empty()) {
      states.emplace_back(defaultState());
    }
    f(*states.begin(), range.start, range.end);
    if (options.wait) {
      taskSet.wait();
    }
    return;
  }
 
  // Adjust down workers if we would have too-small chunks
  if (minItemsPerChunk > 1) {
    size_type maxWorkers = range.size() / minItemsPerChunk;
    if (maxWorkers < maxThreads) {
      maxThreads = static_cast<uint32_t>(maxWorkers);
    }
    if (range.size() / (maxThreads + options.wait) < minItemsPerChunk && range.isAuto()) {
      isStatic = true;
    }
  } else if (range.size() <= N + options.wait) {
    if (range.isAuto()) {
      isStatic = true;
    } else if (!range.isStatic()) {
      maxThreads = range.size() - options.wait;
    }
  }
 
  if (isStatic) {
    detail::parallel_for_staticImpl(
        taskSet,
        states,
        defaultState,
        range,
        std::forward<F>(f),
        static_cast<ssize_t>(maxThreads),
        options.wait,
        options.reuseExistingState);
    return;
  }
 
  // wanting maxThreads workers (potentially including the calling thread), capped by N
  const size_type numToLaunch = std::min<size_type>(maxThreads - options.wait, N);
 
  if (!options.reuseExistingState) {
    states.clear();
  }
 
  size_t numToEmplace = static_cast<size_type>(states.size()) < (numToLaunch + options.wait)
      ? (static_cast<size_t>(numToLaunch) + options.wait) - states.size()
      : 0;
  for (; numToEmplace--;) {
    states.emplace_back(defaultState());
  }
 
  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);
    auto worker = [start = range.start, end = range.end, &index, f, chunkSize, numChunks](auto& s) {
      auto recurseInfo = detail::PerPoolPerThreadInfo::parForRecurse();
 
      while (true) {
        auto cur = index.fetch_add(1, std::memory_order_relaxed);
        if (cur >= numChunks) {
          break;
        }
        auto sidx = static_cast<IntegerT>(start + cur * chunkSize);
        if (cur + 1 == numChunks) {
          f(s, sidx, end);
        } else {
          auto eidx = static_cast<IntegerT>(sidx + chunkSize);
          f(s, sidx, eidx);
        }
      }
    };
 
    auto it = states.begin();
    for (size_type i = 0; i < numToLaunch; ++i) {
      taskSet.schedule([&s = *it++, worker]() { worker(s); });
    }
    worker(*it);
    taskSet.wait();
  } else {
    struct Atomic {
      Atomic() : index(0) {}
      alignas(kCacheLineSize) std::atomic<decltype(numChunks)> index;
      char buffer[kCacheLineSize - sizeof(index)];
    };
 
    void* ptr = detail::alignedMalloc(sizeof(Atomic), alignof(Atomic));
    auto* atm = new (ptr) Atomic();
 
    std::shared_ptr<Atomic> wrapper(atm, detail::AlignedFreeDeleter<Atomic>());
    auto worker = [start = range.start,
                   end = range.end,
                   wrapper = std::move(wrapper),
                   f,
                   chunkSize,
                   numChunks](auto& s) {
      auto recurseInfo = detail::PerPoolPerThreadInfo::parForRecurse();
      while (true) {
        auto cur = wrapper->index.fetch_add(1, std::memory_order_relaxed);
        if (cur >= numChunks) {
          break;
        }
        auto sidx = static_cast<IntegerT>(start + cur * chunkSize);
        if (cur + 1 == numChunks) {
          f(s, sidx, end);
        } else {
          auto eidx = static_cast<IntegerT>(sidx + chunkSize);
          f(s, sidx, eidx);
        }
      }
    };
 
    auto it = states.begin();
    for (size_type i = 0; i < numToLaunch; ++i) {
      taskSet.schedule([&s = *it++, worker]() { worker(s); }, ForceQueuingTag());
    }
  }
}
 
template <typename TaskSetT, typename IntegerT, typename F>
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>
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