tor-browser

thread_pool.h (64831B)

---

// Copyright 2023 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Modified from BSD-licensed code
// Copyright (c) the JPEG XL Project Authors. All rights reserved.
// See https://github.com/libjxl/libjxl/blob/main/LICENSE.

#ifndef HIGHWAY_HWY_CONTRIB_THREAD_POOL_THREAD_POOL_H_
#define HIGHWAY_HWY_CONTRIB_THREAD_POOL_THREAD_POOL_H_

#include <stddef.h>
#include <stdint.h>
#include <stdio.h>  // snprintf
#include <string.h>

#include <array>
#include <atomic>
#include <string>
#include <thread>  // NOLINT
#include <vector>

#include "hwy/aligned_allocator.h"  // HWY_ALIGNMENT
#include "hwy/auto_tune.h"
#include "hwy/base.h"
#include "hwy/cache_control.h"  // Pause
#include "hwy/contrib/thread_pool/futex.h"
#include "hwy/contrib/thread_pool/spin.h"
#include "hwy/contrib/thread_pool/topology.h"
#include "hwy/profiler.h"
#include "hwy/stats.h"
#include "hwy/timer.h"

#if HWY_OS_APPLE
#include <AvailabilityMacros.h>
#endif

#if PROFILER_ENABLED
#include <algorithm>  // std::sort

#include "hwy/bit_set.h"
#endif

namespace hwy {

// Sets the name of the current thread to the format string `format`, which must
// include %d for `thread`. Currently only implemented for pthreads (*nix and
// OSX); Windows involves throwing an exception.
static inline void SetThreadName(const char* format, int thread) {
 char buf[16] = {};  // Linux limit, including \0
 const int chars_written = snprintf(buf, sizeof(buf), format, thread);
 HWY_ASSERT(0 < chars_written &&
            chars_written <= static_cast<int>(sizeof(buf) - 1));

#if (HWY_OS_LINUX && (!defined(__ANDROID__) || __ANDROID_API__ >= 19)) || \
   HWY_OS_FREEBSD
 // Note that FreeBSD pthread_set_name_np does not return a value (#2669).
 HWY_ASSERT(0 == pthread_setname_np(pthread_self(), buf));
#elif HWY_OS_APPLE && (MAC_OS_X_VERSION_MIN_REQUIRED >= 1060)
 // Different interface: single argument, current thread only.
 HWY_ASSERT(0 == pthread_setname_np(buf));
#elif defined(__EMSCRIPTEN__)
 emscripten_set_thread_name(pthread_self(), buf);
#else
 (void)format;
 (void)thread;
#endif
}

// Whether workers should block or spin.
enum class PoolWaitMode : uint8_t { kBlock = 1, kSpin };

enum class Exit : uint32_t { kNone, kLoop, kThread };

// Upper bound on non-empty `ThreadPool` (single-worker pools do not count).
// Turin has 16 clusters. Add one for the across-cluster pool.
HWY_INLINE_VAR constexpr size_t kMaxClusters = 32 + 1;

// Use the last slot so that `PoolWorkerMapping` does not have to know the
// total number of clusters.
HWY_INLINE_VAR constexpr size_t kAllClusters = kMaxClusters - 1;

// Argument to `ThreadPool`: how to map local worker_idx to global.
class PoolWorkerMapping {
public:
 // Backward-compatible mode: returns local worker index.
 PoolWorkerMapping() : cluster_idx_(0), max_cluster_workers_(0) {}
 PoolWorkerMapping(size_t cluster_idx, size_t max_cluster_workers)
     : cluster_idx_(cluster_idx), max_cluster_workers_(max_cluster_workers) {
   HWY_DASSERT(cluster_idx <= kAllClusters);
   // Only use this ctor for the new global worker index mode. If this were
   // zero, we would still return local indices.
   HWY_DASSERT(max_cluster_workers != 0);
 }

 size_t ClusterIdx() const { return cluster_idx_; }
 size_t MaxClusterWorkers() const { return max_cluster_workers_; }

 // Returns global_idx, or unchanged local worker_idx if default-constructed.
 size_t operator()(size_t worker_idx) const {
   if (cluster_idx_ == kAllClusters) {
     const size_t cluster_idx = worker_idx;
     HWY_DASSERT(cluster_idx < kAllClusters);
     // First index within the N-th cluster. The main thread is the first.
     return cluster_idx * max_cluster_workers_;
   }
   HWY_DASSERT(max_cluster_workers_ == 0 || worker_idx < max_cluster_workers_);
   return cluster_idx_ * max_cluster_workers_ + worker_idx;
 }

private:
 size_t cluster_idx_;
 size_t max_cluster_workers_;
};

namespace pool {

#ifndef HWY_POOL_VERBOSITY
#define HWY_POOL_VERBOSITY 0
#endif

static constexpr int kVerbosity = HWY_POOL_VERBOSITY;

// Some CPUs already have more than this many threads, but rather than one
// large pool, we assume applications create multiple pools, ideally per
// cluster (cores sharing a cache), because this improves locality and barrier
// latency. In that case, this is a generous upper bound.
static constexpr size_t kMaxThreads = 127;

// Generates a random permutation of [0, size). O(1) storage.
class ShuffledIota {
public:
 ShuffledIota() : coprime_(1) {}  // for Worker
 explicit ShuffledIota(uint32_t coprime) : coprime_(coprime) {}

 // Returns the next after `current`, using an LCG-like generator.
 uint32_t Next(uint32_t current, const Divisor64& divisor) const {
   HWY_DASSERT(current < divisor.GetDivisor());
   // (coprime * i + current) % size, see https://lemire.me/blog/2017/09/18/.
   return static_cast<uint32_t>(divisor.Remainder(current + coprime_));
 }

 // Returns true if a and b have no common denominator except 1. Based on
 // binary GCD. Assumes a and b are nonzero. Also used in tests.
 static bool CoprimeNonzero(uint32_t a, uint32_t b) {
   const size_t trailing_a = Num0BitsBelowLS1Bit_Nonzero32(a);
   const size_t trailing_b = Num0BitsBelowLS1Bit_Nonzero32(b);
   // If both have at least one trailing zero, they are both divisible by 2.
   if (HWY_MIN(trailing_a, trailing_b) != 0) return false;

   // If one of them has a trailing zero, shift it out.
   a >>= trailing_a;
   b >>= trailing_b;

   for (;;) {
     // Swap such that a >= b.
     const uint32_t tmp_a = a;
     a = HWY_MAX(tmp_a, b);
     b = HWY_MIN(tmp_a, b);

     // When the smaller number is 1, they were coprime.
     if (b == 1) return true;

     a -= b;
     // a == b means there was a common factor, so not coprime.
     if (a == 0) return false;
     a >>= Num0BitsBelowLS1Bit_Nonzero32(a);
   }
 }

 // Returns another coprime >= `start`, or 1 for small `size`.
 // Used to seed independent ShuffledIota instances.
 static uint32_t FindAnotherCoprime(uint32_t size, uint32_t start) {
   if (size <= 2) {
     return 1;
   }

   // Avoids even x for even sizes, which are sure to be rejected.
   const uint32_t inc = (size & 1) ? 1 : 2;

   for (uint32_t x = start | 1; x < start + size * 16; x += inc) {
     if (CoprimeNonzero(x, static_cast<uint32_t>(size))) {
       return x;
     }
   }

   HWY_UNREACHABLE;
 }

 uint32_t coprime_;
};

// 'Policies' suitable for various worker counts and locality. To define a
// new class, add an enum and update `ToString` plus `CallWithConfig`. The
// enumerators must be contiguous so we can iterate over them.
enum class WaitType : uint8_t {
 kBlock,
 kSpin1,
 kSpinSeparate,
 kSentinel  // Must be last.
};

// For printing which is in use.
static inline const char* ToString(WaitType type) {
 switch (type) {
   case WaitType::kBlock:
     return "Block";
   case WaitType::kSpin1:
     return "Single";
   case WaitType::kSpinSeparate:
     return "Separate";
   case WaitType::kSentinel:
     return nullptr;
 }
}

// Parameters governing the main and worker thread behavior. Can be updated at
// runtime via `SetWaitMode`, which calls `SendConfig`. Both have copies which
// are carefully synchronized. 32 bits leave room for two future fields.
// 64 bits would also be fine because this does not go through futex.
struct Config {  // 4 bytes
 static std::vector<Config> AllCandidates(PoolWaitMode wait_mode) {
   std::vector<Config> candidates;

   if (wait_mode == PoolWaitMode::kSpin) {
     std::vector<SpinType> spin_types;
     spin_types.reserve(2);
     spin_types.push_back(DetectSpin());
     // Monitor-based spin may be slower, so also try Pause.
     if (spin_types[0] != SpinType::kPause) {
       spin_types.push_back(SpinType::kPause);
     }

     // All except `kBlock`.
     std::vector<WaitType> wait_types;
     for (size_t wait = 0;; ++wait) {
       const WaitType wait_type = static_cast<WaitType>(wait);
       if (wait_type == WaitType::kSentinel) break;
       if (wait_type != WaitType::kBlock) wait_types.push_back(wait_type);
     }

     candidates.reserve(spin_types.size() * wait_types.size());
     for (const SpinType spin_type : spin_types) {
       for (const WaitType wait_type : wait_types) {
         candidates.emplace_back(spin_type, wait_type);
       }
     }
   } else {
     // kBlock does not use spin, so there is only one candidate.
     candidates.emplace_back(SpinType::kPause, WaitType::kBlock);
   }

   return candidates;
 }

