Add a refcounted thread pool keep-awake API

.claude/skills/scheduling/references/guide/08-benchmarking-and-profiling.md

@@ -28,6 +28,15 @@ problem — too few parallel tasks, a parallel loop that isn't outermost, or a
 memory-bandwidth ceiling — and the profiler's `average threads used` usually
 points the same way.
 
+Pipelines that run many short parallel loops back to back can spend a large
+share of their time waking up the thread pool, because idle workers go to sleep
+after spinning for only a few microseconds. Holding a
+`Halide::Runtime::ThreadPoolKeepAwake` (AOT), `Halide::ThreadPoolKeepAwake`
+(JIT), or `with hl.ThreadPoolKeepAwake():` (Python) around the timed region
+keeps idle workers polling for work instead, which is what a latency-sensitive
+application would do, at the cost of keeping those cores busy. See
+`halide_thread_pool_keep_awake` in `HalideRuntime.h`.
+
 ## The profiler is the main tool
 
 Profiling is turned on by adding `Target::Profile` to the target (for example,

python_bindings/halide/src/halide_/PyPipeline.cpp

@@ -79,6 +79,16 @@ struct PyProfilerScope {
     }
 };
 
+// Owns a ThreadPoolKeepAwake that a with-statement can release early, via
+// exit(), rather than waiting on garbage collection.
+struct PyThreadPoolKeepAwake {
+    std::unique_ptr<ThreadPoolKeepAwake> keep_awake = std::make_unique<ThreadPoolKeepAwake>();
+
+    void exit() {
+        keep_awake.reset();
+    }
+};
+
 }  // namespace
 
 void define_pipeline(py::module &m) {
@@ -438,6 +448,16 @@ void define_pipeline(py::module &m) {
         .def("func_stats", [](const PyProfilerScope &s, const Func &f) { return snapshot(s.get().func_stats(f)); }, py::arg("func"))
         .def("func_stats", [](const PyProfilerScope &s, const std::string &name) { return snapshot(s.get().func_stats(name)); }, py::arg("name"))
         .def("__repr__", [](const PyProfilerScope &s) -> std::string { return "<halide.ProfilerScope>"; });
+
+    py::class_<PyThreadPoolKeepAwake>(m, "ThreadPoolKeepAwake")
+        .def(py::init<>())
+        .def("__enter__", [](PyThreadPoolKeepAwake &k) -> PyThreadPoolKeepAwake & { return k; })
+        .def("__exit__", [](PyThreadPoolKeepAwake &k, const py::object &exc_type, const py::object &exc_value, const py::object &exc_traceback) -> bool {
+            k.exit();
+            return false;
+        })
+        .def("exit", &PyThreadPoolKeepAwake::exit)
+        .def("__repr__", [](const PyThreadPoolKeepAwake &k) -> std::string { return "<halide.ThreadPoolKeepAwake>"; });
 }
 
 }  // namespace PythonBindings

python_bindings/halide/test/correctness/CMakeLists.txt

