third_party/highway/hwy/contrib/thread_pool/thread_pool_test.cc

// 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.

#include "third_party/highway/hwy/contrib/thread_pool/thread_pool.h"

#include <math.h>  // sqrtf
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>

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

#include "third_party/highway/hwy/base.h"  // PopCount
#include "third_party/highway/hwy/contrib/thread_pool/spin.h"
#include "third_party/highway/hwy/contrib/thread_pool/topology.h"
#include "third_party/highway/hwy/tests/hwy_gtest.h"
#include "third_party/highway/hwy/tests/test_util-inl.h"  // AdjustedReps

namespace hwy {
namespace pool {
namespace {

TEST(ThreadPoolTest, TestCoprime) {
  // 1 is coprime with anything
  for (uint32_t i = 1; i < 500; ++i) {
    HWY_ASSERT(ShuffledIota::CoprimeNonzero(1, i));
    HWY_ASSERT(ShuffledIota::CoprimeNonzero(i, 1));
  }

  // Powers of two >= 2 are not coprime
  for (size_t i = 1; i < 20; ++i) {
    for (size_t j = 1; j < 20; ++j) {
      HWY_ASSERT(!ShuffledIota::CoprimeNonzero(1u << i, 1u << j));
    }
  }

  // 2^x and 2^x +/- 1 are coprime
  for (size_t i = 1; i < 30; ++i) {
    const uint32_t pow2 = 1u << i;
    HWY_ASSERT(ShuffledIota::CoprimeNonzero(pow2, pow2 + 1));
    HWY_ASSERT(ShuffledIota::CoprimeNonzero(pow2, pow2 - 1));
    HWY_ASSERT(ShuffledIota::CoprimeNonzero(pow2 + 1, pow2));
    HWY_ASSERT(ShuffledIota::CoprimeNonzero(pow2 - 1, pow2));
  }

  // Random number x * random y (both >= 2) is not co-prime with x nor y.
  RandomState rng;
  for (size_t i = 1; i < 5000; ++i) {
    const uint32_t x = (Random32(&rng) & 0xFFF7) + 2;
    const uint32_t y = (Random32(&rng) & 0xFFF7) + 2;
    HWY_ASSERT(!ShuffledIota::CoprimeNonzero(x * y, x));
    HWY_ASSERT(!ShuffledIota::CoprimeNonzero(x * y, y));
    HWY_ASSERT(!ShuffledIota::CoprimeNonzero(x, x * y));
    HWY_ASSERT(!ShuffledIota::CoprimeNonzero(y, x * y));
  }

  // Primes are all coprime (list from https://oeis.org/A000040)
  static constexpr uint32_t primes[] = {
      2,   3,   5,   7,   11,  13,  17,  19,  23,  29,  31,  37,  41,  43,  47,
      53,  59,  61,  67,  71,  73,  79,  83,  89,  97,  101, 103, 107, 109, 113,
      127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197,
      199, 211, 223, 227, 229, 233, 239, 241, 251, 257, 263, 269, 271};
  for (size_t i = 0; i < sizeof(primes) / sizeof(primes[0]); ++i) {
    for (size_t j = i + 1; j < sizeof(primes) / sizeof(primes[0]); ++j) {
      HWY_ASSERT(ShuffledIota::CoprimeNonzero(primes[i], primes[j]));
      HWY_ASSERT(ShuffledIota::CoprimeNonzero(primes[j], primes[i]));
    }
  }
}

// Ensures `shuffled` visits [0, size) exactly once starting from `current`.
void VerifyPermutation(uint32_t size, const Divisor64& divisor,
                       const ShuffledIota& shuffled, uint32_t current,
                       uint32_t* visited) {
  for (size_t i = 0; i < size; i++) {
    visited[i] = 0;
  }

  for (size_t i = 0; i < size; i++) {
    ++visited[current];
    current = shuffled.Next(current, divisor);
  }

  for (size_t i = 0; i < size; i++) {
    HWY_ASSERT(visited[i] == 1);
  }
}

// Verifies ShuffledIota generates a permutation of [0, size).
TEST(ThreadPoolTest, TestRandomPermutation) {
  constexpr size_t kMaxSize = 40;
  uint32_t visited[kMaxSize];

  // Exhaustive enumeration of size and starting point.
  for (uint32_t size = 1; size < kMaxSize; ++size) {
    const Divisor64 divisor(size);

    const uint32_t coprime = ShuffledIota::FindAnotherCoprime(size, 1);
    const ShuffledIota shuffled(coprime);

    for (uint32_t start = 0; start < size; ++start) {
      VerifyPermutation(size, divisor, shuffled, start, visited);
    }
  }
}

// Verifies multiple ShuffledIota are relatively independent.
TEST(ThreadPoolTest, TestMultiplePermutations) {
  constexpr size_t kMaxSize = 40;
  uint32_t coprimes[kMaxSize];
  // One per ShuffledIota; initially the starting value, then its Next().
  uint32_t current[kMaxSize];

  for (uint32_t size = 1; size < kMaxSize; ++size) {
    const Divisor64 divisor(size);

    // Create `size` ShuffledIota instances with unique coprimes.
    std::vector<ShuffledIota> shuffled;
    for (size_t i = 0; i < size; ++i) {
      coprimes[i] = ShuffledIota::FindAnotherCoprime(
          size, static_cast<uint32_t>((i + 1) * 257 + i * 13));
      shuffled.emplace_back(coprimes[i]);
    }

    // ShuffledIota[i] starts at i to match the worker thread use case.
    for (uint32_t i = 0; i < size; ++i) {
      current[i] = i;
    }

    size_t num_bad = 0;
    uint32_t all_visited[kMaxSize] = {0};

    // For each step, ensure there are few non-unique current[].
    for (size_t step = 0; step < size; ++step) {
      // How many times is each number visited?
      uint32_t visited[kMaxSize] = {0};
      for (size_t i = 0; i < size; ++i) {
        visited[current[i]] += 1;
        all_visited[current[i]] = 1;  // visited at all across all steps?
      }

      // How many numbers are visited multiple times?
      size_t num_contended = 0;
      uint32_t max_contention = 0;
      for (size_t i = 0; i < size; ++i) {
        num_contended += visited[i] > 1;
        max_contention = HWY_MAX(max_contention, visited[i]);
      }

      // Count/print if excessive collisions.
      const size_t expected =
          static_cast<size_t>(sqrtf(static_cast<float>(size)) * 2.0f);
      if ((num_contended > expected) && (max_contention > 3)) {
        ++num_bad;
        if (true) {
          fprintf(stderr, "size %u step %zu contended %zu max contention %u\n",
                  size, step, num_contended, max_contention);
          for (size_t i = 0; i < size; ++i) {
            fprintf(stderr, "  %u\n", current[i]);
          }
          fprintf(stderr, "coprimes\n");
          for (size_t i = 0; i < size; ++i) {
            fprintf(stderr, "  %u\n", coprimes[i]);
          }
        }
      }

      // Advance all ShuffledIota generators.
      for (size_t i = 0; i < size; ++i) {
        current[i] = shuffled[i].Next(current[i], divisor);
      }
    }  // step

    // Ensure each task was visited during at least one step.
    for (size_t i = 0; i < size; ++i) {
      HWY_ASSERT(all_visited[i] != 0);
    }

    if (num_bad != 0) {
      fprintf(stderr, "size %u total bad: %zu\n", size, num_bad);
    }
    HWY_ASSERT(num_bad < kMaxSize / 10);
  }  // size
}

class DoWait {
 public:
  explicit DoWait(Worker& worker) : worker_(worker) {}