 std::string ToString() const {
   char buf[128];
   snprintf(buf, sizeof(buf), "%-14s %-9s", hwy::ToString(spin_type),
            pool::ToString(wait_type));
   return buf;
 }

 Config(SpinType spin_type_in, WaitType wait_type_in)
     : spin_type(spin_type_in), wait_type(wait_type_in) {}
 // Workers initially spin until ThreadPool sends them their actual config.
 Config() : Config(SpinType::kPause, WaitType::kSpinSeparate) {}

 SpinType spin_type;
 WaitType wait_type;
 HWY_MEMBER_VAR_MAYBE_UNUSED uint8_t reserved[2];
};
static_assert(sizeof(Config) == 4, "");

#if PROFILER_ENABLED

// Accumulates timings and stats from main thread and workers.
class Stats {
 // Up to `HWY_ALIGNMENT / 8` slots/offsets, passed to `PerThread`.
 static constexpr size_t kDWait = 0;
 static constexpr size_t kWaitReps = 1;
 static constexpr size_t kTBeforeRun = 2;
 static constexpr size_t kDRun = 3;
 static constexpr size_t kTasksStatic = 4;
 static constexpr size_t kTasksDynamic = 5;
 static constexpr size_t kTasksStolen = 6;
 static constexpr size_t kDFuncStatic = 7;
 static constexpr size_t kDFuncDynamic = 8;
 static constexpr size_t kSentinel = 9;

public:
 Stats() {
   for (size_t thread_idx = 0; thread_idx < kMaxThreads; ++thread_idx) {
     for (size_t offset = 0; offset < kSentinel; ++offset) {
       PerThread(thread_idx, offset) = 0;
     }
   }
   Reset();
 }

 // Called by the N lowest-indexed workers that got one of the N tasks, which
 // includes the main thread because its index is 0.
 // `d_*` denotes "difference" (of timestamps) and thus also duration.
 void NotifyRunStatic(size_t worker_idx, timer::Ticks d_func) {
   if (worker_idx == 0) {  // main thread
     num_run_static_++;
     sum_tasks_static_++;
     sum_d_func_static_ += d_func;
   } else {
     const size_t thread_idx = worker_idx - 1;
     // Defer the sums until `NotifyMainRun` to avoid atomic RMW.
     PerThread(thread_idx, kTasksStatic)++;
     PerThread(thread_idx, kDFuncStatic) += d_func;
   }
 }

 // Called by all workers, including the main thread, regardless of whether
 // they actually stole or even ran a task.
 void NotifyRunDynamic(size_t worker_idx, size_t tasks, size_t stolen,
                       timer::Ticks d_func) {
   if (worker_idx == 0) {  // main thread
     num_run_dynamic_++;
     sum_tasks_dynamic_ += tasks;
     sum_tasks_stolen_ += stolen;
     sum_d_func_dynamic_ += d_func;
   } else {
     const size_t thread_idx = worker_idx - 1;
     // Defer the sums until `NotifyMainRun` to avoid atomic RMW.
     PerThread(thread_idx, kTasksDynamic) += tasks;
     PerThread(thread_idx, kTasksStolen) += stolen;
     PerThread(thread_idx, kDFuncDynamic) += d_func;
   }
 }

 // Called concurrently by non-main worker threads after their `WorkerRun` and
 // before the barrier.
 void NotifyThreadRun(size_t worker_idx, timer::Ticks d_wait, size_t wait_reps,
                      timer::Ticks t_before_run, timer::Ticks d_run) {
   HWY_DASSERT(worker_idx != 0);  // Not called by main thread.
   const size_t thread_idx = worker_idx - 1;
   HWY_DASSERT(PerThread(thread_idx, kDWait) == 0);
   HWY_DASSERT(PerThread(thread_idx, kWaitReps) == 0);
   HWY_DASSERT(PerThread(thread_idx, kTBeforeRun) == 0);
   HWY_DASSERT(PerThread(thread_idx, kDRun) == 0);
   PerThread(thread_idx, kDWait) = d_wait;
   PerThread(thread_idx, kWaitReps) = wait_reps;
   PerThread(thread_idx, kTBeforeRun) = t_before_run;  // For wake latency.
   PerThread(thread_idx, kDRun) = d_run;
 }

 // Called by the main thread after the barrier, whose store-release and
 // load-acquire publishes all prior writes. Note: only the main thread can
 // store `after_barrier`. If workers did, which by definition happens after
 // the barrier, then they would race with this function's reads.
 void NotifyMainRun(size_t num_threads, timer::Ticks t_before_wake,
                    timer::Ticks d_wake, timer::Ticks d_main_run,
                    timer::Ticks d_barrier) {
   HWY_DASSERT(num_threads <= kMaxThreads);

   timer::Ticks min_d_run = ~timer::Ticks{0};
   timer::Ticks max_d_run = 0;
   timer::Ticks sum_d_run = 0;
   for (size_t thread_idx = 0; thread_idx < num_threads; ++thread_idx) {
     sum_tasks_static_ += PerThread(thread_idx, kTasksStatic);
     sum_tasks_dynamic_ += PerThread(thread_idx, kTasksDynamic);
     sum_tasks_stolen_ += PerThread(thread_idx, kTasksStolen);
     sum_d_func_static_ += PerThread(thread_idx, kDFuncStatic);
     sum_d_func_dynamic_ += PerThread(thread_idx, kDFuncDynamic);
     sum_d_wait_ += PerThread(thread_idx, kDWait);
     sum_wait_reps_ += PerThread(thread_idx, kWaitReps);
     const timer::Ticks d_thread_run = PerThread(thread_idx, kDRun);
     min_d_run = HWY_MIN(min_d_run, d_thread_run);
     max_d_run = HWY_MAX(max_d_run, d_thread_run);
     sum_d_run += d_thread_run;
     const timer::Ticks t_before_run = PerThread(thread_idx, kTBeforeRun);

     for (size_t offset = 0; offset < kSentinel; ++offset) {
       PerThread(thread_idx, offset) = 0;
     }

     HWY_DASSERT(t_before_run != 0);
     const timer::Ticks d_latency = t_before_run - t_before_wake;
     sum_wake_latency_ += d_latency;
     max_wake_latency_ = HWY_MAX(max_wake_latency_, d_latency);
   }
   const double inv_avg_d_run =
       static_cast<double>(num_threads) / static_cast<double>(sum_d_run);
   // Ratios of min and max run times to the average, for this pool.Run.
   const double r_min = static_cast<double>(min_d_run) * inv_avg_d_run;
   const double r_max = static_cast<double>(max_d_run) * inv_avg_d_run;

   num_run_++;  // `num_run_*` are incremented by `NotifyRun*`.
   sum_d_run_ += sum_d_run;
   sum_r_min_ += r_min;  // For average across all pool.Run.
   sum_r_max_ += r_max;

   sum_d_wake_ += d_wake;  // `*wake_latency_` are updated above.
   sum_d_barrier_ += d_barrier;

   sum_d_run_ += d_main_run;
   sum_d_run_main_ += d_main_run;
 }

