-
Notifications
You must be signed in to change notification settings - Fork 42
Fan out compute-bound row-scan kernels from 128k rows #424
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
b8ae741
7c59e08
c11d71b
226afcf
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -2502,7 +2502,7 @@ pub fn m4_indices(x: &[f64], y: &[f64], x0: f64, x1: f64, n_buckets: usize) -> V | |
| n_buckets, | ||
| start, | ||
| end, | ||
| par_threads(end - start), | ||
| par_threads_above(end - start, PAR_THRESHOLD_COMPUTE).min(n_buckets), | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. P3: The per-cost-class routing docs updated in this PR enumerate the worker caps for Prompt for AI agents |
||
| ) | ||
| } | ||
|
|
||
|
|
@@ -3174,22 +3174,29 @@ fn bin_2d_count_f32_scalar( | |
| } | ||
| } | ||
|
|
||
| /// Row-scan kernels fan out across cores only past this size, where thread | ||
| /// Row-scan kernels fan out across cores only past these sizes, where thread | ||
| /// spawn + merge cost is well amortized. Threading stays inside the call — | ||
| /// the ABI remains synchronous (engine doc E5). | ||
| /// the ABI remains synchronous (engine doc E5). Compute-bound scans cross over | ||
| /// earlier: at ~2.5–3 ns/row serial, M4 and uniform histogram amortize a | ||
| /// fan-out from ~128k rows, the same crossover zone maps use. | ||
| const PAR_THRESHOLD: usize = 1 << 19; | ||
| const PAR_THRESHOLD_COMPUTE: usize = 1 << 17; | ||
|
|
||
| fn par_threads(n: usize) -> usize { | ||
| par_threads_above(n, PAR_THRESHOLD) | ||
| } | ||
|
|
||
| fn par_threads_above(n: usize, threshold: usize) -> usize { | ||
| static CODSPEED: std::sync::OnceLock<bool> = std::sync::OnceLock::new(); | ||
| if *CODSPEED.get_or_init(|| std::env::var_os("CODSPEED_ENV").is_some()) { | ||
| return 1; | ||
| } | ||
| let cores = std::thread::available_parallelism().map_or(1, |p| p.get()); | ||
| par_threads_for(n, cores) | ||
| par_threads_above_for(n, threshold, cores) | ||
| } | ||
|
|
||
| fn par_threads_for(n: usize, cores: usize) -> usize { | ||
| if n >= PAR_THRESHOLD { | ||
| fn par_threads_above_for(n: usize, threshold: usize, cores: usize) -> usize { | ||
| if n >= threshold { | ||
| cores.clamp(1, MAX_ROW_THREADS) | ||
| } else { | ||
| 1 | ||
|
|
@@ -4596,7 +4603,14 @@ pub fn stratified_sample_mask<T: Copy + Sync + Into<u64>>( | |
| pub fn histogram_uniform(data: &[f64], lo: f64, hi: f64, out: &mut [f64]) -> u64 { | ||
| assert!(x1_gt_x0(lo, hi)); | ||
| assert!(!out.is_empty()); | ||
| histogram_uniform_impl(data, lo, hi, par_threads(data.len()), out) | ||
| histogram_uniform_impl(data, lo, hi, histogram_threads(data.len(), out.len()), out) | ||
| } | ||
|
|
||
| /// Workers for a uniform histogram. Each owns a full bin vector and the merge | ||
| /// is `threads * n_bins`, so points per bin caps the row gate — same rule as | ||
| /// `bin_2d_threads`, whose per-thread grids are this kernel's per-thread bins. | ||
| fn histogram_threads(n: usize, n_bins: usize) -> usize { | ||
| (n / n_bins.max(1)).clamp(1, par_threads_above(n, PAR_THRESHOLD_COMPUTE)) | ||
| } | ||
|
|
||
| /// Per-bin u64 counting shared by the serial and parallel paths. Stays scalar | ||
|
|
@@ -6854,9 +6868,70 @@ mod tests { | |
| // shape): per-thread grids + merge dwarf the scan — stay serial. | ||
| assert_eq!(bin_2d_threads(1 << 20, 1 << 20), 1); | ||
| assert_eq!(bin_2d_threads(2_100_000, 2048 * 2048), 1); | ||
| assert_eq!(par_threads_for(PAR_THRESHOLD - 1, 64), 1); | ||
| assert_eq!(par_threads_for(PAR_THRESHOLD, 64), MAX_ROW_THREADS); | ||
| assert_eq!(par_threads_for(PAR_THRESHOLD, 1), 1); | ||
| assert_eq!( | ||
| par_threads_above_for(PAR_THRESHOLD - 1, PAR_THRESHOLD, 64), | ||
| 1 | ||
| ); | ||
| assert_eq!( | ||
| par_threads_above_for(PAR_THRESHOLD, PAR_THRESHOLD, 64), | ||
| MAX_ROW_THREADS | ||
| ); | ||
| assert_eq!(par_threads_above_for(PAR_THRESHOLD, PAR_THRESHOLD, 1), 1); | ||
| } | ||
|
|
||
| #[test] | ||
| fn histogram_threads_bin_aware() { | ||
| let n = PAR_THRESHOLD_COMPUTE; | ||
| let cap = par_threads_above(n, PAR_THRESHOLD_COMPUTE); | ||
| // Below the compute gate: serial regardless of bin count. | ||
| assert_eq!(histogram_threads(PAR_THRESHOLD_COMPUTE - 1, 8), 1); | ||
| // Chart-sized bin counts (512 is the benchmark config): pure row gate. | ||
| assert_eq!(histogram_threads(n, 512), cap); | ||
| assert_eq!(histogram_threads(n, 8), cap); | ||
| // Bins as numerous as the rows: merge dwarfs the scan — stay serial. | ||
| assert_eq!(histogram_threads(n, n), 1); | ||
| assert_eq!(histogram_threads(200_000, 1_000_000), 1); | ||
| // In between, workers track points per bin. | ||
| assert_eq!(histogram_threads(1 << 18, 1 << 15), 8.min(cap)); | ||
| } | ||
|
|
||
| #[test] | ||
| fn histogram_high_bin_count_matches_serial() { | ||
| // 128k-512k band: high bin counts route serial, and the public path | ||
| // stays bit-identical to the serial impl either way. | ||
| let n = PAR_THRESHOLD_COMPUTE + 1_000; | ||
| let bins = 300_000; | ||
| let data: Vec<f64> = (0..n).map(|i| (i % 977) as f64 * 0.25).collect(); | ||
| assert_eq!(histogram_threads(n, bins), 1); | ||
| let mut serial = vec![0f64; bins]; | ||
| let ts = histogram_uniform_impl(&data, 0.0, 244.0, 1, &mut serial); | ||
| let mut routed = vec![0f64; bins]; | ||
| let tr = histogram_uniform(&data, 0.0, 244.0, &mut routed); | ||
| assert_eq!(ts, tr); | ||
| assert_eq!(serial, routed); | ||
| // Same rows, ordinary bin count: still fans out. | ||
| let cap = par_threads_above(n, PAR_THRESHOLD_COMPUTE); | ||
| assert_eq!(histogram_threads(n, 512), cap); | ||
| let mut small_serial = vec![0f64; 512]; | ||
| let ss = histogram_uniform_impl(&data, 0.0, 244.0, 1, &mut small_serial); | ||
| let mut small_routed = vec![0f64; 512]; | ||
| let sr = histogram_uniform(&data, 0.0, 244.0, &mut small_routed); | ||
| assert_eq!(ss, sr); | ||
| assert_eq!(small_serial, small_routed); | ||
| } | ||
|
|
||
| #[test] | ||
| fn par_thread_gates_route_by_cost_class() { | ||
| const { assert!(PAR_THRESHOLD_COMPUTE < PAR_THRESHOLD) } | ||
| let t_c = PAR_THRESHOLD_COMPUTE; | ||
| for cores in [1, 4, 20, 64] { | ||
| let cap = cores.clamp(1, MAX_ROW_THREADS); | ||
| // Compute-bound kernels fan out from 128k rows. | ||
| assert_eq!(par_threads_above_for(t_c - 1, t_c, cores), 1); | ||
| assert_eq!(par_threads_above_for(t_c, t_c, cores), cap); | ||
| // ...and stay serial in the band the general gate still reserves. | ||
| assert_eq!(par_threads_above_for(t_c, PAR_THRESHOLD, cores), 1); | ||
| } | ||
| } | ||
|
|
||
| #[test] | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
P2: Clustered or duplicate-x M4 windows in the new 128k–512k fan-out band can regress badly: every proposed split scans through the same giant bucket, then only one segment is executed. Consider finding bucket boundaries with a single pass/binary search, or falling back before paying these repeated scans.
Prompt for AI agents