for_each.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 <algorithm>
#include <iterator>
 
#include <dispenso/detail/per_thread_info.h>
#include <dispenso/small_vector.h>
#include <dispenso/task_set.h>
 
namespace dispenso {
 
#if DISPENSO_HAS_CONCEPTS
template <typename F, typename Iter>
concept ForEachFunc = std::invocable<F, decltype(*std::declval<Iter>())>;
#endif // DISPENSO_HAS_CONCEPTS
 
struct ForEachOptions {
  uint32_t maxThreads = std::numeric_limits<int32_t>::max();
  bool wait = true;
};
 
namespace detail {
 
// Random-access iterators: compute chunk boundaries arithmetically, avoiding SmallVector.
template <typename TaskSetT, typename Iter, typename F>
void for_each_n_schedule(
    TaskSetT& tasks,
    Iter start,
    F&& f,
    ssize_t numThreads,
    size_t chunkSize,
    ssize_t transitionIdx,
    size_t smallChunkSize,
    const ForEachOptions& options,
    std::random_access_iterator_tag) {
  ssize_t numToSchedule = options.wait ? numThreads - 1 : numThreads;
 
  if (numToSchedule > 0) {
    tasks.scheduleBulk(
        static_cast<size_t>(numToSchedule),
        [start, &f, chunkSize, smallChunkSize, transitionIdx](size_t idx) {
          ssize_t sidx = static_cast<ssize_t>(idx);
          ssize_t offset;
          ssize_t thisChunkSize;
          if (sidx < transitionIdx) {
            offset = sidx * static_cast<ssize_t>(chunkSize);
            thisChunkSize = static_cast<ssize_t>(chunkSize);
          } else {
            offset = transitionIdx * static_cast<ssize_t>(chunkSize) +
                (sidx - transitionIdx) * static_cast<ssize_t>(smallChunkSize);
            thisChunkSize = static_cast<ssize_t>(smallChunkSize);
          }
          Iter s = start + offset;
          Iter e = s + thisChunkSize;
          return [s, e, f]() {
            auto recurseInfo = PerPoolPerThreadInfo::parForRecurse();
            for (Iter it = s; it != e; ++it) {
              f(*it);
            }
          };
        });
  }
 
  if (options.wait) {
    ssize_t lastIdx = numThreads - 1;
    ssize_t offset;
    ssize_t thisChunkSize;
    if (lastIdx < transitionIdx) {
      offset = lastIdx * static_cast<ssize_t>(chunkSize);
      thisChunkSize = static_cast<ssize_t>(chunkSize);
    } else {
      offset = transitionIdx * static_cast<ssize_t>(chunkSize) +
          (lastIdx - transitionIdx) * static_cast<ssize_t>(smallChunkSize);
      thisChunkSize = static_cast<ssize_t>(smallChunkSize);
    }
    Iter lastStart = start + offset;
    Iter lastEnd = lastStart + thisChunkSize;
    for (Iter it = lastStart; it != lastEnd; ++it) {
      f(*it);
    }
    tasks.wait();
  }
}
 
// Non-random-access iterators: pre-compute boundary iterators into SmallVector.
template <typename TaskSetT, typename Iter, typename F, typename IterCategory>
void for_each_n_schedule(
    TaskSetT& tasks,
    Iter start,
    F&& f,
    ssize_t numThreads,
    size_t chunkSize,
    ssize_t transitionIdx,
    size_t smallChunkSize,
    const ForEachOptions& options,
    IterCategory) {
  SmallVector<Iter, 64> boundaries;
  boundaries.reserve(static_cast<size_t>(numThreads) + 1);
  boundaries.push_back(start);
  for (ssize_t t = 0; t < numThreads; ++t) {
    size_t cs = (t < transitionIdx) ? chunkSize : smallChunkSize;
    Iter next = boundaries[t];
    std::advance(next, static_cast<ptrdiff_t>(cs));
    boundaries.push_back(next);
  }
 
  ssize_t numToSchedule = options.wait ? numThreads - 1 : numThreads;
 
  if (numToSchedule > 0) {
    tasks.scheduleBulk(static_cast<size_t>(numToSchedule), [&boundaries, &f](size_t idx) {
      return [s = boundaries[idx], e = boundaries[idx + 1], f]() {
        auto recurseInfo = PerPoolPerThreadInfo::parForRecurse();
        for (Iter it = s; it != e; ++it) {
          f(*it);
        }
      };
    });
  }
 
  if (options.wait) {
    for (Iter it = boundaries[numThreads - 1]; it != boundaries[numThreads]; ++it) {
      f(*it);
    }
    tasks.wait();
  }
}
 
} // namespace detail
 
template <typename TaskSetT, typename Iter, typename F>
DISPENSO_REQUIRES(ForEachFunc<F, Iter>)
void for_each_n(TaskSetT& tasks, Iter start, size_t n, F&& f, ForEachOptions options = {}) {
  // TODO(bbudge): With options.maxThreads, we might want to allow a small fanout factor in
  // recursive case?
  if (!n || !options.maxThreads || detail::PerPoolPerThreadInfo::isParForRecursive(&tasks.pool())) {
    for (size_t i = 0; i < n; ++i) {
      f(*start);
      ++start;
    }
    if (options.wait) {
      tasks.wait();
    }
    return;
  }
 
  // 0 indicates serial execution per API spec
  int32_t maxThreads = std::max<int32_t>(options.maxThreads, 1);
 
  ssize_t numThreads = std::min<ssize_t>(tasks.numPoolThreads() + options.wait, maxThreads);
  // Reduce threads used if they exceed work to be done.
  numThreads = std::min<ssize_t>(numThreads, n);
 
  auto chunking = detail::staticChunkSize(n, numThreads);
  size_t chunkSize = chunking.ceilChunkSize;
 
  bool perfectlyChunked = chunking.transitionTaskIndex == numThreads;
  ssize_t transitionIdx = chunking.transitionTaskIndex;
  size_t smallChunkSize = chunkSize - !perfectlyChunked;
 
  detail::for_each_n_schedule(
      tasks,
      start,
      std::forward<F>(f),
      numThreads,
      chunkSize,
      transitionIdx,
      smallChunkSize,
      options,
      typename std::iterator_traits<Iter>::iterator_category{});
}
 
template <typename Iter, typename F>
DISPENSO_REQUIRES(ForEachFunc<F, Iter>)
void for_each_n(Iter start, size_t n, F&& f, ForEachOptions options = {}) {
  TaskSet taskSet(globalThreadPool());
  options.wait = true;
  for_each_n(taskSet, start, n, std::forward<F>(f), options);
}
 
template <typename TaskSetT, typename Iter, typename F>
DISPENSO_REQUIRES(ForEachFunc<F, Iter>)
void for_each(TaskSetT& tasks, Iter start, Iter end, F&& f, ForEachOptions options = {}) {
  for_each_n(tasks, start, std::distance(start, end), std::forward<F>(f), options);
}
 
template <typename Iter, typename F>
DISPENSO_REQUIRES(ForEachFunc<F, Iter>)
void for_each(Iter start, Iter end, F&& f, ForEachOptions options = {}) {
  for_each_n(start, std::distance(start, end), std::forward<F>(f), options);
}
 
// TODO(bbudge): Implement ranges versions for these in C++20 (currently not in an env where this
// can be tested)
 
} // namespace dispenso