 void PrintAndReset(size_t num_threads, timer::Ticks d_thread_lifetime_ticks) {
   // This is unconditionally called via `ProfilerFunc`. If the pool was unused
   // in this invocation, skip it.
   if (num_run_ == 0) return;
   HWY_ASSERT(num_run_ == num_run_static_ + num_run_dynamic_);

   const double d_func_static = Seconds(sum_d_func_static_);
   const double d_func_dynamic = Seconds(sum_d_func_dynamic_);
   const double sum_d_run = Seconds(sum_d_run_);
   const double func_div_run = (d_func_static + d_func_dynamic) / sum_d_run;
   if (!(0.95 <= func_div_run && func_div_run <= 1.0)) {
     HWY_WARN("Func time %f should be similar to total run %f.",
              d_func_static + d_func_dynamic, sum_d_run);
   }
   const double sum_d_run_main = Seconds(sum_d_run_main_);
   const double max_wake_latency = Seconds(max_wake_latency_);
   const double sum_d_wait = Seconds(sum_d_wait_);
   const double d_thread_lifetime = Seconds(d_thread_lifetime_ticks);

   const double inv_run = 1.0 / static_cast<double>(num_run_);
   const auto per_run = [inv_run](double sum) { return sum * inv_run; };
   const double avg_d_wake = per_run(Seconds(sum_d_wake_));
   const double avg_wake_latency = per_run(Seconds(sum_wake_latency_));
   const double avg_d_wait = per_run(sum_d_wait);
   const double avg_wait_reps = per_run(static_cast<double>(sum_wait_reps_));
   const double avg_d_barrier = per_run(Seconds(sum_d_barrier_));
   const double avg_r_min = per_run(sum_r_min_);
   const double avg_r_max = per_run(sum_r_max_);

   const size_t num_workers = 1 + num_threads;
   const double avg_tasks_static =
       Avg(sum_tasks_static_, num_run_static_ * num_workers);
   const double avg_tasks_dynamic =
       Avg(sum_tasks_dynamic_, num_run_dynamic_ * num_workers);
   const double avg_steals =
       Avg(sum_tasks_stolen_, num_run_dynamic_ * num_workers);
   const double avg_d_run = sum_d_run / num_workers;

   const double pc_wait = sum_d_wait / d_thread_lifetime * 100.0;
   const double pc_run = sum_d_run / d_thread_lifetime * 100.0;
   const double pc_main = sum_d_run_main / avg_d_run * 100.0;

   const auto us = [](double sec) { return sec * 1E6; };
   const auto ns = [](double sec) { return sec * 1E9; };
   printf(
       "%3zu: %5d x %.2f/%5d x %4.1f tasks, %.2f steals; "
       "wake %7.3f ns, latency %6.3f < %7.3f us, barrier %7.3f us; "
       "wait %.1f us (%6.0f reps, %4.1f%%), balance %4.1f%%-%5.1f%%, "
       "func: %6.3f + %7.3f, "
       "%.1f%% of thread time %7.3f s; main:worker %5.1f%%\n",
       num_threads, num_run_static_, avg_tasks_static, num_run_dynamic_,
       avg_tasks_dynamic, avg_steals, ns(avg_d_wake), us(avg_wake_latency),
       us(max_wake_latency), us(avg_d_barrier), us(avg_d_wait), avg_wait_reps,
       pc_wait, avg_r_min * 100.0, avg_r_max * 100.0, d_func_static,
       d_func_dynamic, pc_run, d_thread_lifetime, pc_main);

   Reset(num_threads);
 }

 void Reset(size_t num_threads = kMaxThreads) {
   num_run_ = 0;
   num_run_static_ = 0;
   num_run_dynamic_ = 0;

   sum_tasks_stolen_ = 0;
   sum_tasks_static_ = 0;
   sum_tasks_dynamic_ = 0;

   sum_d_wake_ = 0;
   sum_wake_latency_ = 0;
   max_wake_latency_ = 0;
   sum_d_wait_ = 0;
   sum_wait_reps_ = 0;
   sum_d_barrier_ = 0;

   sum_d_func_static_ = 0;
   sum_d_func_dynamic_ = 0;
   sum_r_min_ = 0.0;
   sum_r_max_ = 0.0;
   sum_d_run_ = 0;
   sum_d_run_main_ = 0;
   // ctor and `NotifyMainRun` already reset `PerThread`.
 }

private:
 template <typename T>
 static double Avg(T sum, size_t div) {
   return div == 0 ? 0.0 : static_cast<double>(sum) / static_cast<double>(div);
 }

 static constexpr size_t kU64PerLine = HWY_ALIGNMENT / sizeof(uint64_t);

 uint64_t& PerThread(size_t thread_idx, size_t offset) {
   HWY_DASSERT(thread_idx < kMaxThreads);
   HWY_DASSERT(offset < kSentinel);
   return per_thread_[thread_idx * kU64PerLine + offset];
 }

 int32_t num_run_;
 int32_t num_run_static_;
 int32_t num_run_dynamic_;

 int32_t sum_tasks_stolen_;
 int64_t sum_tasks_static_;
 int64_t sum_tasks_dynamic_;

 timer::Ticks sum_d_wake_;
 timer::Ticks sum_wake_latency_;
 timer::Ticks max_wake_latency_;
 timer::Ticks sum_d_wait_;
 uint64_t sum_wait_reps_;
 timer::Ticks sum_d_barrier_;

 timer::Ticks sum_d_func_static_;
 timer::Ticks sum_d_func_dynamic_;
 double sum_r_min_;
 double sum_r_max_;
 timer::Ticks sum_d_run_;
 timer::Ticks sum_d_run_main_;

 // One cache line per pool thread to avoid false sharing.
 uint64_t per_thread_[kMaxThreads * kU64PerLine];
};
// Enables shift rather than multiplication.
static_assert(sizeof(Stats) == (kMaxThreads + 1) * HWY_ALIGNMENT, "Wrong size");

// Non-power of two to avoid 2K aliasing.
HWY_INLINE_VAR constexpr size_t kMaxCallers = 60;

// Per-caller stats, stored in `PerCluster`.
class CallerAccumulator {
public:
 bool Any() const { return calls_ != 0; }

 void Add(size_t tasks, size_t workers, bool is_root, timer::Ticks wait_before,
          timer::Ticks elapsed) {
   calls_++;
   root_ += is_root;
   workers_ += workers;
   min_tasks_ = HWY_MIN(min_tasks_, tasks);
   max_tasks_ = HWY_MAX(max_tasks_, tasks);
   tasks_ += tasks;
   wait_before_ += wait_before;
   elapsed_ += elapsed;
 }

 void AddFrom(const CallerAccumulator& other) {
   calls_ += other.calls_;
   root_ += other.root_;
   workers_ += other.workers_;
   min_tasks_ = HWY_MIN(min_tasks_, other.min_tasks_);
   max_tasks_ = HWY_MAX(max_tasks_, other.max_tasks_);
   tasks_ += other.tasks_;
   wait_before_ += other.wait_before_;
   elapsed_ += other.elapsed_;
 }

 bool operator>(const CallerAccumulator& other) const {
   return elapsed_ > other.elapsed_;
 }

 void PrintAndReset(const char* caller, size_t active_clusters) {
   if (!Any()) return;
   HWY_ASSERT(root_ <= calls_);
   const double inv_calls = 1.0 / static_cast<double>(calls_);
   const double pc_root = static_cast<double>(root_) * inv_calls * 100.0;
   const double avg_workers = static_cast<double>(workers_) * inv_calls;
   const double avg_tasks = static_cast<double>(tasks_) * inv_calls;
   const double avg_tasks_per_worker = avg_tasks / avg_workers;
   const double inv_freq = 1.0 / platform::InvariantTicksPerSecond();
   const double sum_wait_before = static_cast<double>(wait_before_) * inv_freq;
   const double avg_wait_before =
       root_ ? sum_wait_before / static_cast<double>(root_) : 0.0;
   const double elapsed = static_cast<double>(elapsed_) * inv_freq;
   const double avg_elapsed = elapsed * inv_calls;
   const double task_len = avg_elapsed / avg_tasks_per_worker;
   printf(
       "%40s: %7.0f x (%3.0f%%) %2zu clusters, %4.1f workers @ "
       "%5.1f tasks (%5u-%5u), "
       "%5.0f us wait, %6.1E us run (task len %6.1E us), total %6.2f s\n",
       caller, static_cast<double>(calls_), pc_root, active_clusters,
       avg_workers, avg_tasks_per_worker, static_cast<uint32_t>(min_tasks_),
       static_cast<uint32_t>(max_tasks_), avg_wait_before * 1E6,
       avg_elapsed * 1E6, task_len * 1E6, elapsed);
   *this = CallerAccumulator();
 }

 // For the grand total, only print calls and elapsed because averaging the
 // the other stats is not very useful. No need to reset because this is called
 // on a temporary.
 void PrintTotal() {
   if (!Any()) return;
   HWY_ASSERT(root_ <= calls_);
   const double elapsed =
       static_cast<double>(elapsed_) / platform::InvariantTicksPerSecond();
   printf("TOTAL: %7.0f x run %6.2f s\n", static_cast<double>(calls_),
          elapsed);
 }

private:
 int64_t calls_ = 0;
 int64_t root_ = 0;
 uint64_t workers_ = 0;
 uint64_t min_tasks_ = ~uint64_t{0};
 uint64_t max_tasks_ = 0;
 uint64_t tasks_ = 0;
 // both are wall time for root Run, otherwise CPU time.
 timer::Ticks wait_before_ = 0;
 timer::Ticks elapsed_ = 0;
};
static_assert(sizeof(CallerAccumulator) == 64, "");

class PerCluster {
public:
 CallerAccumulator& Get(size_t caller_idx) {
   HWY_DASSERT(caller_idx < kMaxCallers);
   callers_.Set(caller_idx);
   return accumulators_[caller_idx];
 }

