thread.h

Go to the documentation of this file.

#pragma once
 
#include <citor/cancellation.h>
#include <citor/hints.h>
#include <citor/thread_pool.h>
 
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <mutex>
#include <span>
#include <thread>
#include <unordered_map>
#include <utility>
#include <vector>
 
#include "clustering/math/detail/coherence_cache.h"
 
namespace clustering::math {

using OwnedPool = citor::ThreadPool;

inline OwnedPool &sharedPool(std::size_t nJobs);

inline std::size_t clampedJobCount(std::size_t nJobs) noexcept {
  if (nJobs == 0) {
    const std::size_t hw = std::thread::hardware_concurrency();
    return hw == 0 ? std::size_t{1} : hw;
  }
  return nJobs;
}

inline bool shouldSpawnPool(std::size_t totalOps, std::size_t nJobs,
                            std::size_t minOpsPerWorker = std::size_t{1} << 15) noexcept {
  if (nJobs <= 1) {
    return false;
  }
  return (totalOps / nJobs) >= minOpsPerWorker;
}

struct Pool {
  OwnedPool *pool = nullptr;

  [[nodiscard]] std::size_t workerCount() const noexcept {
    if (pool == nullptr) {
      return std::size_t{1};
    }
    return pool->participants();
  }

  [[nodiscard]] static std::size_t workerIndex() noexcept {
    return citor::ThreadPool::workerIndex();
  }

  [[nodiscard]] bool shouldParallelize(std::size_t totalWork, std::size_t minChunk,
                                       std::size_t minTasksPerWorker = 2) const noexcept {
    if (pool == nullptr || minChunk == 0) {
      return false;
    }
    return (totalWork / minChunk) >= (workerCount() * minTasksPerWorker);
  }

  [[nodiscard]] bool shouldParallelizeWork(std::size_t totalOps,
                                           std::size_t minOpsPerWorker = std::size_t{1}
                                                                         << 15) const noexcept {
    if (pool == nullptr) {
      return false;
    }
    return (totalOps / workerCount()) >= minOpsPerWorker;
  }

  [[nodiscard]] std::size_t stealBlocks(std::size_t n,
                                        std::size_t minRowsPerBlock = 256) const noexcept {
    const std::size_t workers = workerCount();
    if (workers <= 1 || n == 0) {
      return workers;
    }
    const std::size_t byRows =
        (minRowsPerBlock == 0) ? n : std::max<std::size_t>(1, n / minRowsPerBlock);
    const std::size_t blocks = std::min(workers * 8, byRows);
    return std::max(blocks, workers);
  }

  template <class HintsT = citor::HintsDefaults, class FnA, class FnB>
  void forkJoin2(FnA &&a, FnB &&b) const {
    if (pool == nullptr) {
      a();
      b();
      return;
    }
    pool->template forkJoin<HintsT>(std::forward<FnA>(a), std::forward<FnB>(b));
  }

  template <class HintsT = citor::HintsDefaults, class Body>
  void parallelForBlocks(std::size_t first, std::size_t last, std::size_t numBlocks, Body body) {
    if (pool == nullptr || first >= last) {
      body(first, last);
      return;
    }
    if (numBlocks != 0) {
      parallelForExactBlocks<HintsT>(first, last, numBlocks, std::move(body));
      return;
    }
    pool->template parallelFor<HintsT>(first, last, body);
  }

  template <class HintsT = citor::HintsDefaults, class Body>
  void parallelForExactBlocks(std::size_t first, std::size_t last, std::size_t numBlocks,
                              Body body) {
    if (first >= last || numBlocks == 0) {
      return;
    }
    if (pool == nullptr || numBlocks == 1) {
      body(first, last);
      return;
    }
    const std::size_t span = last - first;
    auto sliceLo = [first, span, numBlocks](std::size_t s) noexcept {
      return first + ((span * s) / numBlocks);
    };
    pool->template parallelFor<HintsT>(std::size_t{0}, numBlocks,
                                       [&](std::size_t loSlot, std::size_t hiSlot) {
                                         for (std::size_t s = loSlot; s < hiSlot; ++s) {
                                           const std::size_t blockLo = sliceLo(s);
                                           const std::size_t blockHi = sliceLo(s + 1);
                                           body(blockLo, blockHi);
                                         }
                                       });
  }

  template <class HintsT = citor::HintsDefaults, class Body>
  void parallelForExactBlocksWithSlot(std::size_t first, std::size_t last, std::size_t numBlocks,
                                      Body body) {
    if (first >= last || numBlocks == 0) {
      return;
    }
    if (pool == nullptr || numBlocks == 1) {
      body(first, last, std::size_t{0});
      return;
    }
    const std::size_t span = last - first;
    auto sliceLo = [first, span, numBlocks](std::size_t s) noexcept {
      return first + ((span * s) / numBlocks);
    };
    pool->template parallelFor<HintsT>(std::size_t{0}, numBlocks,
                                       [&](std::size_t loSlot, std::size_t hiSlot) {
                                         for (std::size_t s = loSlot; s < hiSlot; ++s) {
                                           body(sliceLo(s), sliceLo(s + 1), s);
                                         }
                                       });
  }

