rlx_cpu/
pool.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Rayon-backed parallel for: `par_for(total, grain, |off, cnt| …)`.
17//!
18//! Replaces the old per-worker Condvar pool with Rayon's work-stealing
19//! scheduler. Same `(offset, count)` chunk API so all existing call
20//! sites (BLAS tiling, SDPA, LayerNorm, …) pick up Rayon without
21//! changes.
22
23#[cfg(not(target_arch = "wasm32"))]
24use rayon::prelude::*;
25#[cfg(not(target_arch = "wasm32"))]
26use std::sync::Once;
27
28#[cfg(not(target_arch = "wasm32"))]
29static POOL_INIT: Once = Once::new();
30
31#[cfg(not(target_arch = "wasm32"))]
32fn ensure_pool() {
33    POOL_INIT.call_once(|| {
34        let cfg = crate::config::RuntimeConfig::global();
35        let n = cfg.pool_workers.max(1);
36        let _ = rayon::ThreadPoolBuilder::new()
37            .num_threads(n)
38            .thread_name(|i| format!("rlx-rayon-{i}"))
39            .build_global();
40    });
41}
42
43/// Total Rayon worker count (configured from [`RuntimeConfig::pool_workers`]).
44///
45/// On wasm there is no thread pool — the browser is single-threaded — so
46/// this is always 1.
47#[cfg(not(target_arch = "wasm32"))]
48pub fn num_threads() -> usize {
49    ensure_pool();
50    rayon::current_num_threads()
51}
52
53#[cfg(target_arch = "wasm32")]
54pub fn num_threads() -> usize {
55    1
56}
57
58/// Parallel for: split `total` items across threads. `f(off, cnt)` is
59/// called once per chunk with disjoint regions.
60///
61/// SAFETY: caller must ensure `f` accesses disjoint memory regions for
62/// different `(offset, count)` pairs.
63/// Parallel `for i in 0..n` over an index range, one task per index.
64/// Serial on wasm (no thread pool). The closure must be `Sync + Send`
65/// because tasks may run on Rayon workers natively.
66#[inline]
67pub fn par_range<F: Fn(usize) + Sync + Send>(n: usize, f: F) {
68    #[cfg(target_arch = "wasm32")]
69    {
70        (0..n).for_each(f);
71    }
72    #[cfg(not(target_arch = "wasm32"))]
73    {
74        (0..n).into_par_iter().for_each(f);
75    }
76}
77
78/// Minimum elements that make threading an element-wise kernel worthwhile.
79/// Below this the rayon hand-off costs more than the work it saves. Kept in
80/// one place so conv/transpose/pool/relu/region kernels share a single,
81/// tunable cutover instead of sprinkling magic constants at each call site.
82pub const ELEMENTWISE_PAR_FLOOR: usize = 4096;
83
84/// Should an op over `work` elements run in parallel? True only with >1 worker
85/// and enough work to keep them busy past the amortization floor. Use this at
86/// every data-parallel kernel call site instead of a hand-picked threshold.
87#[inline]
88pub fn should_parallelize(work: usize) -> bool {
89    num_threads() > 1 && work >= ELEMENTWISE_PAR_FLOOR * 2
90}
91
92/// `par_for` chunk floor for a flat range of `total` elements: aim for one
93/// chunk per worker, never below the amortization floor.
94#[inline]
95pub fn chunk_floor(total: usize) -> usize {
96    (total / num_threads().max(1)).max(ELEMENTWISE_PAR_FLOOR)
97}
98
99/// `par_for` chunk floor for an *outer* loop of `units` independent items
100/// (e.g. (n,c) planes): one chunk per worker, at least one unit.
101#[inline]
102pub fn outer_chunk(units: usize) -> usize {
103    (units / num_threads().max(1)).max(1)
104}
105
106#[inline]
107pub fn par_for<F: Fn(usize, usize) + Sync>(total: usize, min_per_thread: usize, f: &F) {
108    if total == 0 {
109        return;
110    }
111    // wasm is single-threaded: run the whole range inline. (Rayon's
112    // work-stealing scheduler needs OS threads, which the browser lacks.)
113    #[cfg(target_arch = "wasm32")]
114    {
115        let _ = min_per_thread;
116        f(0, total);
117    }
118    #[cfg(not(target_arch = "wasm32"))]
119    {
120        ensure_pool();
121        let grain = min_per_thread.max(1);
122        let n_threads = (total / grain).max(1).min(num_threads());
123        if n_threads <= 1 {
124            f(0, total);
125            return;
126        }
127        let chunk = total.div_ceil(n_threads);
128        (0..n_threads).into_par_iter().for_each(|t| {
129            let off = t * chunk;
130            if off < total {
131                f(off, (off + chunk).min(total) - off);
132            }
133        });
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140    use std::sync::atomic::{AtomicU64, Ordering};
141
142    #[test]
143    fn par_for_sums_correctly() {
144        let data = vec![1.0f32; 10_000];
145        let total = AtomicU64::new(0);
146
147        par_for(data.len(), 100, &|off, cnt| {
148            let partial: f32 = data[off..off + cnt].iter().sum();
149            total.fetch_add(partial.to_bits() as u64, Ordering::Relaxed);
150        });
151
152        assert!(total.load(Ordering::Relaxed) > 0);
153    }
154
155    #[test]
156    fn par_for_small_is_sequential() {
157        let sum = std::sync::atomic::AtomicUsize::new(0);
158        par_for(10, 100, &|off, cnt| {
159            sum.fetch_add(cnt, Ordering::Relaxed);
160            assert_eq!(off + cnt, 10);
161        });
162        assert_eq!(sum.load(Ordering::Relaxed), 10);
163    }
164
165    #[test]
166    fn par_for_exact_sum_many_dispatches() {
167        for &n in &[256usize, 1024, 4097] {
168            let sum = std::sync::atomic::AtomicUsize::new(0);
169            par_for(n, 256, &|off, cnt| {
170                sum.fetch_add(cnt, Ordering::Relaxed);
171                assert!(off + cnt <= n);
172            });
173            assert_eq!(sum.load(Ordering::Relaxed), n);
174        }
175    }
176
177    #[test]
178    fn par_for_concurrent_callers_isolated() {
179        std::thread::scope(|s| {
180            for t in 0..4 {
181                s.spawn(move || {
182                    let n = 4096 + t * 17;
183                    let sum = std::sync::atomic::AtomicUsize::new(0);
184                    par_for(n, 128, &|off, cnt| {
185                        sum.fetch_add(cnt, Ordering::Relaxed);
186                        assert!(off + cnt <= n);
187                    });
188                    assert_eq!(sum.load(Ordering::Relaxed), n);
189                });
190            }
191        });
192    }
193}
rlx_cpu/pool.rs

rlx_cpu/
pool.rs