 template <class Func>
 void ForeachCaller(Func&& func) {
   callers_.Foreach([&](size_t caller_idx) {
     func(caller_idx, accumulators_[caller_idx]);
   });
 }

 // Returns indices (required for `StringTable::Name`) in descending order of
 // elapsed time.
 std::vector<size_t> Sorted() {
   std::vector<size_t> vec;
   vec.reserve(kMaxCallers);
   ForeachCaller([&](size_t caller_idx, CallerAccumulator&) {
     vec.push_back(caller_idx);
   });
   std::sort(vec.begin(), vec.end(), [&](size_t a, size_t b) {
     return accumulators_[a] > accumulators_[b];
   });
   return vec;
 }

 // Caller takes care of resetting `accumulators_`.
 void ResetBits() { callers_ = hwy::BitSet<kMaxCallers>(); }

private:
 CallerAccumulator accumulators_[kMaxCallers];
 hwy::BitSet<kMaxCallers> callers_;
};

// Type-safe wrapper.
class Caller {
public:
 Caller() : idx_(0) {}  // `AddCaller` never returns 0.
 explicit Caller(size_t idx) : idx_(idx) { HWY_DASSERT(idx < kMaxCallers); }
 size_t Idx() const { return idx_; }

private:
 size_t idx_;
};

// Singleton, shared by all ThreadPool.
class Shared {
public:
 static HWY_DLLEXPORT Shared& Get();  // Thread-safe.

 Stopwatch MakeStopwatch() const { return Stopwatch(timer_); }
 Stopwatch& LastRootEnd() { return last_root_end_; }

 // Thread-safe. Calls with the same `name` return the same `Caller`.
 Caller AddCaller(const char* name) { return Caller(callers_.Add(name)); }

 PerCluster& Cluster(size_t cluster_idx) {
   HWY_DASSERT(cluster_idx < kMaxClusters);
   return per_cluster_[cluster_idx];
 }

 // Called from the main thread via `Profiler::PrintResults`.
 void PrintAndReset() {
   // Start counting pools (= one per cluster) invoked by each caller.
   size_t active_clusters[kMaxCallers] = {};
   per_cluster_[0].ForeachCaller(
       [&](size_t caller_idx, CallerAccumulator& acc) {
         active_clusters[caller_idx] = acc.Any();
       });
   // Reduce per-cluster accumulators into the first cluster.
   for (size_t cluster_idx = 1; cluster_idx < kMaxClusters; ++cluster_idx) {
     per_cluster_[cluster_idx].ForeachCaller(
         [&](size_t caller_idx, CallerAccumulator& acc) {
           active_clusters[caller_idx] += acc.Any();
           per_cluster_[0].Get(caller_idx).AddFrom(acc);
           acc = CallerAccumulator();
         });
     per_cluster_[cluster_idx].ResetBits();
   }

   CallerAccumulator total;
   for (size_t caller_idx : per_cluster_[0].Sorted()) {
     CallerAccumulator& acc = per_cluster_[0].Get(caller_idx);
     total.AddFrom(acc);  // must be before PrintAndReset.
     acc.PrintAndReset(callers_.Name(caller_idx), active_clusters[caller_idx]);
   }
   total.PrintTotal();
   per_cluster_[0].ResetBits();
 }

private:
 Shared()  // called via Get().
     : last_root_end_(timer_),
       send_config(callers_.Add("SendConfig")),
       dtor(callers_.Add("PoolDtor")),
       print_stats(callers_.Add("PrintStats")) {
   Profiler::Get().AddFunc(this, [this]() { PrintAndReset(); });
   // Can skip `RemoveFunc` because the singleton never dies.
 }

 const Timer timer_;
 Stopwatch last_root_end_;

 PerCluster per_cluster_[kMaxClusters];
 StringTable<kMaxCallers> callers_;

public:
 // Returned from `callers_.Add`:
 Caller send_config;
 Caller dtor;
 Caller print_stats;
};

#else

struct Stats {
 void NotifyRunStatic(size_t, timer::Ticks) {}
 void NotifyRunDynamic(size_t, size_t, size_t, timer::Ticks) {}
 void NotifyThreadRun(size_t, timer::Ticks, size_t, timer::Ticks,
                      timer::Ticks) {}
 void NotifyMainRun(size_t, timer::Ticks, timer::Ticks, timer::Ticks,
                    timer::Ticks) {}
 void PrintAndReset(size_t, timer::Ticks) {}
 void Reset(size_t = kMaxThreads) {}
};

struct Caller {};

class Shared {
public:
 static HWY_DLLEXPORT Shared& Get();  // Thread-safe.

 Stopwatch MakeStopwatch() const { return Stopwatch(timer_); }

 Caller AddCaller(const char*) { return Caller(); }

private:
 Shared() {}

 const Timer timer_;

public:
 Caller send_config;
 Caller dtor;
 Caller print_stats;
};

#endif  // PROFILER_ENABLED

// Per-worker state used by both main and worker threads. `ThreadFunc`
// (threads) and `ThreadPool` (main) have a few additional members of their own.
class alignas(HWY_ALIGNMENT) Worker {  // HWY_ALIGNMENT bytes
 static constexpr size_t kMaxVictims = 4;

 static constexpr auto kAcq = std::memory_order_acquire;
 static constexpr auto kRel = std::memory_order_release;

 bool OwnsGlobalIdx() const {
#if PROFILER_ENABLED
   if (global_idx_ >= profiler::kMaxWorkers) {
     HWY_WARN("Windows-only bug? global_idx %zu >= %zu.", global_idx_,
              profiler::kMaxWorkers);
   }
#endif  // PROFILER_ENABLED
   // Across-cluster pool owns all except the main thread, which is reserved by
   // profiler.cc.
   if (cluster_idx_ == kAllClusters) return global_idx_ != 0;
   // Within-cluster pool owns all except *its* main thread, because that is
   // owned by the across-cluster pool.
   return worker_ != 0;
 }

public:
 Worker(const size_t worker, const size_t num_threads,
        const PoolWorkerMapping mapping, const Divisor64& div_workers,
        const Stopwatch& stopwatch)
     : workers_(this - worker),
       worker_(worker),
       num_threads_(num_threads),
       stopwatch_(stopwatch),
       // If `num_threads == 0`, we might be in an inner pool and must use
       // the `global_idx` we are currently running on.
       global_idx_(num_threads == 0 ? Profiler::GlobalIdx() : mapping(worker)),
       cluster_idx_(mapping.ClusterIdx()) {
   HWY_DASSERT(IsAligned(this, HWY_ALIGNMENT));
   HWY_DASSERT(worker <= num_threads);
   const size_t num_workers = static_cast<size_t>(div_workers.GetDivisor());
   num_victims_ = static_cast<uint32_t>(HWY_MIN(kMaxVictims, num_workers));

   // Increase gap between coprimes to reduce collisions.
   const uint32_t coprime = ShuffledIota::FindAnotherCoprime(
       static_cast<uint32_t>(num_workers),
       static_cast<uint32_t>((worker + 1) * 257 + worker * 13));
   const ShuffledIota shuffled_iota(coprime);

   // To simplify `WorkerRun`, this worker is the first to 'steal' from.
   victims_[0] = static_cast<uint32_t>(worker);
   for (uint32_t i = 1; i < num_victims_; ++i) {
     victims_[i] = shuffled_iota.Next(victims_[i - 1], div_workers);
     HWY_DASSERT(victims_[i] != worker);
   }

   HWY_IF_CONSTEXPR(PROFILER_ENABLED) {
     if (HWY_LIKELY(OwnsGlobalIdx())) {
       Profiler::Get().ReserveWorker(global_idx_);
     }
   }
 }

 ~Worker() {
   HWY_IF_CONSTEXPR(PROFILER_ENABLED) {
     if (HWY_LIKELY(OwnsGlobalIdx())) {
       Profiler::Get().FreeWorker(global_idx_);
     }
   }
 }

 // Placement-newed by `WorkerLifecycle`, we do not expect any copying.
 Worker(const Worker&) = delete;
 Worker& operator=(const Worker&) = delete;

 size_t Index() const { return worker_; }
 // For work stealing.
 Worker* AllWorkers() { return workers_; }
 const Worker* AllWorkers() const { return workers_; }
 size_t NumThreads() const { return num_threads_; }

 size_t GlobalIdx() const { return global_idx_; }
 size_t ClusterIdx() const { return cluster_idx_; }

 void SetStartTime() { stopwatch_.Reset(); }
 timer::Ticks ElapsedTime() { return stopwatch_.Elapsed(); }

 // ------------------------ Per-worker storage for `SendConfig`

 Config NextConfig() const { return next_config_; }
 // Called during `SendConfig` by workers and now also the main thread. This
 // avoids a separate `ThreadPool` member which risks going out of sync.
 void SetNextConfig(Config copy) { next_config_ = copy; }