@@ -32,6 +32,7 @@ set(tests
     serialization.py
     struct_type.py
     target.py
+    thread_pool_keep_awake.py
     tuple_select.py
     type.py
     user_context_test.py

python_bindings/halide/test/correctness/thread_pool_keep_awake.py

@@ -0,0 +1,63 @@
+import halide as hl
+
+
+def make_pipeline():
+    x, y = hl.Var("x"), hl.Var("y")
+    f = hl.Func("f")
+    f[x, y] = x * 3 + y
+    f.parallel(y)
+    return f
+
+
+def check(buf, width, height):
+    for y in range(height):
+        for x in range(width):
+            assert buf[x, y] == x * 3 + y, (x, y, buf[x, y])
+
+
+def test_with_statement():
+    f = make_pipeline()
+    with hl.ThreadPoolKeepAwake() as keep_awake:
+        assert repr(keep_awake) == "<halide.ThreadPoolKeepAwake>"
+        for i in range(100):
+            check(f.realize([16, 1 + i % 8]), 16, 1 + i % 8)
+        # Holders nest.
+        with hl.ThreadPoolKeepAwake():
+            check(f.realize([16, 8]), 16, 8)
+        check(f.realize([16, 8]), 16, 8)
+    check(f.realize([16, 8]), 16, 8)
+
+
+def test_exit():
+    f = make_pipeline()
+    keep_awake = hl.ThreadPoolKeepAwake()
+    check(f.realize([16, 8]), 16, 8)
+    # Releasing is idempotent, and leaving the with-statement afterwards
+    # doesn't release a second reference.
+    keep_awake.exit()
+    keep_awake.exit()
+    with hl.ThreadPoolKeepAwake() as k:
+        k.exit()
+        check(f.realize([16, 8]), 16, 8)
+    check(f.realize([16, 8]), 16, 8)
+
+
+def test_exception():
+    f = make_pipeline()
+    try:
+        with hl.ThreadPoolKeepAwake():
+            check(f.realize([16, 8]), 16, 8)
+            raise ValueError("expected")
+    except ValueError:
+        pass
+    check(f.realize([16, 8]), 16, 8)
+
+
+def main():
+    test_with_statement()
+    test_exit()
+    test_exception()
+
+
+if __name__ == "__main__":
+    main()

src/JITModule.cpp

@@ -720,6 +720,15 @@ int JITModule::set_num_threads(int n) const {
     return 1;
 }
 
+int JITModule::thread_pool_keep_awake(bool keep_awake) const {
+    std::map<std::string, Symbol>::const_iterator f =
+        exports().find("halide_thread_pool_keep_awake");
+    if (f != exports().end()) {
+        return (reinterpret_bits<int (*)(bool)>(f->second.address))(keep_awake);
+    }
+    return 0;
+}
+
 bool JITModule::compiled() const {
     return jit_module->JIT != nullptr;
 }
@@ -732,6 +741,12 @@ JITHandlers active_handlers;
 int64_t default_cache_size;
 int default_trace_file = -1;
 
+// The number of references held through JITSharedRuntime::thread_pool_keep_awake,
+// and whether the MainShared runtime currently holds one reference on their
+// behalf. Guarded by shared_runtimes_mutex.
+int keep_awake_count = 0;
+bool keep_awake_applied = false;
+
 void merge_handlers(JITHandlers &base, const JITHandlers &addins) {
     if (addins.custom_print) {
         base.custom_print = addins.custom_print;
@@ -1216,6 +1231,11 @@ JITModule &make_module(llvm::Module *for_module, Target target,
                 runtime.set_trace_file(default_trace_file);
             }
 
+            if (keep_awake_count > 0) {
+                runtime.thread_pool_keep_awake(true);
+                keep_awake_applied = true;
+            }
+
             runtime.jit_module->name = "MainShared";
         } else {
             runtime.jit_module->name = "GPU";
@@ -1452,6 +1472,13 @@ void JITSharedRuntime::populate_jit_handlers(JITUserContext *jit_user_context, c
 void JITSharedRuntime::release_all() {
     std::scoped_lock lock(shared_runtimes_mutex);
 
+    // Compiled pipelines may keep the old runtime alive, so let its thread
+    // pool sleep. A new runtime picks up the reference when it's created.
+    if (keep_awake_applied) {
+        shared_runtimes(MainShared).thread_pool_keep_awake(false);
+        keep_awake_applied = false;
+    }
+
     for (int i = MaxRuntimeKind; i > 0; i--) {
         shared_runtimes((RuntimeKind)(i - 1)) = JITModule();
     }
@@ -1500,6 +1527,27 @@ int JITSharedRuntime::set_num_threads(int n) {
     return shared_runtimes(MainShared).set_num_threads(n);
 }
 
+int JITSharedRuntime::thread_pool_keep_awake(bool keep_awake) {
+    std::scoped_lock lock(shared_runtimes_mutex);
+    // The MainShared runtime holds a single reference while keep_awake_count
+    // is positive, so that references can be acquired before it exists.
+    if (keep_awake) {
+        if (keep_awake_count++ == 0 && shared_runtimes(MainShared).compiled()) {
+            shared_runtimes(MainShared).thread_pool_keep_awake(true);
+            keep_awake_applied = true;
+        }
+    } else {
+        user_assert(keep_awake_count > 0)
+            << "JITSharedRuntime::thread_pool_keep_awake(false) called without a "
+            << "matching JITSharedRuntime::thread_pool_keep_awake(true).\n";
+        if (--keep_awake_count == 0 && keep_awake_applied) {
+            shared_runtimes(MainShared).thread_pool_keep_awake(false);
+            keep_awake_applied = false;
+        }
+    }
+    return keep_awake_count;
+}
+
 void *JITSharedRuntime::find_symbol(const Target &target, const std::string &name) {
     for (const JITModule &m : JITSharedRuntime::get(nullptr, target, false)) {
         JITModule::Symbol sym = m.find_symbol_by_name(name);
@@ -1646,4 +1694,13 @@ void JITFuncCallContext::finalize(int exit_status) {
 }
 
 }  // namespace Internal
+
+ThreadPoolKeepAwake::ThreadPoolKeepAwake() {
+    Internal::JITSharedRuntime::thread_pool_keep_awake(true);
+}
+
+ThreadPoolKeepAwake::~ThreadPoolKeepAwake() {
+    Internal::JITSharedRuntime::thread_pool_keep_awake(false);
+}
+
 }  // namespace Halide

src/JITModule.h

@@ -196,6 +196,24 @@ struct JITUserContext {
     JITHandlers handlers;
 };
 
+/** Hold a reference on the JIT thread pool's keep-awake count for the lifetime
+ * of this object, so that back-to-back small parallel loops don't pay to wake
+ * up the pool's threads. Instances may be nested, and may live on different
+ * threads. Idle threads keep their cores busy while any instance is alive, so
+ * hold one only around phases that run many short parallel loops. See
+ * halide_thread_pool_keep_awake in HalideRuntime.h for details. This is the
+ * JIT equivalent of Halide::Runtime::ThreadPoolKeepAwake, which is for AOT
+ * code. See also Internal::JITSharedRuntime::thread_pool_keep_awake. */
+class ThreadPoolKeepAwake {
+public:
+    ThreadPoolKeepAwake();
+    ~ThreadPoolKeepAwake();
+    ThreadPoolKeepAwake(const ThreadPoolKeepAwake &) = delete;
+    ThreadPoolKeepAwake &operator=(const ThreadPoolKeepAwake &) = delete;
+    ThreadPoolKeepAwake(ThreadPoolKeepAwake &&) = delete;
+    ThreadPoolKeepAwake &operator=(ThreadPoolKeepAwake &&) = delete;
+};
+
 namespace Internal {
 
 class JITModuleContents;
@@ -314,6 +332,11 @@ struct JITModule {
     /** See JITSharedRuntime::set_num_threads */
     int set_num_threads(int) const;
 
+    /** Call this module's halide_thread_pool_keep_awake, if it has one, and
+     * return its result. Returns 0 otherwise. Prefer
+     * JITSharedRuntime::thread_pool_keep_awake. */
+    int thread_pool_keep_awake(bool) const;
+
     /** Return true if compile_module has been called on this module. */
     bool compiled() const;
 };
@@ -365,6 +388,16 @@ class JITSharedRuntime {
      * number. */
     static int set_num_threads(int);
 
+    /** Acquire (true) or release (false) a reference on the keep-awake count of
+     * the shared JIT runtime's thread pool, and return the new number of
+     * references held through this function. See halide_thread_pool_keep_awake
+     * in HalideRuntime.h. The references are tracked here, so they may be
+     * acquired before any JIT code has run: they are applied when the shared
+     * runtime is created, and carried over to a new shared runtime after
+     * release_all. Releasing a reference that isn't held is a user error. Most
+     * code should use Halide::ThreadPoolKeepAwake instead. */
+    static int thread_pool_keep_awake(bool);
+
     /** Search the shared JIT runtime for `target` for a symbol with the
      * given name. Returns the first match's address, or nullptr if no
      * runtime module exports it. JIT shared runtimes are created

src/runtime/HalideRuntime.h

@@ -423,6 +423,45 @@ extern int halide_get_num_threads();
 extern int halide_set_num_threads(int n);
 // @}
 
+/** Acquire (keep_awake = true) or release (keep_awake = false) a reference on
+ * the default thread pool's process-wide keep-awake count, and return the new
+ * count.
+ *
+ * Normally, a thread pool thread that runs out of work spins briefly and then
+ * goes to sleep, so a parallel loop that starts after the pool has been idle
+ * for more than a few microseconds pays to wake its threads up. While the
+ * keep-awake count is positive, idle threads (workers, and threads waiting for
+ * parallel loops they started to finish) instead keep polling for work,
+ * calling halide_thread_yield() between polls. This makes back-to-back small
+ * parallel loops, separated by serial work, much cheaper. The cost is that
+ * those threads keep their cores busy while the count is held, even with no
+ * work to do, so hold it only around phases that run many short parallel loops
+ * (e.g. a benchmark, or a model's inference loop) and release it afterwards.
+ * Workers blocked on a semaphore of an async pipeline still go to sleep.
+ *
+ * Holders compose: the pool stays awake until every reference is released.
+ * When the count returns to zero, spinning threads go to sleep on their next
+ * poll. As a backstop against leaked references, and against competing for
+ * cores with other thread pools for long periods, an idle thread goes to sleep
+ * anyway after polling a few thousand times without finding work (a few
+ * milliseconds, depending on the cost of halide_thread_yield); the next
+ * parallel loop wakes it as usual.
+ *
+ * Threads that are already asleep when the first reference is acquired start
+ * polling once the next parallel loop wakes them. Releasing a reference that
+ * isn't held is an error: it calls halide_error, leaves the count at zero,
+ * and returns halide_error_code_generic_error.
+ *
+ * The count is independent of halide_set_num_threads, and survives
+ * halide_shutdown_thread_pool (which still wakes and joins all threads). This
+ * function may be called from any thread. Custom thread pools (see
+ * halide_set_custom_do_par_for) are not affected, and on platforms without
+ * threads this only maintains the count. C++ code can use the RAII wrapper
+ * Halide::Runtime::ThreadPoolKeepAwake. In JIT code, use
+ * Halide::ThreadPoolKeepAwake instead (hl.ThreadPoolKeepAwake in Python).
+ */
+extern int halide_thread_pool_keep_awake(bool keep_awake);
+
 /** Halide calls these functions to allocate and free memory. To
  * replace in AOT code, use the halide_set_custom_malloc and
  * halide_set_custom_free, or (on platforms that support weak
@@ -2523,6 +2562,38 @@ HALIDE_ALWAYS_INLINE constexpr halide_type_t halide_type_of<int64_t>() {
 
 #ifndef COMPILING_HALIDE_RUNTIME
 
+namespace Halide {
+namespace Runtime {
+
+/** Hold a reference on the default thread pool's keep-awake count (see
+ * halide_thread_pool_keep_awake) for the lifetime of this object. Instances
+ * may be nested, and may live on different threads. For example:
+ \code
+ {
+     Halide::Runtime::ThreadPoolKeepAwake keep_awake;
+     for (int i = 0; i < 1000; i++) {
+         my_small_parallel_pipeline(...);
+     }
+ }
+ \endcode
+ * For JIT-compiled code, use Halide::ThreadPoolKeepAwake instead. */
+class ThreadPoolKeepAwake {
+public:
+    ThreadPoolKeepAwake() {
+        (void)halide_thread_pool_keep_awake(true);
+    }
+    ~ThreadPoolKeepAwake() {
+        (void)halide_thread_pool_keep_awake(false);
+    }
+    ThreadPoolKeepAwake(const ThreadPoolKeepAwake &) = delete;
+    ThreadPoolKeepAwake &operator=(const ThreadPoolKeepAwake &) = delete;
+    ThreadPoolKeepAwake(ThreadPoolKeepAwake &&) = delete;
+    ThreadPoolKeepAwake &operator=(ThreadPoolKeepAwake &&) = delete;
+};
+
+}  // namespace Runtime
+}  // namespace Halide
+
 // These structures are used by `function_info_header` files
 // (generated by passing `-e function_info_header` to a Generator).
 // The generated files contain documentation on the proper usage.

src/runtime/fake_thread_pool.cpp

@@ -1,5 +1,6 @@
 #include "HalideRuntime.h"
 #include "printer.h"
+#include "runtime_atomics.h"
 
 extern "C" {
 
@@ -71,6 +72,7 @@ WEAK halide_semaphore_init_t custom_semaphore_init = halide_default_semaphore_in
 WEAK halide_semaphore_try_acquire_t custom_semaphore_try_acquire = halide_default_semaphore_try_acquire;
 WEAK halide_semaphore_release_t custom_semaphore_release = halide_default_semaphore_release;
 WEAK halide_mutex_array halide_fake_mutex_array;
+WEAK int keep_awake_count = 0;
 
 }  // namespace Internal
 }  // namespace Runtime
@@ -131,6 +133,26 @@ WEAK int halide_get_num_threads() {
     return 1;
 }
 
+WEAK int halide_thread_pool_keep_awake(bool keep_awake) {
+    // There are no threads to keep awake, but keep the count so that
+    // unbalanced calls are diagnosed as they are with a real thread pool.
+    using namespace Halide::Runtime::Internal;
+    if (keep_awake) {
+        return Synchronization::atomic_add_fetch_sequentially_consistent(&keep_awake_count, 1);
+    }
+    int count;
+    Synchronization::atomic_load_relaxed(&keep_awake_count, &count);
+    int desired;
+    do {
+        if (count <= 0) {
+            halide_error(nullptr, "halide_thread_pool_keep_awake(false) called without a matching halide_thread_pool_keep_awake(true).\n");
+            return halide_error_code_generic_error;
+        }
+        desired = count - 1;
+    } while (!Synchronization::atomic_cas_weak_relacq_relaxed(&keep_awake_count, &count, &desired));
+    return desired;
+}
+
 WEAK halide_do_task_t halide_set_custom_do_task(halide_do_task_t f) {
     halide_do_task_t result = custom_do_task;
     custom_do_task = f;