diff --git a/.claude/skills/scheduling/references/guide/08-benchmarking-and-profiling.md b/.claude/skills/scheduling/references/guide/08-benchmarking-and-profiling.md index bd868160cfeb..e13ba9d3b35f 100644 --- a/.claude/skills/scheduling/references/guide/08-benchmarking-and-profiling.md +++ b/.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, diff --git a/python_bindings/halide/src/halide_/PyPipeline.cpp b/python_bindings/halide/src/halide_/PyPipeline.cpp index 2a07a8aecb4e..66db9f65cd3d 100644 --- a/python_bindings/halide/src/halide_/PyPipeline.cpp +++ b/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 keep_awake = std::make_unique(); + + 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 ""; }); + + py::class_(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 ""; }); } } // namespace PythonBindings diff --git a/python_bindings/halide/test/correctness/CMakeLists.txt b/python_bindings/halide/test/correctness/CMakeLists.txt index b6e1156405b0..6ae3ac4a10a2 100644 --- a/python_bindings/halide/test/correctness/CMakeLists.txt +++ b/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 diff --git a/python_bindings/halide/test/correctness/thread_pool_keep_awake.py b/python_bindings/halide/test/correctness/thread_pool_keep_awake.py new file mode 100644 index 000000000000..84b23194bdf1 --- /dev/null +++ b/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) == "" + 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() diff --git a/src/JITModule.cpp b/src/JITModule.cpp index 4ec9894759fc..0fe61f3b0b47 100644 --- a/src/JITModule.cpp +++ b/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::const_iterator f = + exports().find("halide_thread_pool_keep_awake"); + if (f != exports().end()) { + return (reinterpret_bits(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 diff --git a/src/JITModule.h b/src/JITModule.h index b54d88aeb0fd..f4c1ef679204 100644 --- a/src/JITModule.h +++ b/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 diff --git a/src/runtime/HalideRuntime.h b/src/runtime/HalideRuntime.h index 6c186ab8e4d0..a2f4dcff7ce3 100644 --- a/src/runtime/HalideRuntime.h +++ b/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() { #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. diff --git a/src/runtime/fake_thread_pool.cpp b/src/runtime/fake_thread_pool.cpp index fe8c852897f1..4ad22fceccfc 100644 --- a/src/runtime/fake_thread_pool.cpp +++ b/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; diff --git a/src/runtime/runtime_api.cpp b/src/runtime/runtime_api.cpp index 0599225f816a..05d031befbb4 100644 --- a/src/runtime/runtime_api.cpp +++ b/src/runtime/runtime_api.cpp @@ -218,6 +218,7 @@ extern "C" __attribute__((used)) void *halide_runtime_api_functions[] = { (void *)&halide_start_clock, (void *)&halide_start_timer_chain, (void *)&halide_string_to_string, + (void *)&halide_thread_pool_keep_awake, (void *)&halide_trace, (void *)&halide_trace_helper, (void *)&halide_uint64_to_string, diff --git a/src/runtime/thread_pool_common.h b/src/runtime/thread_pool_common.h index abf5ebb3ba02..df9fe0cfbff2 100644 --- a/src/runtime/thread_pool_common.h +++ b/src/runtime/thread_pool_common.h @@ -17,6 +17,28 @@ namespace Halide { namespace Runtime { namespace Internal { +// The process-wide keep-awake reference count; see +// halide_thread_pool_keep_awake. It lives outside work_queue so that +// halide_shutdown_thread_pool, which resets work_queue, doesn't drop +// references that are still held. +WEAK int keep_awake_count = 0; + +// While the keep-awake count is held, how many times an idle thread polls for +// work (each poll calls halide_thread_yield) before going to sleep anyway. This +// is a backstop against leaked references and against taking cores from other +// thread pools for long periods. Waking a thread costs tens of microseconds, so +// once a thread has been idle for a few milliseconds, spinning any longer saves +// at most about 1% of the idle time. With a halide_thread_yield that costs +// about a microsecond (a sched_yield-style syscall), this is a few +// milliseconds. +constexpr int keep_awake_max_spins = 1 << 12; + +ALWAYS_INLINE bool keep_awake_held() { + int count; + Synchronization::atomic_load_relaxed(&keep_awake_count, &count); + return count > 0; +} + // A condition variable, augmented with a bit of spinning on an atomic counter // before going to sleep for real. This helps reduce overhead at the end of a // parallel for loop when idle worker threads are waiting for other threads to @@ -27,11 +49,14 @@ struct halide_cond_with_spinning { void wait(halide_mutex *mutex) { // First spin for a bit, checking the counter for another thread to bump - // it. + // it. While the keep-awake count is held, keep spinning (up to + // keep_awake_max_spins) instead of going to sleep. uintptr_t initial; Synchronization::atomic_load_relaxed(&counter, &initial); halide_mutex_unlock(mutex); - for (int spin = 0; spin < 40; spin++) { + for (int spin = 0; + spin < 40 || (spin < keep_awake_max_spins && keep_awake_held()); + spin++) { halide_thread_yield(); uintptr_t current; Synchronization::atomic_load_relaxed(&counter, ¤t); @@ -825,10 +850,34 @@ WEAK int halide_get_num_threads() { return n; } +WEAK int halide_thread_pool_keep_awake(bool keep_awake) { + if (keep_awake) { + // Idle threads that are already asleep start spinning once the next + // parallel loop wakes them. + return Synchronization::atomic_add_fetch_sequentially_consistent(&keep_awake_count, 1); + } + + // Threads spinning because of the count notice it dropping to zero on + // their next poll and go to sleep, so there's nothing to wake here. + 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 void halide_shutdown_thread_pool() { if (work_queue.initialized) { // Wake everyone up and tell them the party's over and it's time - // to go home + // to go home. This includes threads spinning because the + // keep-awake count is held. The count itself is preserved, so the + // threads of a restarted pool stay awake too. halide_mutex_lock(&work_queue.mutex); work_queue.shutdown = true; diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index 026792711ab9..3c5015462c6b 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -679,6 +679,7 @@ tests( rfactor_aligned_split_phases.cpp ring_buffer.cpp stream_compaction.cpp + thread_pool_keep_awake.cpp thread_safety.cpp tracing_thread_ids.cpp transitive_in.cpp diff --git a/test/correctness/thread_pool_keep_awake.cpp b/test/correctness/thread_pool_keep_awake.cpp new file mode 100644 index 000000000000..23d37cf7fbd9 --- /dev/null +++ b/test/correctness/thread_pool_keep_awake.cpp @@ -0,0 +1,184 @@ +#include "Halide.h" + +#include +#include +#include + +using namespace Halide; +using Halide::Internal::JITSharedRuntime; + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + fprintf(stderr, "%s:%d: Check failed: %s\n", __FILE__, __LINE__, #cond); \ + exit(1); \ + } \ + } while (0) + +namespace { + +using keep_awake_fn = int (*)(bool); + +// The number of references held through JITSharedRuntime. Only meaningful +// when no other thread is changing it. +int front_end_count() { + int c = JITSharedRuntime::thread_pool_keep_awake(true); + CHECK(c >= 1); + CHECK(JITSharedRuntime::thread_pool_keep_awake(false) == c - 1); + return c - 1; +} + +// The keep-awake count of a shared runtime's thread pool. +int runtime_count(keep_awake_fn f) { + int c = f(true); + CHECK(c >= 1); + CHECK(f(false) == c - 1); + return c - 1; +} + +keep_awake_fn find_keep_awake(const Target &t) { + return (keep_awake_fn)JITSharedRuntime::find_symbol(t, "halide_thread_pool_keep_awake"); +} + +// A small pipeline with nested parallelism and an async producer. +Callable make_pipeline(const Target &t, Param &offset) { + Var x, y, xo, xi; + Func producer, consumer, output; + producer(x, y) = x * y + offset; + consumer(x, y) = producer(x - 1, y) + producer(x + 1, y); + output(x, y) = consumer(x, y) + 1; + + output.parallel(y); + consumer.compute_at(output, y).split(x, xo, xi, 4).parallel(xo); + producer.compute_at(output, y).async(); + return output.compile_to_callable({offset}, t); +} + +void run(Callable &c, int offset) { + Buffer out(32, 8); + CHECK(c(offset, out) == 0); + out.for_each_element([&](int x, int y) { + CHECK(out(x, y) == 2 * x * y + 2 * offset + 1); + }); +} + +} // namespace + +int main(int argc, char **argv) { + Target t = get_jit_target_from_environment(); + if (t.arch == Target::WebAssembly) { + printf("[SKIP] WebAssembly does not support async() yet.\n"); + return 0; + } + + Param offset; + { + // Hold a reference before the shared runtime exists. It's applied + // when the runtime is created. + ThreadPoolKeepAwake early; + CHECK(front_end_count() == 1); + + Callable pipeline = make_pipeline(t, offset); + run(pipeline, 0); + keep_awake_fn keep_awake = find_keep_awake(t); + CHECK(keep_awake); + CHECK(runtime_count(keep_awake) == 1); + + // References may nest and live on other threads. The shared runtime + // holds one reference on behalf of all of them. + std::vector threads; + for (int i = 0; i < 4; i++) { + threads.emplace_back([&pipeline, i]() { + ThreadPoolKeepAwake a; + { + ThreadPoolKeepAwake b; + for (int j = 0; j < 100; j++) { + run(pipeline, i + j); + } + } + for (int j = 0; j < 100; j++) { + run(pipeline, i - j); + } + }); + } + for (auto &th : threads) { + th.join(); + } + CHECK(front_end_count() == 1); + CHECK(runtime_count(keep_awake) == 1); + + // Many tiny parallel loops, with and without an extra reference. + for (int i = 0; i < 500; i++) { + if (i % 100 < 50) { + ThreadPoolKeepAwake extra; + run(pipeline, i); + } else { + run(pipeline, i); + } + } + + // Change the number of threads while held. + int old_threads = JITSharedRuntime::set_num_threads(1); + for (int n : {1, 4, 2, 8, 3}) { + JITSharedRuntime::set_num_threads(n); + CHECK(JITSharedRuntime::get_num_threads() == n); + for (int i = 0; i < 50; i++) { + run(pipeline, i); + } + } + JITSharedRuntime::set_num_threads(old_threads); + + // Shut down the thread pool while held. The count survives. + using shutdown_fn = void (*)(); + shutdown_fn shutdown = + (shutdown_fn)JITSharedRuntime::find_symbol(t, "halide_shutdown_thread_pool"); + CHECK(shutdown); + shutdown(); + CHECK(runtime_count(keep_awake) == 1); + run(pipeline, 7); + + // Releasing the shared runtimes releases the old runtime's reference, + // and a new runtime picks it up again. The old runtime stays alive + // because the pipeline compiled with it still refers to it. + JITSharedRuntime::release_all(); + CHECK(runtime_count(keep_awake) == 0); + run(pipeline, 8); + CHECK(front_end_count() == 1); + + Callable new_pipeline = make_pipeline(t, offset); + run(new_pipeline, 9); + keep_awake_fn new_keep_awake = find_keep_awake(t); + CHECK(new_keep_awake && new_keep_awake != keep_awake); + CHECK(runtime_count(new_keep_awake) == 1); + CHECK(runtime_count(keep_awake) == 0); + + // Releasing a nested reference doesn't release the runtime's. + { + ThreadPoolKeepAwake nested; + } + CHECK(runtime_count(new_keep_awake) == 1); + } + CHECK(front_end_count() == 0); + keep_awake_fn keep_awake = find_keep_awake(t); + CHECK(keep_awake); + CHECK(runtime_count(keep_awake) == 0); + +#if HALIDE_WITH_EXCEPTIONS + if (Halide::exceptions_enabled()) { + // Releasing a reference that isn't held is an error. + bool threw = false; + try { + JITSharedRuntime::thread_pool_keep_awake(false); + } catch (const Halide::CompileError &e) { + printf("Expected error: %s\n", e.what()); + threw = true; + } + CHECK(threw); + CHECK(front_end_count() == 0); + CHECK(runtime_count(keep_awake) == 0); + } +#endif + + printf("Success!\n"); + return 0; +} diff --git a/test/generator/CMakeLists.txt b/test/generator/CMakeLists.txt index 24eea79fd321..354ea7d23349 100644 --- a/test/generator/CMakeLists.txt +++ b/test/generator/CMakeLists.txt @@ -1012,6 +1012,20 @@ _add_halide_aot_tests(shuffler) _add_halide_libraries(templated) _add_halide_aot_tests(templated) +# thread_pool_keep_awake_aottest.cpp +# thread_pool_keep_awake_generator.cpp +_add_halide_libraries( + thread_pool_keep_awake + # Requires threading support, not yet available for wasm tests + ENABLE_IF NOT ${_USING_WASM} +) +_add_halide_aot_tests( + thread_pool_keep_awake + # Requires threading support, not yet available for wasm tests + ENABLE_IF NOT ${_USING_WASM} + GROUPS multithreaded +) + # thread_pool_lost_wakeup_aottest.cpp # thread_pool_lost_wakeup_generator.cpp _add_halide_libraries( diff --git a/test/generator/thread_pool_keep_awake_aottest.cpp b/test/generator/thread_pool_keep_awake_aottest.cpp new file mode 100644 index 000000000000..37762f2403c4 --- /dev/null +++ b/test/generator/thread_pool_keep_awake_aottest.cpp @@ -0,0 +1,390 @@ +#include "HalideBuffer.h" +#include "HalideRuntime.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "thread_pool_keep_awake.h" + +using namespace Halide::Runtime; + +static_assert(!std::is_copy_constructible_v); +static_assert(!std::is_copy_assignable_v); +static_assert(!std::is_move_constructible_v); + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + fprintf(stderr, "%s:%d: Check failed: %s\n", __FILE__, __LINE__, #cond); \ + exit(1); \ + } \ + } while (0) + +namespace { + +std::atomic errors_reported{0}; +std::string last_error; + +void record_error(void *user_context, const char *msg) { + errors_reported++; + last_error = msg; +} + +// Only meaningful when no other thread is changing the count. +int keep_awake_count() { + int c = halide_thread_pool_keep_awake(true); + CHECK(c >= 1); + CHECK(halide_thread_pool_keep_awake(false) == c - 1); + return c - 1; +} + +// Busy-wait for about the given number of microseconds, to leave a gap +// between parallel loops like a real workload's serial work would. +void serial_gap(int us) { + auto start = std::chrono::steady_clock::now(); + while (std::chrono::steady_clock::now() - start < std::chrono::microseconds(us)) { + } +} + +void run_pipeline(int offset, int w = 32, int h = 8) { + Buffer out(w, h); + CHECK(thread_pool_keep_awake(offset, out) == 0); + for (int y = 0; y < h; y++) { + for (int x = 0; x < w; x++) { + CHECK(out(x, y) == 2 * x * y + 2 * offset + 1); + } + } +} + +struct ParForClosure { + int *hits; + int inner; + int fail_at; +}; + +int inner_task(void *user_context, int idx, uint8_t *closure) { + ((int *)closure)[idx]++; + return 0; +} + +int outer_task(void *user_context, int idx, uint8_t *closure) { + ParForClosure *c = (ParForClosure *)closure; + if (idx == c->fail_at) { + return halide_error_code_generic_error; + } + c->hits[idx]++; + if (c->inner > 0) { + // Nested parallelism. + std::vector inner_hits(c->inner, 0); + int ret = halide_do_par_for(user_context, inner_task, 0, c->inner, (uint8_t *)inner_hits.data()); + if (ret != 0) { + return ret; + } + for (int h : inner_hits) { + if (h != 1) { + return halide_error_code_generic_error; + } + } + } + return 0; +} + +void run_par_for(int n, int inner = 0) { + std::vector hits(n, 0); + ParForClosure c{hits.data(), inner, -1}; + CHECK(halide_do_par_for(nullptr, outer_task, 0, n, (uint8_t *)&c) == 0); + for (int h : hits) { + CHECK(h == 1); + } + + // Errors must still propagate. + if (n > 1) { + c.fail_at = n / 2; + CHECK(halide_do_par_for(nullptr, outer_task, 0, n, (uint8_t *)&c) == halide_error_code_generic_error); + } +} + +constexpr int pc_size = 64; + +struct ProducerConsumerState { + halide_semaphore_t sem; + int data[pc_size]; + int consumed[pc_size]; +}; + +int producer_task(void *user_context, int min, int extent, uint8_t *closure, void *task_parent) { + ProducerConsumerState *s = (ProducerConsumerState *)closure; + for (int i = min; i < min + extent; i++) { + if (i % 16 == 0) { + // Some nested parallelism inside a task. + std::vector inner_hits(8, 0); + if (halide_do_par_for(user_context, inner_task, 0, 8, (uint8_t *)inner_hits.data()) != 0) { + return halide_error_code_generic_error; + } + } + s->data[i] = 3 * i + 1; + halide_semaphore_release(&s->sem, 1); + } + return 0; +} + +int consumer_task(void *user_context, int min, int extent, uint8_t *closure, void *task_parent) { + ProducerConsumerState *s = (ProducerConsumerState *)closure; + for (int i = min; i < min + extent; i++) { + s->consumed[i] = s->data[i]; + } + return 0; +} + +void run_producer_consumer() { + ProducerConsumerState s; + memset(&s, 0, sizeof(s)); + halide_semaphore_init(&s.sem, 0); + + halide_semaphore_acquire_t acquire = {&s.sem, 1}; + halide_parallel_task_t tasks[2]; + memset(tasks, 0, sizeof(tasks)); + // The consumer goes first, so that threads find it blocked on its + // semaphore before the producer gets going. + tasks[0].fn = consumer_task; + tasks[0].closure = (uint8_t *)&s; + tasks[0].name = "consumer"; + tasks[0].semaphores = &acquire; + tasks[0].num_semaphores = 1; + tasks[0].min = 0; + tasks[0].extent = pc_size; + tasks[0].serial = true; + tasks[1].fn = producer_task; + tasks[1].closure = (uint8_t *)&s; + tasks[1].name = "producer"; + tasks[1].min = 0; + tasks[1].extent = pc_size; + tasks[1].serial = true; + + CHECK(halide_do_parallel_tasks(nullptr, 2, tasks, nullptr) == 0); + for (int i = 0; i < pc_size; i++) { + CHECK(s.consumed[i] == 3 * i + 1); + } +} + +void run_everything(int i) { + run_pipeline(i); + run_par_for(1 + i % 13); + if (i % 4 == 0) { + run_par_for(6, 3); + run_producer_consumer(); + } +} + +void test_counting() { + CHECK(keep_awake_count() == 0); + CHECK(halide_thread_pool_keep_awake(true) == 1); + CHECK(halide_thread_pool_keep_awake(true) == 2); + CHECK(halide_thread_pool_keep_awake(false) == 1); + CHECK(halide_thread_pool_keep_awake(false) == 0); + + { + ThreadPoolKeepAwake a; + CHECK(keep_awake_count() == 1); + { + ThreadPoolKeepAwake b; + CHECK(keep_awake_count() == 2); + { + ThreadPoolKeepAwake c; + CHECK(keep_awake_count() == 3); + } + CHECK(keep_awake_count() == 2); + } + CHECK(keep_awake_count() == 1); + } + CHECK(keep_awake_count() == 0); +} + +void test_unbalanced_release() { + halide_error_handler_t old_handler = halide_set_error_handler(record_error); + + errors_reported = 0; + CHECK(halide_thread_pool_keep_awake(false) == halide_error_code_generic_error); + CHECK(errors_reported == 1); + CHECK(last_error.find("halide_thread_pool_keep_awake") != std::string::npos); + // The count stays at zero, and works as normal afterwards. + CHECK(halide_thread_pool_keep_awake(true) == 1); + CHECK(halide_thread_pool_keep_awake(false) == 0); + CHECK(halide_thread_pool_keep_awake(false) == halide_error_code_generic_error); + CHECK(errors_reported == 2); + CHECK(keep_awake_count() == 0); + run_everything(0); + + halide_set_error_handler(old_handler); +} + +void test_many_small_loops() { + // Alternate between holding and not holding the count, with and + // without serial gaps between the loops, including one long enough + // for idle threads to give up and go to sleep anyway. + std::mt19937 rng(0); + for (int phase = 0; phase < 4; phase++) { + if (phase % 2) { + ThreadPoolKeepAwake keep_awake; + for (int i = 0; i < 300; i++) { + run_everything(i); + serial_gap(phase == 1 ? 0 : (int)(rng() % 200)); + } + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + run_everything(1); + } else { + for (int i = 0; i < 300; i++) { + run_everything(i); + serial_gap(phase == 0 ? 0 : (int)(rng() % 200)); + } + } + } + CHECK(keep_awake_count() == 0); +} + +void test_small_then_big() { + // While the count is held, workers that a small loop doesn't need are + // still demoted to the B team, where they spin rather than sleep. Mix small + // and big loops around acquiring and releasing the count; results must be + // correct throughout, and demoted threads must still pick up work. + halide_set_num_threads(0); + const int threads = halide_get_num_threads(); + for (int round = 0; round < 20; round++) { + // Without the count, a small loop after a big one demotes workers. + run_par_for(2 * threads); + run_par_for(1 + round % 2); + serial_gap(round % 3 == 0 ? 100 : 0); + { + ThreadPoolKeepAwake keep_awake; + for (int i = 0; i < 50; i++) { + run_par_for(1 + i % 2); + serial_gap(i % 5 == 0 ? 100 : 0); + run_par_for(2 * threads); + if (i % 10 == 0) { + run_par_for(4, 3); + run_producer_consumer(); + } + } + run_par_for(1); + } + // After releasing, idle workers sleep again as usual. + run_par_for(1); + run_par_for(2 * threads); + } + CHECK(keep_awake_count() == 0); +} + +void test_set_num_threads() { + halide_set_num_threads(0); + const int default_threads = halide_get_num_threads(); + ThreadPoolKeepAwake keep_awake; + for (int n : {1, 2, 4, 1, 8, 3, 16, 2, 0, 5}) { + halide_set_num_threads(n); + CHECK(halide_get_num_threads() == (n == 0 ? default_threads : n)); + for (int i = 0; i < 40; i++) { + run_everything(i); + serial_gap(i % 3 == 0 ? 50 : 0); + } + } + halide_set_num_threads(default_threads); + CHECK(keep_awake_count() == 1); +} + +void test_shutdown() { + { + ThreadPoolKeepAwake keep_awake; + // Shut down while idle threads are kept awake. + run_everything(0); + halide_shutdown_thread_pool(); + CHECK(keep_awake_count() == 1); + // The pool restarts, with its count still held. + run_everything(4); + halide_shutdown_thread_pool(); + CHECK(keep_awake_count() == 1); + // A shutdown with no pool running is a no-op. + halide_shutdown_thread_pool(); + run_everything(8); + } + CHECK(keep_awake_count() == 0); + halide_shutdown_thread_pool(); + run_everything(12); + + // Acquire the first reference while the pool is shut down. + halide_shutdown_thread_pool(); + { + ThreadPoolKeepAwake keep_awake; + run_everything(16); + } + halide_shutdown_thread_pool(); + CHECK(keep_awake_count() == 0); +} + +void test_concurrent() { + // Several threads use the pool and hold or release references + // concurrently, while other threads toggle the count and the number of + // threads. + std::atomic stop{false}; + halide_set_num_threads(0); + const int default_threads = halide_get_num_threads(); + + std::thread toggler([&]() { + while (!stop) { + ThreadPoolKeepAwake keep_awake; + serial_gap(20); + } + }); + + std::thread resizer([&]() { + std::mt19937 rng(1); + while (!stop) { + halide_set_num_threads(1 + rng() % 8); + serial_gap(100); + } + }); + + std::vector users; + for (int t = 0; t < 3; t++) { + users.emplace_back([t]() { + for (int i = 0; i < 100; i++) { + if ((i / 10 + t) % 2) { + ThreadPoolKeepAwake keep_awake; + run_everything(i); + } else { + run_everything(i); + } + } + }); + } + for (auto &u : users) { + u.join(); + } + stop = true; + toggler.join(); + resizer.join(); + halide_set_num_threads(default_threads); + CHECK(keep_awake_count() == 0); +} + +} // namespace + +int main(int argc, char **argv) { + test_counting(); + test_unbalanced_release(); + test_many_small_loops(); + test_small_then_big(); + test_set_num_threads(); + test_shutdown(); + test_concurrent(); + CHECK(keep_awake_count() == 0); + + printf("Success!\n"); + return 0; +} diff --git a/test/generator/thread_pool_keep_awake_generator.cpp b/test/generator/thread_pool_keep_awake_generator.cpp new file mode 100644 index 000000000000..e4a2def2855a --- /dev/null +++ b/test/generator/thread_pool_keep_awake_generator.cpp @@ -0,0 +1,28 @@ +#include "Halide.h" + +namespace { + +class ThreadPoolKeepAwake : public Halide::Generator { +public: + Input offset{"offset"}; + Output> output{"output"}; + + void generate() { + // A small pipeline with nested parallelism and an async producer, so + // that it exercises do_par_for, do_parallel_tasks, and semaphores. + Var x{"x"}, y{"y"}, xo{"xo"}, xi{"xi"}; + Func producer{"producer"}, consumer{"consumer"}; + + producer(x, y) = x * y + offset; + consumer(x, y) = producer(x - 1, y) + producer(x + 1, y); + output(x, y) = consumer(x, y) + 1; + + output.parallel(y); + consumer.compute_at(output, y).split(x, xo, xi, 4).parallel(xo); + producer.compute_at(output, y).async(); + } +}; + +} // namespace + +HALIDE_REGISTER_GENERATOR(ThreadPoolKeepAwake, thread_pool_keep_awake) diff --git a/test/performance/CMakeLists.txt b/test/performance/CMakeLists.txt index 5d1b6e38098b..94c567b07e19 100644 --- a/test/performance/CMakeLists.txt +++ b/test/performance/CMakeLists.txt @@ -51,6 +51,7 @@ tests( split_storage.cpp split_storage_bank_conflict.cpp stack_vs_heap.cpp + thread_pool_keep_awake.cpp thread_safe_jit_callable.cpp ) diff --git a/test/performance/thread_pool_keep_awake.cpp b/test/performance/thread_pool_keep_awake.cpp new file mode 100644 index 000000000000..85ed09c72797 --- /dev/null +++ b/test/performance/thread_pool_keep_awake.cpp @@ -0,0 +1,84 @@ +#include "Halide.h" + +#include +#include +#include +#include + +using namespace Halide; + +// Measures the fork/join overhead of small parallel loops separated by serial +// work, with and without Halide::ThreadPoolKeepAwake. Timing on shared CI +// machines is too noisy to assert on, so this only reports the numbers. + +namespace { + +using Clock = std::chrono::steady_clock; + +void busy_wait(int us) { + auto start = Clock::now(); + while (Clock::now() - start < std::chrono::microseconds(us)) { + } +} + +// Median time per call of a parallel loop of num_tasks trivial tasks, in +// microseconds, with gap_us of serial work between calls. +double fork_join_us(Callable &c, Buffer &out, int gap_us) { + const int samples = 2000; + std::vector times; + times.reserve(samples); + for (int i = 0; i < samples + 100; i++) { + auto t0 = Clock::now(); + if (c(out) != 0) { + printf("Pipeline failed\n"); + exit(1); + } + auto t1 = Clock::now(); + if (i >= 100) { + times.push_back(std::chrono::duration(t1 - t0).count()); + } + busy_wait(gap_us); + } + std::sort(times.begin(), times.end()); + return times[times.size() / 2]; +} + +} // namespace + +int main(int argc, char **argv) { + Target t = get_jit_target_from_environment(); + if (t.arch == Target::WebAssembly) { + printf("[SKIP] Performance tests are meaningless and/or misleading under WebAssembly interpreter.\n"); + return 0; + } + + Var x; + Func f; + f(x) = x * 2; + f.parallel(x); + Callable c = f.compile_to_callable({}, t); + + // One task per thread. The thread pool picks its size on first use. + Buffer out(1); + if (c(out) != 0) { + printf("Pipeline failed\n"); + return 1; + } + const int num_threads = std::max(Internal::JITSharedRuntime::get_num_threads(), 1); + out = Buffer(num_threads); + + printf("threads %d\n", num_threads); + printf("%8s %14s %14s\n", "gap (us)", "default (us)", "awake (us)"); + for (int gap : {0, 20, 100}) { + double asleep = fork_join_us(c, out, gap); + double awake; + { + ThreadPoolKeepAwake keep_awake; + awake = fork_join_us(c, out, gap); + } + printf("%8d %14.2f %14.2f\n", gap, asleep, awake); + } + + printf("Success!\n"); + return 0; +}