 Exit GetExit() const { return exit_; }
 void SetExit(Exit exit) { exit_ = exit; }

 uint32_t WorkerEpoch() const { return worker_epoch_; }
 uint32_t AdvanceWorkerEpoch() { return ++worker_epoch_; }

 // ------------------------ Task assignment

 // Called from the main thread.
 void SetRange(const uint64_t begin, const uint64_t end) {
   my_begin_.store(begin, kRel);
   my_end_.store(end, kRel);
 }

 uint64_t MyEnd() const { return my_end_.load(kAcq); }

 Span<const uint32_t> Victims() const {
   return hwy::Span<const uint32_t>(victims_.data(),
                                    static_cast<size_t>(num_victims_));
 }

 // Returns the next task to execute. If >= MyEnd(), it must be skipped.
 uint64_t WorkerReserveTask() {
   // TODO(janwas): replace with cooperative work-stealing.
   return my_begin_.fetch_add(1, std::memory_order_relaxed);
 }

 // ------------------------ Waiter: Threads wait for tasks

 // WARNING: some `Wait*` do not set this for all Worker instances. For
 // example, `WaitType::kBlock` only uses the first worker's `Waiter` because
 // one futex can wake multiple waiters. Hence we never load this directly
 // without going through `Wait*` policy classes, and must ensure all threads
 // use the same wait mode.

 const std::atomic<uint32_t>& Waiter() const { return wait_epoch_; }
 std::atomic<uint32_t>& MutableWaiter() { return wait_epoch_; }  // futex
 void StoreWaiter(uint32_t epoch) { wait_epoch_.store(epoch, kRel); }

 // ------------------------ Barrier: Main thread waits for workers

 // For use by `HasReached` and `UntilReached`.
 const std::atomic<uint32_t>& Barrier() const { return barrier_epoch_; }
 // Setting to `epoch` signals that the worker has reached the barrier.
 void StoreBarrier(uint32_t epoch) { barrier_epoch_.store(epoch, kRel); }

private:
 // Set by `SetRange`:
 std::atomic<uint64_t> my_begin_;
 std::atomic<uint64_t> my_end_;

 Worker* const workers_;
 const size_t worker_;
 const size_t num_threads_;

 Stopwatch stopwatch_;  // Reset by `SetStartTime`.
 const size_t global_idx_;
 const size_t cluster_idx_;

 // Use u32 to match futex.h. These must start at the initial value of
 // `worker_epoch_`.
 std::atomic<uint32_t> wait_epoch_{1};
 std::atomic<uint32_t> barrier_epoch_{1};

 uint32_t num_victims_;  // <= kPoolMaxVictims
 std::array<uint32_t, kMaxVictims> victims_;

 // Written and read by the same thread, hence not atomic.
 Config next_config_;
 Exit exit_ = Exit::kNone;
 // thread_pool_test requires nonzero epoch.
 uint32_t worker_epoch_ = 1;

 HWY_MEMBER_VAR_MAYBE_UNUSED uint8_t
     padding_[HWY_ALIGNMENT - 56 - 6 * sizeof(void*) - sizeof(victims_)];
};
static_assert(sizeof(Worker) == HWY_ALIGNMENT, "");

// Creates/destroys `Worker` using preallocated storage. See comment at
// `ThreadPool::worker_bytes_` for why we do not dynamically allocate.
class WorkerLifecycle {  // 0 bytes
public:
 // Placement new for `Worker` into `storage` because its ctor requires
 // the worker index. Returns array of all workers.
 static Worker* Init(uint8_t* storage, size_t num_threads,
                     PoolWorkerMapping mapping, const Divisor64& div_workers,
                     Shared& shared) {
   Worker* workers = new (storage)
       Worker(0, num_threads, mapping, div_workers, shared.MakeStopwatch());
   for (size_t worker = 1; worker <= num_threads; ++worker) {
     new (Addr(storage, worker)) Worker(worker, num_threads, mapping,
                                        div_workers, shared.MakeStopwatch());
     // Ensure pointer arithmetic is the same (will be used in Destroy).
     HWY_DASSERT(reinterpret_cast<uintptr_t>(workers + worker) ==
                 reinterpret_cast<uintptr_t>(Addr(storage, worker)));
   }

   // Publish non-atomic stores in `workers`.
   std::atomic_thread_fence(std::memory_order_release);

   return workers;
 }

 static void Destroy(Worker* workers, size_t num_threads) {
   for (size_t worker = 0; worker <= num_threads; ++worker) {
     workers[worker].~Worker();
   }
 }

private:
 static uint8_t* Addr(uint8_t* storage, size_t worker) {
   return storage + worker * sizeof(Worker);
 }
};

// Stores arguments to `Run`: the function and range of task indices. Set by
// the main thread, read by workers including the main thread.
class Tasks {
 static constexpr auto kAcq = std::memory_order_acquire;

 // Signature of the (internal) function called from workers(s) for each
 // `task` in the [`begin`, `end`) passed to Run(). Closures (lambdas) do not
 // receive the first argument, which points to the lambda object.
 typedef void (*RunFunc)(const void* opaque, uint64_t task, size_t worker);

public:
 Tasks() { HWY_DASSERT(IsAligned(this, 8)); }

 template <class Closure>
 void Set(uint64_t begin, uint64_t end, const Closure& closure) {
   constexpr auto kRel = std::memory_order_release;
   // `TestTasks` and `SetWaitMode` call this with `begin == end`.
   HWY_DASSERT(begin <= end);
   begin_.store(begin, kRel);
   end_.store(end, kRel);
   func_.store(static_cast<RunFunc>(&CallClosure<Closure>), kRel);
   opaque_.store(reinterpret_cast<const void*>(&closure), kRel);
 }

 // Assigns workers their share of `[begin, end)`. Called from the main
 // thread; workers are initializing or waiting for a command.
 // Negligible CPU time.
 static void DivideRangeAmongWorkers(const uint64_t begin, const uint64_t end,
                                     const Divisor64& div_workers,
                                     Worker* workers) {
   const size_t num_workers = static_cast<size_t>(div_workers.GetDivisor());
   HWY_DASSERT(num_workers > 1);  // Else Run() runs on the main thread.
   HWY_DASSERT(begin <= end);
   const size_t num_tasks = static_cast<size_t>(end - begin);

   // Assigning all remainders to the last worker causes imbalance. We instead
   // give one more to each worker whose index is less. This may be zero when
   // called from `TestTasks`.
   const size_t min_tasks = static_cast<size_t>(div_workers.Divide(num_tasks));
   const size_t remainder =
       static_cast<size_t>(div_workers.Remainder(num_tasks));

   uint64_t my_begin = begin;
   for (size_t worker = 0; worker < num_workers; ++worker) {
     const uint64_t my_end = my_begin + min_tasks + (worker < remainder);
     workers[worker].SetRange(my_begin, my_end);
     my_begin = my_end;
   }
   HWY_DASSERT(my_begin == end);
 }

 // Runs the worker's assigned range of tasks, plus work stealing if needed.
 void WorkerRun(Worker* worker, const Shared& shared, Stats& stats) const {
   if (NumTasks() > worker->NumThreads() + 1) {
     WorkerRunDynamic(worker, shared, stats);
   } else {
     WorkerRunStatic(worker, shared, stats);
   }
 }

private:
 // Special case for <= 1 task per worker, where stealing is unnecessary.
 void WorkerRunStatic(Worker* worker, const Shared& shared,
                      Stats& stats) const {
   const uint64_t begin = begin_.load(kAcq);
   const uint64_t end = end_.load(kAcq);
   HWY_DASSERT(begin <= end);
   const size_t index = worker->Index();

   const uint64_t task = begin + index;
   // We might still have more workers than tasks, so check first.
   if (HWY_LIKELY(task < end)) {
     const void* opaque = Opaque();
     const RunFunc func = Func();
     Stopwatch stopwatch = shared.MakeStopwatch();
     func(opaque, task, index);
     stats.NotifyRunStatic(index, stopwatch.Elapsed());
   }
 }

