Don't spin on the main mutex while waiting for new work

src/runtime/HalideRuntime.h

@@ -194,7 +194,7 @@ struct halide_mutex {
 
 /** Cross platform condition variable. Must be initialized to 0. */
 struct halide_cond {
-    uintptr_t _private[1];
+    uintptr_t _private[2];
 };
 
 /** A basic set of mutex and condition variable functions, which call

src/runtime/synchronization_common.h

@@ -816,11 +816,15 @@ struct wait_parking_control final : public parking_control {
 
 class fast_cond {
     uintptr_t state = 0;
+    uintptr_t counter = 0;
 
 public:
     ALWAYS_INLINE void signal() {
         if_tsan_pre_signal(this);
 
+        // Release any spinning waiters
+        atomic_fetch_add_acquire_release(&counter, (uintptr_t)1);
+
         uintptr_t val;
         atomic_load_relaxed(&state, &val);
         if (val == 0) {
@@ -834,6 +838,10 @@ class fast_cond {
 
     ALWAYS_INLINE void broadcast() {
         if_tsan_pre_signal(this);
+
+        // Release any spinning waiters
+        atomic_fetch_add_acquire_release(&counter, (uintptr_t)1);
+
         uintptr_t val;
         atomic_load_relaxed(&state, &val);
         if (val == 0) {
@@ -846,6 +854,35 @@ class fast_cond {
     }
 
     ALWAYS_INLINE void wait(fast_mutex *mutex) {
+        // Spin for a bit, waiting to see if someone else calls signal or
+        // broadcast.
+        uintptr_t initial;
+        atomic_load_relaxed(&counter, &initial);
+        mutex->unlock();
+        spin_control spinner;
+        while (spinner.should_spin()) {
+            halide_thread_yield();
+            uintptr_t current;
+            atomic_load_relaxed(&counter, &current);
+            if (current != initial) {
+                mutex->lock();
+                return;
+            }
+        }
+
+        mutex->lock();  // Can locking and then immediately waiting be
+                        // optimized?
+
+        // Check one final time with the lock held. This guarantees we won't
+        // miss an increment of the counter because it is only ever incremented
+        // with the lock held.
+        uintptr_t current;
+        atomic_load_relaxed(&counter, &current);
+        if (current != initial) {
+            return;
+        }
+
+        // Go to sleep until signaled
         wait_parking_control control(&state, mutex);
         uintptr_t result = control.park((uintptr_t)this);
         if (result != (uintptr_t)mutex) {

src/runtime/thread_pool_common.h

@@ -203,9 +203,6 @@ WEAK void dump_job_state() {
 WEAK void worker_thread(void *);
 
 WEAK void worker_thread_already_locked(work *owned_job) {
-    int spin_count = 0;
-    const int max_spin_count = 40;
-
     while (owned_job ? owned_job->running() : !work_queue.shutdown) {
         work *job = work_queue.jobs;
         work **prev_ptr = &work_queue.jobs;
@@ -283,38 +280,24 @@ WEAK void worker_thread_already_locked(work *owned_job) {
         if (!job) {
             // There is no runnable job. Go to sleep.
             if (owned_job) {
-                if (spin_count++ < max_spin_count) {
-                    // Give the workers a chance to finish up before sleeping
-                    halide_mutex_unlock(&work_queue.mutex);
-                    halide_thread_yield();
-                    halide_mutex_lock(&work_queue.mutex);
-                } else {
-                    work_queue.owners_sleeping++;
-                    owned_job->owner_is_sleeping = true;
-                    halide_cond_wait(&work_queue.wake_owners, &work_queue.mutex);
-                    owned_job->owner_is_sleeping = false;
-                    work_queue.owners_sleeping--;
-                }
+                work_queue.owners_sleeping++;
+                owned_job->owner_is_sleeping = true;
+                halide_cond_wait(&work_queue.wake_owners, &work_queue.mutex);
+                owned_job->owner_is_sleeping = false;
+                work_queue.owners_sleeping--;
             } else {
                 work_queue.workers_sleeping++;
                 if (work_queue.a_team_size > work_queue.target_a_team_size) {
                     // Transition to B team
                     work_queue.a_team_size--;
                     halide_cond_wait(&work_queue.wake_b_team, &work_queue.mutex);
                     work_queue.a_team_size++;
-                } else if (spin_count++ < max_spin_count) {
-                    // Spin waiting for new work
-                    halide_mutex_unlock(&work_queue.mutex);
-                    halide_thread_yield();
-                    halide_mutex_lock(&work_queue.mutex);
                 } else {
                     halide_cond_wait(&work_queue.wake_a_team, &work_queue.mutex);
                 }
                 work_queue.workers_sleeping--;
             }
             continue;
-        } else {
-            spin_count = 0;
         }
 
         log_message("Working on job " << job->task.name);

test/performance/parallel_scenarios.cpp

@@ -31,31 +31,36 @@ int main(int argc, char **argv) {
 
     int native_threads = Halide::Internal::JITSharedRuntime::get_num_threads();
 
+    std::map<std::tuple<bool, bool, int, int>, std::vector<float>> results;
+
     auto bench = [&](bool m, bool c, int i, int o) {
-        const int num_samples = 128;
         const int memory_limit = m ? max_memory : 128;
 
+        auto now = std::chrono::high_resolution_clock::now;
+        auto to_ns = [](auto delta) { return 1e9 * std::chrono::duration<float>(delta).count(); };
+
         auto bench_one = [&]() {
-            auto t1 = std::chrono::high_resolution_clock::now();
+            auto t1 = now();
             callable(i, o, memory_limit, in, out);
-            auto t2 = std::chrono::high_resolution_clock::now();
-            return 1e9 * std::chrono::duration<float>(t2 - t1).count() / (i * o);
+            auto t2 = now();
+            return to_ns(t2 - t1) / (i * o);
         };
 
-        std::vector<float> times(num_samples);
+        const int num_tasks = 8;
+        const int min_samples = 32;
+
+        std::vector<float> times[num_tasks];
         if (c) {
             Halide::Tools::ThreadPool<void> thread_pool;
-            const int num_tasks = 8;
-            const int samples_per_task = num_samples / num_tasks;
             Halide::Internal::JITSharedRuntime::set_num_threads(num_tasks * native_threads);
             std::vector<std::future<void>> futures(num_tasks);
             for (size_t t = 0; t < futures.size(); t++) {
                 futures[t] = thread_pool.async(
                     [&](size_t t) {
                         bench_one();
-                        for (int s = 0; s < samples_per_task; s++) {
-                            size_t idx = t * samples_per_task + s;
-                            times[idx] = bench_one();
+                        auto t_start = now();
+                        while (to_ns(now() - t_start) < 1e7 || times[t].size() < min_samples / num_tasks) {
+                            times[t].push_back(bench_one());
                         }
                     },
                     t);
@@ -66,32 +71,43 @@ int main(int argc, char **argv) {
         } else {
             Halide::Internal::JITSharedRuntime::set_num_threads(native_threads);
             bench_one();
-            for (int s = 0; s < num_samples; s++) {
-                times[s] = bench_one();
+            auto t_start = now();
+            while (to_ns(now() - t_start) < 1e7 || times[0].size() < min_samples) {
+                times[0].push_back(bench_one());
             }
         }
-        std::sort(times.begin(), times.end());
-        printf("%d %d %d %d ", m, c, i, o);
-        const int n = 8;
-        int off = (num_samples / n) / 2;
-        for (int i = 0; i < n; i++) {
-            printf("%g ", times[off + (num_samples * i) / n]);
+
+        std::vector<float> &r = results[{m, c, i, o}];
+        for (int i = 0; i < num_tasks; i++) {
+            r.insert(r.end(), times[i].begin(), times[i].end());
         }
-        printf("\n");
     };
 
     // The output is designed to be copy-pasted into a spreadsheet, not read by a human
-    printf("memory_bound contended inner outer t0 t1 t2 t3 t4 t5 t7\n");
-    for (bool contended : {false, true}) {
-        for (bool memory_bound : {false, true}) {
-            for (int i : {1 << 0, 1 << 6, 1 << 12, 1 << 18}) {
-                for (int o : {1, 2, 4, 8, 16, 32, 64, 128, 256}) {
-                    bench(memory_bound, contended, i, o);
+    printf("memory_bound contended inner outer num_samples 10%% 20%% 30%% 40%% 50%% 60%% 70%% 80%% 90%%\n");
+    for (int repeat = 0; repeat < 10; repeat++) {
+        for (bool contended : {false, true}) {
+            for (bool memory_bound : {false, true}) {
+                for (int i : {1 << 6, 1 << 9, 1 << 12, 1 << 15}) {
+                    for (int o : {1, 2, 4, 8, 16, 32, 64, 128, 256}) {
+                        bench(memory_bound, contended, i, o);
+                    }
                 }
             }
         }
     }
 
+    for (auto p : results) {
+        auto &times = p.second;
+        std::sort(times.begin(), times.end());
+        auto [m, c, i, o] = p.first;
+        printf("%d %d %d %d %d ", m, c, i, o, (int)times.size());
+        for (int decile = 10; decile <= 90; decile += 10) {
+            printf("%g ", times[(decile * times.size()) / 100]);
+        }
+        printf("\n");
+    }
+
     printf("Success!\n");
 
     return 0;