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 <dispenso/detail/per_thread_info.h>
#include <dispenso/task_set.h>
 
namespace dispenso {
 
struct ForEachOptions {
  uint32_t maxThreads = std::numeric_limits<int32_t>::max();
  bool wait = true;
};
 
template <typename TaskSetT, typename Iter, typename F>
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;
 
  // (!perfectlyChunked) ? chunking.transitionTaskIndex : numThreads - 1;
  ssize_t firstLoopLen = chunking.transitionTaskIndex - perfectlyChunked;
 
  ssize_t t;
  for (t = 0; t < firstLoopLen; ++t) {
    Iter next = start;
    std::advance(next, chunkSize);
    tasks.schedule([start, next, f]() {
      auto recurseInfo = detail::PerPoolPerThreadInfo::parForRecurse();
      for (Iter it = start; it != next; ++it) {
        f(*it);
      }
    });
    start = next;
  }
 
  // Reduce the remaining chunk sizes by 1.
  chunkSize -= !perfectlyChunked;
  // Finish submitting all but the last item.
  for (; t < numThreads - 1; ++t) {
    Iter next = start;
    std::advance(next, chunkSize);
    tasks.schedule([start, next, f]() {
      auto recurseInfo = detail::PerPoolPerThreadInfo::parForRecurse();
      for (Iter it = start; it != next; ++it) {
        f(*it);
      }
    });
    start = next;
  }
 
  Iter end = start;
  std::advance(end, chunkSize);
 
  if (options.wait) {
    for (Iter it = start; it != end; ++it) {
      f(*it);
    }
    tasks.wait();
  } else {
    tasks.schedule(
        [start, end, f]() {
          auto recurseInfo = detail::PerPoolPerThreadInfo::parForRecurse();
          for (Iter it = start; it != end; ++it) {
            f(*it);
          }
        },
        ForceQueuingTag());
  }
}
 
template <typename Iter, typename F>
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>
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>
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