 // Must be called for each `worker` in [0, num_workers).
 //
 // A prior version of this code attempted to assign only as much work as a
 // worker will actually use. As with OpenMP's 'guided' strategy, we assigned
 // remaining/(k*num_threads) in each iteration. Although the worst-case
 // imbalance is bounded, this required several rounds of work allocation, and
 // the atomic counter did not scale to > 30 threads.
 //
 // We now use work stealing instead, where already-finished workers look for
 // and perform work from others, as if they were that worker. This deals with
 // imbalances as they arise, but care is required to reduce contention. We
 // randomize the order in which threads choose victims to steal from.
 void WorkerRunDynamic(Worker* worker, const Shared& shared,
                       Stats& stats) const {
   Worker* workers = worker->AllWorkers();
   const size_t index = worker->Index();
   const RunFunc func = Func();
   const void* opaque = Opaque();

   size_t sum_tasks = 0;
   size_t sum_stolen = 0;
   timer::Ticks sum_d_func = 0;
   // For each worker in random order, starting with our own, attempt to do
   // all their work.
   for (uint32_t victim : worker->Victims()) {
     Worker* other_worker = workers + victim;

     // Until all of other_worker's work is done:
     const uint64_t other_end = other_worker->MyEnd();
     for (;;) {
       // The worker that first sets `task` to `other_end` exits this loop.
       // After that, `task` can be incremented up to `num_workers - 1` times,
       // once per other worker.
       const uint64_t task = other_worker->WorkerReserveTask();
       if (HWY_UNLIKELY(task >= other_end)) {
         hwy::Pause();  // Reduce coherency traffic while stealing.
         break;
       }
       Stopwatch stopwatch = shared.MakeStopwatch();
       // Pass the index we are actually running on; this is important
       // because it is the TLS index for user code.
       func(opaque, task, index);
       sum_tasks++;
       sum_stolen += worker != other_worker;
       sum_d_func += stopwatch.Elapsed();
     }
   }
   stats.NotifyRunDynamic(index, sum_tasks, sum_stolen, sum_d_func);
 }

 size_t NumTasks() const {
   return static_cast<size_t>(end_.load(kAcq) - begin_.load(kAcq));
 }

 const void* Opaque() const { return opaque_.load(kAcq); }
 RunFunc Func() const { return func_.load(kAcq); }

 // Calls closure(task, worker). Signature must match `RunFunc`.
 template <class Closure>
 static void CallClosure(const void* opaque, uint64_t task, size_t worker) {
   (*reinterpret_cast<const Closure*>(opaque))(task, worker);
 }

 std::atomic<uint64_t> begin_;
 std::atomic<uint64_t> end_;
 std::atomic<RunFunc> func_;
 std::atomic<const void*> opaque_;
};
static_assert(sizeof(Tasks) == 16 + 2 * sizeof(void*), "");

// ------------------------------ Threads wait, main wakes them

// Considerations:
// - uint32_t storage per `Worker` so we can use `futex.h`.
// - avoid atomic read-modify-write. These are implemented on x86 using a LOCK
//   prefix, which interferes with other cores' cache-coherency transactions
//   and drains our core's store buffer. We use only store-release and
//   load-acquire. Although expressed using `std::atomic`, these are normal
//   loads/stores in the strong x86 memory model.
// - prefer to avoid resetting the state. "Sense-reversing" (flipping a flag)
//   would work, but we we prefer an 'epoch' counter because it is more useful
//   and easier to understand/debug, and as fast.

// Both the main thread and each worker maintain their own counter, which are
// implicitly synchronized by the barrier. To wake, the main thread does a
// store-release, and each worker does a load-acquire. The policy classes differ
// in whether they block or spin (with pause/monitor to reduce power), and
// whether workers check their own counter or a shared one.
//
// All methods are const because they only use storage in `Worker`, and we
// prefer to pass const-references to empty classes to enable type deduction.

// Futex: blocking reduces apparent CPU usage, but has higher wake latency.
struct WaitBlock {
 // Wakes all workers by storing the current `epoch`.
 void WakeWorkers(Worker* workers, const uint32_t epoch) const {
   HWY_DASSERT(epoch != 0);
   workers[1].StoreWaiter(epoch);
   WakeAll(workers[1].MutableWaiter());  // futex: expensive syscall
 }

 // Waits until `WakeWorkers(_, epoch)` has been called.
 template <class Spin>
 size_t UntilWoken(const Worker& worker, const Spin& /*spin*/) const {
   HWY_DASSERT(worker.Index() != 0);  // main is 0
   const uint32_t epoch = worker.WorkerEpoch();
   const Worker* workers = worker.AllWorkers();
   BlockUntilDifferent(epoch - 1, workers[1].Waiter());
   return 1;  // iterations
 }
};

// Single u32: single store by the main thread. All worker threads poll this
// one cache line and thus have it in a shared state, which means the store
// will invalidate each of them, leading to more transactions than SpinSeparate.
struct WaitSpin1 {
 void WakeWorkers(Worker* workers, const uint32_t epoch) const {
   workers[1].StoreWaiter(epoch);
 }

 // Returns the number of spin-wait iterations.
 template <class Spin>
 size_t UntilWoken(const Worker& worker, const Spin& spin) const {
   HWY_DASSERT(worker.Index() != 0);  // main is 0
   const Worker* workers = worker.AllWorkers();
   const uint32_t epoch = worker.WorkerEpoch();
   return spin.UntilEqual(epoch, workers[1].Waiter());
 }
};

// Separate u32 per thread: more stores for the main thread, but each worker
// only polls its own cache line, leading to fewer cache-coherency transactions.
struct WaitSpinSeparate {
 void WakeWorkers(Worker* workers, const uint32_t epoch) const {
   for (size_t thread = 0; thread < workers->NumThreads(); ++thread) {
     workers[1 + thread].StoreWaiter(epoch);
   }
 }

 template <class Spin>
 size_t UntilWoken(const Worker& worker, const Spin& spin) const {
   HWY_DASSERT(worker.Index() != 0);  // main is 0
   const uint32_t epoch = worker.WorkerEpoch();
   return spin.UntilEqual(epoch, worker.Waiter());
 }
};

// Calls unrolled code selected by all config enums.
template <class Func, typename... Args>
HWY_INLINE void CallWithConfig(const Config& config, Func&& func,
                              Args&&... args) {
 switch (config.wait_type) {
   case WaitType::kBlock:
     return func(SpinPause(), WaitBlock(), std::forward<Args>(args)...);
   case WaitType::kSpin1:
     return CallWithSpin(config.spin_type, func, WaitSpin1(),
                         std::forward<Args>(args)...);
   case WaitType::kSpinSeparate:
     return CallWithSpin(config.spin_type, func, WaitSpinSeparate(),
                         std::forward<Args>(args)...);
   case WaitType::kSentinel:
     HWY_UNREACHABLE;
 }
}

// ------------------------------ Barrier: Main thread waits for workers

// Similar to `WaitSpinSeparate`, a store-release of the same local epoch
// counter serves as a "have arrived" flag that does not require resetting.
class Barrier {
public:
 void WorkerReached(Worker& worker, uint32_t epoch) const {
   HWY_DASSERT(worker.Index() != 0);  // main is 0
   worker.StoreBarrier(epoch);
 }

 // Returns true if `worker` (can be the main thread) reached the barrier.
 bool HasReached(const Worker* worker, uint32_t epoch) const {
   const uint32_t barrier = worker->Barrier().load(std::memory_order_acquire);
   HWY_DASSERT(barrier <= epoch);
   return barrier == epoch;
 }

 // Main thread loops over each worker. A "group of 2 or 4" barrier was not
 // competitive on Skylake, Granite Rapids and Zen5.
 template <class Spin>
 void UntilReached(size_t num_threads, Worker* workers, const Spin& spin,
                   uint32_t epoch) const {
   workers[0].StoreBarrier(epoch);  // for main thread HasReached.

   for (size_t i = 0; i < num_threads; ++i) {
     // TODO: log number of spin-wait iterations.
     (void)spin.UntilEqual(epoch, workers[1 + i].Barrier());
   }
 }
};

// In debug builds, detects when functions are re-entered.
class BusyFlag {
public:
 void Set() { HWY_DASSERT(!busy_.test_and_set()); }
 void Clear() { HWY_IF_CONSTEXPR(HWY_IS_DEBUG_BUILD) busy_.clear(); }

private:
 std::atomic_flag busy_ = ATOMIC_FLAG_INIT;
};

}  // namespace pool

// Highly efficient parallel-for, intended for workloads with thousands of
// fork-join regions which consist of calling tasks[t](i) for a few hundred i,
// using dozens of threads.
//
// To reduce scheduling overhead, we assume that tasks are statically known and
// that threads do not schedule new work themselves. This allows us to avoid
// queues and only store a counter plus the current task. The latter is a
// pointer to a lambda function, without the allocation/indirection required for
// `std::function`.
//
// To reduce fork/join latency, we choose an efficient barrier, optionally
// enable spin-waits via `SetWaitMode`, and avoid any mutex/lock. We largely
// even avoid atomic RMW operations (LOCK prefix): currently for the wait and
// barrier, in future hopefully also for work stealing.
//
// To eliminate false sharing and enable reasoning about cache line traffic, the
// class is aligned and holds all worker state.
//
// For load-balancing, we use work stealing in random order.
class alignas(HWY_ALIGNMENT) ThreadPool {
 // Used to initialize `num_threads_` from the ctor argument.
 static size_t ClampedNumThreads(size_t num_threads) {
   // Upper bound is required for `worker_bytes_`.
   if (HWY_UNLIKELY(num_threads > pool::kMaxThreads)) {
     HWY_WARN("ThreadPool: clamping num_threads %zu to %zu.", num_threads,
              pool::kMaxThreads);
     num_threads = pool::kMaxThreads;
   }
   return num_threads;
 }

public:
 // This typically includes hyperthreads, hence it is a loose upper bound.
 // -1 because these are in addition to the main thread.
 static size_t MaxThreads() {
   LogicalProcessorSet lps;
   // This is OS dependent, but more accurate if available because it takes
   // into account restrictions set by cgroups or numactl/taskset.
   if (GetThreadAffinity(lps)) {
     return lps.Count() - 1;
   }
   return static_cast<size_t>(std::thread::hardware_concurrency() - 1);
 }

