From fb24763a15e6be34782e147f04d03e6030f36b4f Mon Sep 17 00:00:00 2001
From: Robert O'Callahan <rocallahan@google.com>
Date: Wed, 28 Jan 2026 18:50:23 +0000
Subject: [PATCH] Add `work_pool_size`, `IntRange`, `item_range_for_worker`,
 and `ThreadIndex`

We'll use these later in this PR.
---
 kernel/threading.cc | 35 +++++++++++++++++++++++++++++++++++
 kernel/threading.h  | 29 +++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/kernel/threading.cc b/kernel/threading.cc
index dcc044c89..8c9bfb390 100644
--- a/kernel/threading.cc
+++ b/kernel/threading.cc
@@ -17,6 +17,20 @@ static int get_max_threads()
 	return max_threads;
 }
 
+static int init_work_units_per_thread_override()
+{
+	const char *v = getenv("YOSYS_WORK_UNITS_PER_THREAD");
+	if (v == nullptr)
+		return 0;
+	return atoi(v);
+}
+
+static int get_work_units_per_thread_override()
+{
+	static int work_units_per_thread = init_work_units_per_thread_override();
+	return work_units_per_thread;
+}
+
 void DeferredLogs::flush()
 {
 	for (auto &m : logs)
@@ -37,6 +51,14 @@ int ThreadPool::pool_size(int reserved_cores, int max_worker_threads)
 #endif
 }
 
+int ThreadPool::work_pool_size(int reserved_cores, int work_units, int work_units_per_thread)
+{
+	int work_units_per_thread_override = get_work_units_per_thread_override();
+	if (work_units_per_thread_override > 0)
+		work_units_per_thread = work_units_per_thread_override;
+	return pool_size(reserved_cores, work_units / work_units_per_thread);
+}
+
 ThreadPool::ThreadPool(int pool_size, std::function<void(int)> b)
 	: body(std::move(b))
 {
@@ -57,4 +79,17 @@ ThreadPool::~ThreadPool()
 #endif
 }
 
+IntRange item_range_for_worker(int num_items, int thread_num, int num_threads)
+{
+	if (num_threads <= 1) {
+		return {0, num_items};
+	}
+	int items_per_thread = num_items / num_threads;
+	int extra_items = num_items % num_threads;
+	// The first `extra_items` threads get one extra item.
+	int start = thread_num * items_per_thread + std::min(thread_num, extra_items);
+	int end = (thread_num + 1) * items_per_thread + std::min(thread_num + 1, extra_items);
+	return {start, end};
+}
+
 YOSYS_NAMESPACE_END
diff --git a/kernel/threading.h b/kernel/threading.h
index b8cd62f87..eb068bb20 100644
--- a/kernel/threading.h
+++ b/kernel/threading.h
@@ -131,6 +131,11 @@ public:
 	// The result may be 0.
 	static int pool_size(int reserved_cores, int max_worker_threads);
 
+	// Computes the number of worker threads to use, by dividing work_units among threads.
+	// For testing purposes you can set YOSYS_WORK_UNITS_PER_THREAD to override `work_units_per_thread`.
+	// The result may be 0.
+	static int work_pool_size(int reserved_cores, int work_units, int work_units_per_thread);
+
 	// Create a pool of threads running the given closure (parameterized by thread number).
 	// `pool_size` must be the result of a `pool_size()` call.
 	ThreadPool(int pool_size, std::function<void(int)> b);
@@ -154,6 +159,30 @@ private:
 #endif
 };
 
+// A range of integers [start_, end_) that can be iterated over with a
+// C++ range-based for loop.
+struct IntRange {
+	int start_;
+	int end_;
+	struct Int {
+		int v;
+		int operator*() const { return v; }
+		Int &operator++() { ++v; return *this; }
+		bool operator!=(const Int &other) const { return v != other.v; }
+	};
+	Int begin() const { return {start_}; }
+	Int end() const { return {end_}; }
+};
+// Divides some number of items into `num_threads` subranges and returns the
+// `thread_num`'th subrange. If `num_threads` is zero, returns the whole range.
+IntRange item_range_for_worker(int num_items, int thread_num, int num_threads);
+
+// A type that encapsulates the index of a thread in some list of threads. Useful for
+// stronger typechecking and code readability.
+struct ThreadIndex {
+	int thread_num;
+};
+
 template <class T>
 class ConcurrentStack
 {