  template <class Spin, class Wait>
  void operator()(const Spin& spin, const Wait& wait) const {
    wait.UntilWoken(worker_, spin);
  }

 private:
  Worker& worker_;
};

class DoWakeWorkers {
 public:
  explicit DoWakeWorkers(Worker* workers) : workers_(workers) {}

  template <class Spin, class Wait>
  void operator()(const Spin&, const Wait& wait) const {
    wait.WakeWorkers(workers_, workers_[0].WorkerEpoch());
  }

 private:
  Worker* const workers_;
};

// Verifies that waiter(s) can be woken by another thread.
TEST(ThreadPoolTest, TestWaiter) {
  if (!hwy::HaveThreadingSupport()) return;

  // Not actual threads, but we allocate and loop over this many workers.
  for (size_t num_threads = 1; num_threads < 6; ++num_threads) {
    const size_t num_workers = 1 + num_threads;
    auto storage = hwy::AllocateAligned<uint8_t>(num_workers * sizeof(Worker));
    HWY_ASSERT(storage);
    const Divisor64 div_workers(num_workers);
    Shared& shared = Shared::Get();  // already calls ReserveWorker(0).

    for (WaitType wait_type :
         {WaitType::kBlock, WaitType::kSpin1, WaitType::kSpinSeparate}) {
      Worker* workers = pool::WorkerLifecycle::Init(
          storage.get(), num_threads, PoolWorkerMapping(), div_workers, shared);

      alignas(8) const Config config(SpinType::kPause, wait_type);

      // This thread acts as the "main thread", which will wake the actual main
      // and all its worker instances.
      std::thread thread(
          [&]() { CallWithConfig(config, DoWakeWorkers(workers)); });

      // main is 0
      for (size_t worker = 1; worker < num_workers; ++worker) {
        CallWithConfig(config, DoWait(workers[1]));
      }
      thread.join();

      pool::WorkerLifecycle::Destroy(workers, num_workers);
    }
  }
}

// Ensures all tasks are run. Similar to TestPool below but without threads.
TEST(ThreadPoolTest, TestTasks) {
  for (size_t num_threads = 1; num_threads <= 8; ++num_threads) {
    const size_t num_workers = num_threads + 1;
    auto storage = hwy::AllocateAligned<uint8_t>(num_workers * sizeof(Worker));
    HWY_ASSERT(storage);
    const Divisor64 div_workers(num_workers);
    Shared& shared = Shared::Get();
    Stats stats;
    Worker* workers = WorkerLifecycle::Init(
        storage.get(), num_threads, PoolWorkerMapping(), div_workers, shared);

    constexpr uint64_t kMaxTasks = 20;
    uint64_t mementos[kMaxTasks];  // non-atomic, no threads involved.
    for (uint64_t num_tasks = 0; num_tasks < 20; ++num_tasks) {
      for (uint64_t begin = 0; begin < AdjustedReps(32); ++begin) {
        const uint64_t end = begin + num_tasks;

        ZeroBytes(mementos, sizeof(mementos));
        const auto func = [begin, end, &mementos](uint64_t task,
                                                  size_t /*worker*/) {
          HWY_ASSERT(begin <= task && task < end);

          // Store mementos ensure we visited each task.
          mementos[task - begin] = 1000 + task;
        };
        Tasks tasks;
        tasks.Set(begin, end, func);

        Tasks::DivideRangeAmongWorkers(begin, end, div_workers, workers);
        // The `tasks < workers` special case requires running by all workers.
        for (size_t worker = 0; worker < num_workers; ++worker) {
          tasks.WorkerRun(workers + worker, shared, stats);
        }

        // Ensure all tasks were run.
        for (uint64_t task = begin; task < end; ++task) {
          HWY_ASSERT_EQ(1000 + task, mementos[task - begin]);
        }
      }
    }

    WorkerLifecycle::Destroy(workers, num_workers);
  }
}

// Ensures task parameter is in bounds, every parameter is reached,
// pool can be reused (multiple consecutive Run calls), pool can be destroyed
// (joining with its threads), num_threads=0 works (runs on current thread).
TEST(ThreadPoolTest, TestPool) {
  if (!hwy::HaveThreadingSupport()) return;

  constexpr uint64_t kMaxTasks = 20;
  static std::atomic<uint64_t> mementos[kMaxTasks];
  static std::atomic<uint64_t> a_begin;
  static std::atomic<uint64_t> a_end;
  static std::atomic<uint64_t> a_num_workers;

  // Called by pool; sets mementos and runs a nested but serial Run.
  const auto func = [](uint64_t task, size_t worker) {
    HWY_ASSERT(worker < a_num_workers.load());
    const uint64_t begin = a_begin.load(std::memory_order_acquire);
    const uint64_t end = a_end.load(std::memory_order_acquire);

    if (!(begin <= task && task < end)) {
      HWY_ABORT("Task %d not in [%d, %d]", static_cast<int>(task),
                static_cast<int>(begin), static_cast<int>(end));
    }

    // Store mementos ensure we visited each task.
    mementos[task - begin].store(1000 + task);

    // Re-entering Run is fine on a 0-worker pool. Note that this must be
    // per-thread so that it gets the `global_idx` it is running on.
    hwy::ThreadPool inner(0);
    inner.Run(begin, end,
              [begin, end](uint64_t inner_task, size_t inner_worker) {
                HWY_ASSERT(inner_worker == 0);
                HWY_ASSERT(begin <= inner_task && inner_task < end);
              });
  };

  for (size_t num_threads = 0; num_threads <= 6; num_threads += 3) {
    hwy::ThreadPool pool(HWY_MIN(ThreadPool::MaxThreads(), num_threads));
    a_num_workers.store(pool.NumWorkers());
    for (bool spin : {true, false}) {
      pool.SetWaitMode(spin ? PoolWaitMode::kSpin : PoolWaitMode::kBlock);

      for (uint64_t num_tasks = 0; num_tasks < kMaxTasks; ++num_tasks) {
        for (uint64_t all_begin = 0; all_begin < AdjustedReps(32);
             ++all_begin) {
          const uint64_t all_end = all_begin + num_tasks;
          a_begin.store(all_begin, std::memory_order_release);
          a_end.store(all_end, std::memory_order_release);

          for (size_t i = 0; i < kMaxTasks; ++i) {
            mementos[i].store(0);
          }

          pool.Run(all_begin, all_end, func);

          for (uint64_t task = all_begin; task < all_end; ++task) {
            const uint64_t expected = 1000 + task;
            const uint64_t actual = mementos[task - all_begin].load();
            if (expected != actual) {
              HWY_ABORT(
                  "threads %zu, tasks %d: task not run, expected %d, got %d\n",
                  num_threads, static_cast<int>(num_tasks),
                  static_cast<int>(expected), static_cast<int>(actual));
            }
          }
        }
      }
    }
  }
}

// Debug tsan builds seem to generate incorrect codegen for [&] of atomics, so
// use a pointer to a state object instead.
struct SmallAssignmentState {
  // (Avoid mutex because it may perturb the worker thread scheduling)
  std::atomic<uint64_t> num_tasks{0};
  std::atomic<uint64_t> num_workers{0};
  std::atomic<uint64_t> id_bits{0};
  std::atomic<uint64_t> num_calls{0};
};

// Verify "thread" parameter when processing few tasks.
TEST(ThreadPoolTest, TestSmallAssignments) {
  if (!hwy::HaveThreadingSupport()) return;

  static SmallAssignmentState state;

  for (size_t num_threads :
       {size_t{0}, size_t{1}, size_t{3}, size_t{5}, size_t{8}}) {
    ThreadPool pool(HWY_MIN(ThreadPool::MaxThreads(), num_threads));
    state.num_workers.store(pool.NumWorkers());

    for (size_t mul = 1; mul <= 2; ++mul) {
      const size_t num_tasks = pool.NumWorkers() * mul;
      state.num_tasks.store(num_tasks);
      state.id_bits.store(0);
      state.num_calls.store(0);

      pool.Run(0, num_tasks, [](uint64_t task, size_t worker) {
        HWY_ASSERT(task < state.num_tasks.load());
        HWY_ASSERT(worker < state.num_workers.load());

        state.num_calls.fetch_add(1);

        uint64_t bits = state.id_bits.load();
        while (!state.id_bits.compare_exchange_weak(bits,
                                                    bits | (1ULL << worker))) {
        }
      });

      // Correct number of tasks.
      const uint64_t actual_calls = state.num_calls.load();
      HWY_ASSERT(num_tasks == actual_calls);

      const size_t num_participants = PopCount(state.id_bits.load());
      // <= because some workers may not manage to run any tasks.
      HWY_ASSERT(num_participants <= pool.NumWorkers());
    }
  }
}

struct Counter {
  Counter() {
    // Suppress "unused-field" warning.
    (void)padding;
  }
  void Assimilate(const Counter& victim) { counter += victim.counter; }
  std::atomic<uint64_t> counter{0};
  uint64_t padding[15];
};

// Can switch between any wait mode, and multiple times.
TEST(ThreadPoolTest, TestWaitMode) {
  if (!hwy::HaveThreadingSupport()) return;

  ThreadPool pool(9);
  RandomState rng;
  for (size_t i = 0; i < 100; ++i) {
    pool.SetWaitMode((Random32(&rng) & 1u) ? PoolWaitMode::kSpin
                                    : PoolWaitMode::kBlock);
  }
}

TEST(ThreadPoolTest, TestCounter) {
  if (!hwy::HaveThreadingSupport()) return;

  const size_t kNumThreads = 12;
  ThreadPool pool(kNumThreads);
  for (PoolWaitMode mode : {PoolWaitMode::kSpin, PoolWaitMode::kBlock}) {
    pool.SetWaitMode(mode);
    alignas(128) Counter counters[1+kNumThreads];

    const uint64_t kNumTasks = kNumThreads * 19;
    pool.Run(0, kNumTasks,
             [&counters](const uint64_t task, const size_t worker) {
               counters[worker].counter.fetch_add(task);
             });

    uint64_t expected = 0;
    for (uint64_t i = 0; i < kNumTasks; ++i) {
      expected += i;
    }

    for (size_t i = 1; i < pool.NumWorkers(); ++i) {
      counters[0].Assimilate(counters[i]);
    }
    HWY_ASSERT_EQ(expected, counters[0].counter.load());
  }
}

}  // namespace
}  // namespace pool
}  // namespace hwy

HWY_TEST_MAIN();