 // `num_threads` is the number of *additional* threads to spawn, which should
 // not exceed `MaxThreads()`. Note that the main thread also performs work.
 // `mapping` indicates how to map local worker_idx to global.
 ThreadPool(size_t num_threads,
            PoolWorkerMapping mapping = PoolWorkerMapping())
     : num_threads_(ClampedNumThreads(num_threads)),
       div_workers_(1 + num_threads_),
       shared_(pool::Shared::Get()),  // on first call, calls ReserveWorker(0)!
       workers_(pool::WorkerLifecycle::Init(worker_bytes_, num_threads_,
                                            mapping, div_workers_, shared_)) {
   // Leaves the default wait mode as `kBlock`, which means futex, because
   // spinning only makes sense when threads are pinned and wake latency is
   // important, so it must explicitly be requested by calling `SetWaitMode`.
   for (PoolWaitMode mode : {PoolWaitMode::kSpin, PoolWaitMode::kBlock}) {
     wait_mode_ = mode;  // for AutoTuner
     AutoTuner().SetCandidates(
         pool::Config::AllCandidates(mode));
   }

   // Skip empty pools because they do not update stats anyway.
   if (num_threads_ > 0) {
     Profiler::Get().AddFunc(this, [this]() { PrintStats(); });
   }

   threads_.reserve(num_threads_);
   for (size_t thread = 0; thread < num_threads_; ++thread) {
     threads_.emplace_back(
         ThreadFunc(workers_[1 + thread], tasks_, shared_, stats_));
   }

   // Threads' `Config` defaults to spinning. Change to `kBlock` (see above).
   // This also ensures all threads have started before we return, so that
   // startup latency is billed to the ctor, not the first `Run`.
   SendConfig(AutoTuner().Candidates()[0]);
 }

 // If we created threads, waits for them all to exit.
 ~ThreadPool() {
   // There is no portable way to request threads to exit like `ExitThread` on
   // Windows, otherwise we could call that from `Run`. Instead, we must cause
   // the thread to wake up and exit. We can just use `Run`.
   (void)RunWithoutAutotune(
       0, NumWorkers(), shared_.dtor,
       [this](HWY_MAYBE_UNUSED uint64_t task, size_t worker) {
         HWY_DASSERT(task == worker);
         workers_[worker].SetExit(Exit::kThread);
       });

   for (std::thread& thread : threads_) {
     HWY_DASSERT(thread.joinable());
     thread.join();
   }

   if (num_threads_ > 0) {
     Profiler::Get().RemoveFunc(this);
   }

   pool::WorkerLifecycle::Destroy(workers_, num_threads_);
 }

 ThreadPool(const ThreadPool&) = delete;
 ThreadPool& operator&(const ThreadPool&) = delete;

 // Returns number of Worker, i.e., one more than the largest `worker`
 // argument. Useful for callers that want to allocate thread-local storage.
 size_t NumWorkers() const {
   return static_cast<size_t>(div_workers_.GetDivisor());
 }

 // `mode` defaults to `kBlock`, which means futex. Switching to `kSpin`
 // reduces fork-join overhead especially when there are many calls to `Run`,
 // but wastes power when waiting over long intervals. Must not be called
 // concurrently with any `Run`, because this uses the same waiter/barrier.
 void SetWaitMode(PoolWaitMode mode) {
   wait_mode_ = mode;
   SendConfig(AutoTuneComplete() ? *AutoTuner().Best()
                                 : AutoTuner().NextConfig());
 }

 // For printing which is in use.
 pool::Config config() const { return workers_[0].NextConfig(); }

 bool AutoTuneComplete() const { return AutoTuner().Best(); }
 Span<CostDistribution> AutoTuneCosts() { return AutoTuner().Costs(); }

 static pool::Caller AddCaller(const char* name) {
   return pool::Shared::Get().AddCaller(name);
 }

 // parallel-for: Runs `closure(task, worker)` on workers for every `task` in
 // `[begin, end)`. Note that the unit of work should be large enough to
 // amortize the function call overhead, but small enough that each worker
 // processes a few tasks. Thus each `task` is usually a loop.
 //
 // Not thread-safe - concurrent parallel-for in the same `ThreadPool` are
 // forbidden unless `NumWorkers() == 1` or `end <= begin + 1`.
 template <class Closure>
 void Run(uint64_t begin, uint64_t end, pool::Caller caller,
          const Closure& closure) {
   AutoTuneT& auto_tuner = AutoTuner();
   // Already finished tuning: run without time measurement.
   if (HWY_LIKELY(auto_tuner.Best())) {
     // Don't care whether threads ran, we are done either way.
     (void)RunWithoutAutotune(begin, end, caller, closure);
     return;
   }

   // Not yet finished: measure time and notify autotuner.
   Stopwatch stopwatch(shared_.MakeStopwatch());
   // Skip update if threads didn't actually run.
   if (!RunWithoutAutotune(begin, end, caller, closure)) return;
   auto_tuner.NotifyCost(stopwatch.Elapsed());

   pool::Config next = auto_tuner.NextConfig();  // may be overwritten below
   if (auto_tuner.Best()) {  // just finished
     next = *auto_tuner.Best();
     HWY_IF_CONSTEXPR(pool::kVerbosity >= 1) {
       const size_t idx_best = static_cast<size_t>(
           auto_tuner.Best() - auto_tuner.Candidates().data());
       HWY_DASSERT(idx_best < auto_tuner.Costs().size());
       auto& AT = auto_tuner.Costs()[idx_best];
       const double best_cost = AT.EstimateCost();
       HWY_DASSERT(best_cost > 0.0);  // will divide by this below

       Stats s_ratio;
       for (size_t i = 0; i < auto_tuner.Costs().size(); ++i) {
         if (i == idx_best) continue;
         const double cost = auto_tuner.Costs()[i].EstimateCost();
         s_ratio.Notify(static_cast<float>(cost / best_cost));
       }

       fprintf(stderr,
               "Pool %3zu: %s %8.0f +/- %6.0f. Gain %.2fx [%.2fx, %.2fx]\n",
               NumWorkers(), auto_tuner.Best()->ToString().c_str(), best_cost,
               AT.Stddev(), s_ratio.GeometricMean(),
               static_cast<double>(s_ratio.Min()),
               static_cast<double>(s_ratio.Max()));
     }
   }
   SendConfig(next);
 }

 // Backward-compatible version without Caller.
 template <class Closure>
 void Run(uint64_t begin, uint64_t end, const Closure& closure) {
   Run(begin, end, pool::Caller(), closure);
 }

private:
 // Called via `CallWithConfig`.
 struct MainWakeAndBarrier {
   template <class Spin, class Wait>
   void operator()(const Spin& spin, const Wait& wait, pool::Worker& main,
                   const pool::Tasks& tasks, const pool::Shared& shared,
                   pool::Stats& stats) const {
     const pool::Barrier barrier;
     pool::Worker* workers = main.AllWorkers();
     HWY_DASSERT(&main == main.AllWorkers());  // main is first.
     const size_t num_threads = main.NumThreads();
     const uint32_t epoch = main.AdvanceWorkerEpoch();

     HWY_IF_CONSTEXPR(HWY_IS_DEBUG_BUILD) {
       for (size_t i = 0; i < 1 + num_threads; ++i) {
         HWY_DASSERT(!barrier.HasReached(workers + i, epoch));
       }
     }

     Stopwatch stopwatch(shared.MakeStopwatch());
     const timer::Ticks t_before_wake = stopwatch.Origin();
     wait.WakeWorkers(workers, epoch);
     const timer::Ticks d_wake = stopwatch.Elapsed();

     // Also perform work on the main thread before the barrier.
     tasks.WorkerRun(&main, shared, stats);
     const timer::Ticks d_run = stopwatch.Elapsed();

     // Spin-waits until all worker *threads* (not `main`, because it already
     // knows it is here) called `WorkerReached`.
     barrier.UntilReached(num_threads, workers, spin, epoch);
     const timer::Ticks d_barrier = stopwatch.Elapsed();
     stats.NotifyMainRun(main.NumThreads(), t_before_wake, d_wake, d_run,
                         d_barrier);

     HWY_IF_CONSTEXPR(HWY_IS_DEBUG_BUILD) {
       for (size_t i = 0; i < 1 + num_threads; ++i) {
         HWY_DASSERT(barrier.HasReached(workers + i, epoch));
       }
     }

     // Threads are or will soon be waiting `UntilWoken`, which serves as the
     // 'release' phase of the barrier.
   }
 };