  template <class HintsT = citor::HintsDefaults, class Body>
  void parallelForChunks(std::size_t numChunks, Body body) {
    if (pool == nullptr || numChunks == 0) {
      for (std::size_t c = 0; c < numChunks; ++c) {
        body(c);
      }
      return;
    }
    pool->template parallelFor<HintsT>(std::size_t{0}, numChunks,
                                       [&](std::size_t lo, std::size_t hi) {
                                         for (std::size_t c = lo; c < hi; ++c) {
                                           body(c);
                                         }
                                       });
  }

  template <class HintsT = citor::HintsDefaults, class T, class Map, class Combine>
  [[nodiscard]] T parallelReduce(std::size_t first, std::size_t last, T init, Map map,
                                 Combine combine) {
    if (first >= last) {
      return init;
    }
    if (pool == nullptr) {
      return combine(std::move(init), map(first, last));
    }
    return pool->template parallelReduce<HintsT>(first, last, std::move(init), std::move(map),
                                                 std::move(combine));
  }

  template <class HintsT = citor::HintsDefaults, class Phase>
  void parallelRunPlex(std::size_t nPhases, std::size_t n, Phase phaseFn) {
    auto noPrePhase = [](std::size_t /*phaseIdx*/) noexcept {};
    parallelRunPlex<HintsT>(nPhases, n, std::move(phaseFn), noPrePhase);
  }

  template <class HintsT = citor::HintsDefaults, class T, class BodyFn, class PrefixFn>
  T parallelScan(std::size_t n, T identity, BodyFn body, PrefixFn prefix) {
    if (n == 0) {
      return identity;
    }
    if (pool == nullptr) {
      T partial = body(std::size_t{0}, std::size_t{0}, n, identity, static_cast<T *>(nullptr));
      return prefix(std::move(identity), std::move(partial));
    }
    return pool->template parallelScan<HintsT>(n, std::move(identity), std::move(body),
                                               std::move(prefix));
  }

  template <class HintsT = citor::HintsDefaults, class T, class PrefixFn>
  [[nodiscard]] T inclusiveScan(std::span<const T> in, std::span<T> out, T identity,
                                PrefixFn prefix) {
    if (pool == nullptr) {
      T acc = identity;
      for (std::size_t i = 0; i < in.size(); ++i) {
        acc = prefix(acc, in[i]);
        out[i] = acc;
      }
      return acc;
    }
    return pool->template inclusiveScan<HintsT>(in, out, std::move(identity), std::move(prefix));
  }

  template <class HintsT = citor::HintsDefaults, class Phase, class PrePhase>
  void parallelRunPlex(std::size_t nPhases, std::size_t n, Phase phaseFn, PrePhase prePhaseFn,
                       citor::CancellationToken tok = citor::CancellationToken{}) {
    if (nPhases == 0) {
      return;
    }
    if (pool == nullptr) {
      // Single-slot inline emulation when no pool is attached; mirrors the backend's
      // phase-boundary token check.
      for (std::size_t p = 0; p < nPhases; ++p) {
        if (tok.stop_requested()) {
          return;
        }
        prePhaseFn(p);
        phaseFn(p, std::uint32_t{0}, std::size_t{0}, n, static_cast<void *>(nullptr));
      }
      return;
    }
    pool->template runPlex<HintsT>(nPhases, n, std::forward<Phase>(phaseFn),
                                   std::forward<PrePhase>(prePhaseFn), std::move(tok));
  }
};
 
inline OwnedPool &sharedPool(std::size_t nJobs) {
  const std::size_t effective = clampedJobCount(nJobs);
  // Inline-static map: shared across translation units thanks to the inline keyword on the
  // enclosing function. Mutex protects the registry against first-time-init races; lookups
  // are O(1) and gated by Python's GIL or other caller-side serialization in practice.
  static std::unordered_map<std::size_t, std::unique_ptr<OwnedPool>> registry;
  static std::mutex registryMutex;
  const std::scoped_lock guard{registryMutex};
  auto &slot = registry[effective];
  if (!slot) {
    // Seed citor's coherence-probe cache from disk before constructing the pool
    // so a short-lived or single-fit process skips the live inter-core
    // calibration. On a cold cache the ctor runs the probe; persist the fresh
    // result so the next process replays it.
    const bool seeded = detail::importPersistedCoherenceProbe(effective);
    slot = std::make_unique<OwnedPool>(effective);
    if (!seeded) {
      detail::exportPersistedCoherenceProbe(*slot, effective);
    }
  }
  return *slot;
}
 
} // namespace clustering::math