perf(analysis): IDW and kriging rows on several cores

crates/analysis/src/kriging.rs

@@ -313,6 +313,19 @@ pub fn kriging(
     params: &VariogramParams,
     neighbours: usize,
     name: &str,
+) -> Result<Raster, String> {
+    kriging_on(None, points, extent, cell_deg, params, neighbours, name)
+}
+
+/// [`kriging`] on `threads` threads, or as many as its work calls for.
+fn kriging_on(
+    threads: Option<usize>,
+    points: &[(Coord, f64)],
+    extent: &Extent,
+    cell_deg: f64,
+    params: &VariogramParams,
+    neighbours: usize,
+    name: &str,
 ) -> Result<Raster, String> {
     if points.len() < 2 {
         return Err("kriging needs at least two points with values".into());
@@ -334,14 +347,17 @@ pub fn kriging(
         return Err(format!("{w} × {h} cells from {} points is too much work; use a larger cell size", points.len()));
     }
 
-    let mut pred_data = Vec::with_capacity(w * h);
-    let mut var_data = Vec::with_capacity(w * h);
-    let mut near: Vec<(f64, f64, Coord)> = Vec::with_capacity(points.len());
-    let mut buffers = KrigingBuffers::new();
-
-    for row in 0..h {
+    // Each row on its own, on several threads when the work is large (#1252); a thread has its
+    // own distances and system buffers.
+    let scratch_bytes = points.len().saturating_mul(std::mem::size_of::<(f64, f64, Coord)>());
+    let threads =
+        threads.unwrap_or_else(|| crate::parallel::threads_for_rows(h, w.saturating_mul(h).saturating_mul(work_per_cell), scratch_bytes));
+    let scratch = || (Vec::<(f64, f64, Coord)>::with_capacity(points.len()), KrigingBuffers::new());
+    let (mut pred_data, mut var_data) = (vec![NODATA; w * h], vec![NODATA; w * h]);
+    let rows = pred_data.chunks_mut(w).zip(var_data.chunks_mut(w));
+    crate::parallel::for_each_row(rows, threads, scratch, |(near, buffers), row, (pred, var): (&mut [f32], &mut [f32])| {
         let kx = (extent.max_y - (row as f64 + 0.5) * cell_deg).to_radians().cos().abs().max(1e-6);
-        for col in 0..w {
+        for (col, (pred, var)) in pred.iter_mut().zip(var.iter_mut()).enumerate() {
             let c = Coord::new(extent.min_x + (col as f64 + 0.5) * cell_deg, extent.max_y - (row as f64 + 0.5) * cell_deg);
             near.clear();
             near.extend(points.iter().map(|(p, v)| {
@@ -356,11 +372,11 @@ pub fn kriging(
                 near.truncate(k);
             }
 
-            let (val, var) = ordinary_kriging_point(c, &near, params, kx, &mut buffers);
-            pred_data.push(if val.is_finite() { val as f32 } else { NODATA });
-            var_data.push(if var.is_finite() { var as f32 } else { NODATA });
+            let (val, var_est) = ordinary_kriging_point(c, near, params, kx, buffers);
+            *pred = if val.is_finite() { val as f32 } else { NODATA };
+            *var = if var_est.is_finite() { var_est as f32 } else { NODATA };
         }
-    }
+    });
 
     let gt = GeoTransform { origin_x: extent.min_x, origin_y: extent.max_y, pixel_width: cell_deg, pixel_height: -cell_deg };
     Raster::new(name, w as u32, h as u32, vec![pred_data, var_data], gt, Crs::LonLat, Some(NODATA as f64), ColorHint::Gray, "Float32")
@@ -436,6 +452,26 @@ mod tests {
         }
     }
 
+    /// Rows made on several threads give the same surfaces, bit for bit, as one thread (#1252).
+    #[test]
+    fn kriging_on_several_threads_is_kriging_on_one() {
+        let pts: Vec<(Coord, f64)> = (0..300)
+            .map(|i| {
+                let (x, y) = ((i * 37 % 300) as f64 / 30.0, (i * 91 % 300) as f64 / 30.0);
+                (Coord::new(x, 40.0 + y), (x * 1.3).sin() * 100.0 + y * 7.0)
+            })
+            .collect();
+        let extent = Extent::new(0.0, 40.0, 10.0, 50.0);
+        let params = auto_variogram(&pts, &extent, VariogramModel::Spherical);
+        let one = kriging_on(Some(1), &pts, &extent, 0.2, &params, 12, "k").unwrap();
+        for threads in [2, 7, 16] {
+            let many = kriging_on(Some(threads), &pts, &extent, 0.2, &params, 12, "k").unwrap();
+            for (a, b) in one.bands.iter().zip(&many.bands) {
+                assert!(a.data.iter().map(|v| v.to_bits()).eq(b.data.iter().map(|v| v.to_bits())), "{threads} threads");
+            }
+        }
+    }
+
     #[test]
     fn test_high_latitude_metric_consistency() {
         let lat = 60.0;

crates/analysis/src/parallel.rs

@@ -12,13 +12,77 @@ pub(crate) const BATCH: usize = 16;
 /// Each worker's stack, as large as the main thread's usually is.
 const STACK_BYTES: usize = 8 << 20;
 
+/// The work below which a raster tool's rows are made on one thread: cell-point pairs, say,
+/// about 30 ns each for IDW, so some 5 ms of it, against tens of microseconds to start a thread.
+pub(crate) const MIN_ROW_WORK: usize = 200_000;
+/// The scratch memory a raster tool's threads may take together, beyond one thread's.
+const SCRATCH_BUDGET: usize = 256 << 20;
+
 /// How many threads to make `items` items on: one below `min_items`, always one in the browser,
 /// and no more than there are batches.
 pub(crate) fn threads_for(items: usize, min_items: usize) -> usize {
     if cfg!(target_arch = "wasm32") || items < min_items {
         return 1;
     }
-    std::thread::available_parallelism().map_or(1, |n| n.get()).min(MAX_THREADS).min(items.div_ceil(BATCH)).max(1)
+    cores().min(items.div_ceil(BATCH)).max(1)
+}
+
+/// How many threads to make `rows` raster rows on, `work` in all, each thread taking
+/// `scratch_bytes` of scratch: one below [`MIN_ROW_WORK`] and in the browser, no more than there
+/// are rows, and no more than [`SCRATCH_BUDGET`] holds (a scratch per point, for millions of
+/// points, is large).
+pub(crate) fn threads_for_rows(rows: usize, work: usize, scratch_bytes: usize) -> usize {
+    if cfg!(target_arch = "wasm32") || work < MIN_ROW_WORK {
+        return 1;
+    }
+    let affordable = 1 + SCRATCH_BUDGET / scratch_bytes.max(1);
+    cores().min(rows).min(affordable).max(1)
+}
+
+/// The cores to use, at most [`MAX_THREADS`].
+fn cores() -> usize {
+    std::thread::available_parallelism().map_or(1, |n| n.get()).min(MAX_THREADS)
+}
+
+/// `f` on each of a raster's `rows` (its output rows, `chunks_mut` of the bands it writes) with
+/// the row's number, on `threads` threads ([`threads_for_rows`]), this one among them. Each row
+/// is written where it lies, so the output is held once, as on one thread. A row is costly
+/// enough to be taken one at a time, so rows that cost more than others are shared out too.
+/// Each thread makes its own `scratch()` once; one that panics stops the others after the rows
+/// they are making, and its panic is raised again here.
+pub(crate) fn for_each_row<T: Send, S>(
+    rows: impl Iterator<Item = T> + Send,
+    threads: usize,
+    scratch: impl Fn() -> S + Sync,
+    f: impl Fn(&mut S, usize, T) + Sync,
+) {
+    use std::sync::atomic::{AtomicBool, Ordering};
+    if threads <= 1 {
+        let mut s = scratch();
+        rows.enumerate().for_each(|(i, row)| f(&mut s, i, row));
+        return;
+    }
+    let rows = std::sync::Mutex::new(rows.enumerate());
+    let failed = AtomicBool::new(false);
+    let work = || {
+        let _stop = StopOnPanic(&failed);
+        let mut s = scratch();
+        while !failed.load(Ordering::Relaxed) {
+            // The lock is held only to take a row; a poisoned one means a thread panicked.
+            let Some((i, row)) = rows.lock().ok().and_then(|mut r| r.next()) else { break };
+            f(&mut s, i, row);
+        }
+    };
+    std::thread::scope(|scope| {
+        let others: Vec<_> =
+            (1..threads).filter_map(|_| std::thread::Builder::new().stack_size(STACK_BYTES).spawn_scoped(scope, work).ok()).collect();
+        work();
+        for h in others {
+            if let Err(e) = h.join() {
+                std::panic::resume_unwind(e);
+            }
+        }
+    });
 }
 
 /// `f` on each of `items` with its index, the results in their order: on several threads when
@@ -45,26 +109,51 @@ pub(crate) fn try_map_on<T: Sync, S, R: Send, E: Send>(
     threads: usize,
     scratch: impl Fn() -> S + Sync,
     f: impl Fn(&mut S, usize, &T) -> Result<R, E> + Sync,
+) -> Result<Vec<R>, E> {
+    try_each_on(items.len(), threads, scratch, |s, i| f(s, i, &items[i]))
+}
+
+/// Sets the stop flag when the thread it lives on panics, so the others stop after the items
+/// they are making rather than make the rest for nothing.
+struct StopOnPanic<'a>(&'a std::sync::atomic::AtomicBool);
+
+impl Drop for StopOnPanic<'_> {
+    fn drop(&mut self) {
+        if std::thread::panicking() {
+            self.0.store(true, std::sync::atomic::Ordering::Relaxed);
+        }
+    }
+}
+
+/// [`try_map_on`] over the indices `0..len`.
+fn try_each_on<S, R: Send, E: Send>(
+    len: usize,
+    threads: usize,
+    scratch: impl Fn() -> S + Sync,
+    f: impl Fn(&mut S, usize) -> Result<R, E> + Sync,
 ) -> Result<Vec<R>, E> {
     use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
-    if threads <= 1 || items.len() <= 1 {
+    if threads <= 1 || len <= 1 {
         let mut s = scratch();
-        return items.iter().enumerate().map(|(i, item)| f(&mut s, i, item)).collect();
+        return (0..len).map(|i| f(&mut s, i)).collect();
     }
     let next = AtomicUsize::new(0);
     let failed = AtomicBool::new(false);
     let work = || {
+        let _stop = StopOnPanic(&failed);
         let mut s = scratch();
         let mut made = Vec::new();
         while !failed.load(Ordering::Relaxed) {
             let from = next.fetch_add(BATCH, Ordering::Relaxed);
-            let Some(batch) = items.get(from..items.len().min(from + BATCH)).filter(|b| !b.is_empty()) else { break };
-            for (k, item) in batch.iter().enumerate() {
-                match f(&mut s, from + k, item) {
-                    Ok(r) => made.push((from + k, r)),
+            if from >= len {
+                break;
+            }
+            for i in from..len.min(from.saturating_add(BATCH)) {
+                match f(&mut s, i) {
+                    Ok(r) => made.push((i, r)),
                     Err(e) => {
                         failed.store(true, Ordering::Relaxed);
-                        return Err((from + k, e));
+                        return Err((i, e));
                     }
                 }
             }
@@ -82,7 +171,7 @@ pub(crate) fn try_map_on<T: Sync, S, R: Send, E: Send>(
             parts.push(h.join().unwrap_or_else(|e| std::panic::resume_unwind(e)));
         }
     });
-    let mut made: Vec<(usize, R)> = Vec::with_capacity(items.len());
+    let mut made: Vec<(usize, R)> = Vec::with_capacity(len);
     let mut first: Option<(usize, E)> = None;
     for part in parts {
         match part {
@@ -217,6 +306,80 @@ mod tests {
         assert!(ran.load(std::sync::atomic::Ordering::Relaxed) < items.len());
     }
 
+    #[test]
+    fn each_row_is_written_in_its_place() {
+        let alone: Vec<usize> = (0..37 * 5).map(|i| i / 5 * 3 + i % 5).collect();
+        for threads in [1, 2, 5, 16, 64] {
+            let mut out = vec![0usize; 37 * 5];
+            for_each_row(
+                out.chunks_mut(5),
+                threads,
+                || (),
+                |_, r, row: &mut [usize]| {
+                    row.iter_mut().enumerate().for_each(|(c, o)| *o = r * 3 + c);
+                },
+            );
+            assert_eq!(out, alone, "{threads} threads");
+        }
+        let mut none: Vec<u8> = Vec::new();
+        for_each_row(none.chunks_mut(5), 4, || (), |_, _, _: &mut [u8]| panic!("no rows"));
+        // Small work, and the browser: one thread; never more threads than rows.
+        assert_eq!(threads_for_rows(100, MIN_ROW_WORK - 1, 0), 1);
+        assert!(threads_for_rows(3, usize::MAX, 0) <= 3);
+        // Scratch larger than the budget: one thread more than it holds, none beyond.
+        assert!(threads_for_rows(100, usize::MAX, SCRATCH_BUDGET) <= 2);
+        assert_eq!(threads_for_rows(100, usize::MAX, usize::MAX), 1);
+    }
+
+    #[test]
+    fn a_workers_panic_stops_the_other_threads() {
+        // A worker panics on its first item; the others stop after the items they are making,
+        // far short of all of them.
+        let items: Vec<u32> = (0..2_000).collect();
+        let caller = std::thread::current().id();
+        let ran = std::sync::atomic::AtomicUsize::new(0);
+        let r = std::panic::catch_unwind(|| {
+            map_on(
+                &items,
+                4,
+                || (),
+                |_, _, x| {
+                    ran.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+                    if std::thread::current().id() != caller {
+                        panic!("a worker's panic");
+                    }
+                    std::thread::sleep(std::time::Duration::from_millis(1));
+                    *x
+                },
+            )
+        });
+        assert!(r.is_err(), "no worker ran an item");
+        assert!(ran.load(std::sync::atomic::Ordering::Relaxed) < items.len() / 2, "{}", ran.load(std::sync::atomic::Ordering::Relaxed));
+    }
+
+    #[test]
+    fn a_workers_panic_on_a_row_reaches_the_caller_and_stops_the_others() {
+        let caller = std::thread::current().id();
+        let ran = std::sync::atomic::AtomicUsize::new(0);
+        let mut out = vec![0u8; 2_000];
+        let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+            for_each_row(
+                out.chunks_mut(1),
+                4,
+                || (),
+                |_, _, _: &mut [u8]| {
+                    ran.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+                    if std::thread::current().id() != caller {
+                        panic!("a worker's panic");
+                    }
+                    std::thread::sleep(std::time::Duration::from_millis(1));
+                },
+            )
+        }));
+        assert!(r.is_err(), "no worker took a row");
+        assert!(ran.load(std::sync::atomic::Ordering::Relaxed) < 1_000, "{}", ran.load(std::sync::atomic::Ordering::Relaxed));
+    }
+
     #[test]
     fn no_more_threads_than_batches() {
         assert_eq!(threads_for(10, 256), 1);