 // Called by `std::thread`. Could also be a lambda.
 class ThreadFunc {
   // Functor called by `CallWithConfig`. Loops until `SendConfig` changes the
   // Spin or Wait policy or the pool is destroyed.
   struct WorkerLoop {
     template <class Spin, class Wait>
     void operator()(const Spin& spin, const Wait& wait, pool::Worker& worker,
                     pool::Tasks& tasks, const pool::Shared& shared,
                     pool::Stats& stats) const {
       do {
         // Main worker also calls this, so their epochs match.
         const uint32_t epoch = worker.AdvanceWorkerEpoch();

         Stopwatch stopwatch(shared.MakeStopwatch());

         const size_t wait_reps = wait.UntilWoken(worker, spin);
         const timer::Ticks d_wait = stopwatch.Elapsed();
         const timer::Ticks t_before_run = stopwatch.Origin();

         tasks.WorkerRun(&worker, shared, stats);
         const timer::Ticks d_run = stopwatch.Elapsed();
         stats.NotifyThreadRun(worker.Index(), d_wait, wait_reps, t_before_run,
                               d_run);

         // Notify barrier after `WorkerRun`. Note that we cannot send an
         // after-barrier timestamp, see above.
         pool::Barrier().WorkerReached(worker, epoch);
         // Check after `WorkerReached`, otherwise the main thread deadlocks.
       } while (worker.GetExit() == Exit::kNone);
     }
   };

  public:
   ThreadFunc(pool::Worker& worker, pool::Tasks& tasks,
              const pool::Shared& shared, pool::Stats& stats)
       : worker_(worker), tasks_(tasks), shared_(shared), stats_(stats) {}

   void operator()() {
     // Ensure main thread's writes are visible (synchronizes with fence in
     // `WorkerLifecycle::Init`).
     std::atomic_thread_fence(std::memory_order_acquire);

     HWY_DASSERT(worker_.Index() != 0);  // main is 0
     SetThreadName("worker%03zu", static_cast<int>(worker_.Index() - 1));

     worker_.SetStartTime();
     Profiler& profiler = Profiler::Get();
     profiler.SetGlobalIdx(worker_.GlobalIdx());
     // No Zone here because it would only exit after `GetExit`, which may be
     // after the main thread's `PROFILER_END_ROOT_RUN`, and thus too late to
     // be counted. Instead, `ProfilerFunc` records the elapsed time.

     // Loop termination via `GetExit` is triggered by `~ThreadPool`.
     for (;;) {
       // Uses the initial config, or the last one set during WorkerRun.
       CallWithConfig(worker_.NextConfig(), WorkerLoop(), worker_, tasks_,
                      shared_, stats_);

       // Exit or reset the flag and return to WorkerLoop with a new config.
       if (worker_.GetExit() == Exit::kThread) break;
       worker_.SetExit(Exit::kNone);
     }

     profiler.SetGlobalIdx(~size_t{0});

     // Defer `FreeWorker` until workers are destroyed to ensure the profiler
     // is not still using the worker.
   }

  private:
   pool::Worker& worker_;
   pool::Tasks& tasks_;
   const pool::Shared& shared_;
   pool::Stats& stats_;
 };

 void PrintStats() {
   // Total run time from all non-main threads.
   std::atomic<timer::Ticks> sum_thread_elapsed{0};
   (void)RunWithoutAutotune(
       0, NumWorkers(), shared_.print_stats,
       [this, &sum_thread_elapsed](HWY_MAYBE_UNUSED uint64_t task,
                                   size_t worker) {
         HWY_DASSERT(task == worker);
         // Skip any main thread(s) because they did not init the stopwatch.
         if (worker != 0) {
           sum_thread_elapsed.fetch_add(workers_[worker].ElapsedTime());
         }
       });
   const timer::Ticks thread_total =
       sum_thread_elapsed.load(std::memory_order_acquire);
   stats_.PrintAndReset(num_threads_, thread_total);
 }

 // Returns whether threads were used. If not, there is no need to update
 // the autotuner config.
 template <class Closure>
 bool RunWithoutAutotune(uint64_t begin, uint64_t end, pool::Caller caller,
                         const Closure& closure) {
   pool::Worker& main = workers_[0];

   const size_t num_tasks = static_cast<size_t>(end - begin);
   const size_t num_workers = NumWorkers();

   // If zero or one task, or no extra threads, run on the main thread without
   // setting any member variables, because we may be re-entering Run.
   if (HWY_UNLIKELY(num_tasks <= 1 || num_workers == 1)) {
     for (uint64_t task = begin; task < end; ++task) {
       closure(task, /*worker=*/0);
     }
     return false;
   }

   busy_.Set();

#if PROFILER_ENABLED
   const bool is_root = PROFILER_IS_ROOT_RUN();
   Stopwatch stopwatch(shared_.MakeStopwatch());
   const timer::Ticks wait_before =
       is_root ? shared_.LastRootEnd().Elapsed() : 0;
#endif

   tasks_.Set(begin, end, closure);

   // More than one task per worker: use work stealing.
   if (HWY_LIKELY(num_tasks > num_workers)) {
     pool::Tasks::DivideRangeAmongWorkers(begin, end, div_workers_, workers_);
   }

   // Runs `MainWakeAndBarrier` with the first worker slot.
   CallWithConfig(config(), MainWakeAndBarrier(), main, tasks_, shared_,
                  stats_);

#if PROFILER_ENABLED
   pool::CallerAccumulator& acc =
       shared_.Cluster(main.ClusterIdx()).Get(caller.Idx());
   acc.Add(num_tasks, num_workers, is_root, wait_before, stopwatch.Elapsed());
   if (is_root) {
     PROFILER_END_ROOT_RUN();
     shared_.LastRootEnd().Reset();
   }
#else
   (void)caller;
#endif

   busy_.Clear();
   return true;
 }

 // Sends `next_config` to workers:
 // - Main wakes threads using the current config.
 // - Threads copy `next_config` into their `Worker` during `WorkerRun`.
 // - Threads notify the (same) barrier and already wait for the next wake
 //   using `next_config`.
 HWY_NOINLINE void SendConfig(pool::Config next_config) {
   (void)RunWithoutAutotune(
       0, NumWorkers(), shared_.send_config,
       [this, next_config](HWY_MAYBE_UNUSED uint64_t task, size_t worker) {
         HWY_DASSERT(task == worker);  // one task per worker
         workers_[worker].SetNextConfig(next_config);
         workers_[worker].SetExit(Exit::kLoop);
       });

   // All have woken and are, or will be, waiting per `next_config`. Now we
   // can entirely switch the main thread's config for the next wake.
   workers_[0].SetNextConfig(next_config);
 }

 using AutoTuneT = AutoTune<pool::Config, 30>;
 AutoTuneT& AutoTuner() {
   static_assert(static_cast<size_t>(PoolWaitMode::kBlock) == 1, "");
   return auto_tune_[static_cast<size_t>(wait_mode_) - 1];
 }
 const AutoTuneT& AutoTuner() const {
   return auto_tune_[static_cast<size_t>(wait_mode_) - 1];
 }

 const size_t num_threads_;  // not including main thread
 const Divisor64 div_workers_;
 pool::Shared& shared_;
 pool::Worker* const workers_;  // points into `worker_bytes_`

 alignas(HWY_ALIGNMENT) pool::Stats stats_;

 // This is written by the main thread and read by workers, via reference
 // passed to `ThreadFunc`. Padding ensures that the workers' cache lines are
 // not unnecessarily invalidated when the main thread writes other members.
 alignas(HWY_ALIGNMENT) pool::Tasks tasks_;
 HWY_MEMBER_VAR_MAYBE_UNUSED char
     padding_[HWY_ALIGNMENT - sizeof(pool::Tasks)];

 pool::BusyFlag busy_;

 // Unmodified after ctor, but cannot be const because we call thread::join().
 std::vector<std::thread> threads_;

 PoolWaitMode wait_mode_;
 AutoTuneT auto_tune_[2];  // accessed via `AutoTuner`

 // Last because it is large. Store inside `ThreadPool` so that callers can
 // bind it to the NUMA node's memory. Not stored inside `WorkerLifecycle`
 // because that class would be initialized after `workers_`.
 alignas(HWY_ALIGNMENT) uint8_t
     worker_bytes_[sizeof(pool::Worker) * (pool::kMaxThreads + 1)];
};

}  // namespace hwy

#endif  // HIGHWAY_HWY_CONTRIB_THREAD_POOL_THREAD_POOL_H_