grafos_std/cpu_shared/
partition.rs

1//! Range/tile partition helpers for shared-memory tasklets.
2//!
3//! Pure integer math — available on every target, both host and wasm32,
4//! with or without the `std` feature.
5
6/// Divide `0..total` into `worker_count` contiguous chunks and return
7/// the chunk owned by `worker_index`. Distribution handles remainders
8/// by giving the first `total % worker_count` workers one extra item.
9///
10/// This is the right choice for most workloads: the union of all workers
11/// covers `0..total` exactly, and chunk sizes differ by at most one. For
12/// algorithms that require every worker to process the exact same number
13/// of items (and want to fail loudly otherwise), see [`range_even`].
14pub fn range(total: usize, worker_index: u16, worker_count: u16) -> core::ops::Range<usize> {
15    assert!(worker_count > 0, "worker_count must be > 0");
16    assert!(worker_index < worker_count, "worker_index out of range");
17    let wc = worker_count as usize;
18    let wi = worker_index as usize;
19    let base = total / wc;
20    let rem = total % wc;
21    let start = wi * base + wi.min(rem);
22    let len = base + if wi < rem { 1 } else { 0 };
23    start..start + len
24}
25
26/// Error returned by [`range_even`] when `total` is not evenly divisible
27/// by `worker_count`.
28///
29/// Kept deliberately small and `Copy` so it can be propagated through
30/// `no_std` guest code without any allocation.
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub struct UnevenError {
33    pub total: usize,
34    pub worker_count: u16,
35}
36
37impl core::fmt::Display for UnevenError {
38    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
39        write!(
40            f,
41            "partition::range_even: total {} is not evenly divisible by worker_count {}",
42            self.total, self.worker_count
43        )
44    }
45}
46
47/// Strict even-split variant of [`range`].
48///
49/// Returns `Ok(slice)` — identical to `range(total, worker_index, worker_count)`
50/// — when `total % worker_count == 0`, so every worker owns exactly the same
51/// number of items. Returns [`UnevenError`] otherwise.
52///
53/// Use this when your algorithm genuinely requires an even split (e.g. SIMD
54/// kernels that assume a fixed per-worker tile size) and you would rather
55/// fail loudly than silently process an uneven distribution. For workloads
56/// that can absorb a one-item skew — histograms, reductions, map-style
57/// passes — prefer [`range`].
58pub fn range_even(
59    total: usize,
60    worker_index: u16,
61    worker_count: u16,
62) -> Result<core::ops::Range<usize>, UnevenError> {
63    assert!(worker_count > 0, "worker_count must be > 0");
64    assert!(worker_index < worker_count, "worker_index out of range");
65    if total % (worker_count as usize) != 0 {
66        return Err(UnevenError {
67            total,
68            worker_count,
69        });
70    }
71    Ok(range(total, worker_index, worker_count))
72}
73
74/// A 2D image tile.
75#[derive(Clone, Copy, Debug, PartialEq, Eq)]
76pub struct Tile {
77    pub x: u32,
78    pub y: u32,
79    pub w: u32,
80    pub h: u32,
81}
82
83/// Iterate the tiles assigned to `worker_index`. Tiles are produced in
84/// row-major order across the `(width, height)` image with cell size
85/// `(tile_w, tile_h)`. Worker `i` receives every `worker_count`-th tile
86/// starting at index `i`, so the union of all workers covers the image
87/// exactly once.
88pub fn tiles(
89    width: u32,
90    height: u32,
91    tile_w: u32,
92    tile_h: u32,
93    worker_index: u16,
94    worker_count: u16,
95) -> impl Iterator<Item = Tile> {
96    assert!(tile_w > 0 && tile_h > 0);
97    assert!(worker_count > 0 && worker_index < worker_count);
98    let cols = width.div_ceil(tile_w);
99    let rows = height.div_ceil(tile_h);
100    let total = (cols as usize) * (rows as usize);
101    let wi = worker_index as usize;
102    let wc = worker_count as usize;
103    (0..total).filter(move |i| i % wc == wi).map(move |i| {
104        let cx = (i as u32) % cols;
105        let cy = (i as u32) / cols;
106        let x = cx * tile_w;
107        let y = cy * tile_h;
108        let w = (tile_w).min(width - x);
109        let h = (tile_h).min(height - y);
110        Tile { x, y, w, h }
111    })
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn range_even_succeeds_when_total_divides_evenly() {
120        assert_eq!(range_even(64, 0, 4), Ok(0..16));
121        assert_eq!(range_even(64, 1, 4), Ok(16..32));
122        assert_eq!(range_even(64, 2, 4), Ok(32..48));
123        assert_eq!(range_even(64, 3, 4), Ok(48..64));
124    }
125
126    #[test]
127    fn range_even_returns_error_when_total_does_not_divide_evenly() {
128        // 64 items across 3 workers is the canonical histogram footgun.
129        assert_eq!(
130            range_even(64, 0, 3),
131            Err(UnevenError {
132                total: 64,
133                worker_count: 3,
134            })
135        );
136        assert_eq!(
137            range_even(10, 2, 4),
138            Err(UnevenError {
139                total: 10,
140                worker_count: 4,
141            })
142        );
143    }
144
145    #[test]
146    fn range_even_handles_corner_cases() {
147        // Zero items divides evenly by any worker_count.
148        assert_eq!(range_even(0, 0, 1), Ok(0..0));
149        assert_eq!(range_even(0, 1, 4), Ok(0..0));
150        // Single item, single worker.
151        assert_eq!(range_even(1, 0, 1), Ok(0..1));
152        // Single item across two workers does NOT divide evenly.
153        assert_eq!(
154            range_even(1, 0, 2),
155            Err(UnevenError {
156                total: 1,
157                worker_count: 2,
158            })
159        );
160    }
161}
grafos_std/cpu_shared/partition.rs

grafos_std/cpu_shared/
partition.rs
