/home/noah/src/trueno/src/brick/batch.rs
Line | Count | Source |
1 | | //! Batch Splitting and Work Distribution |
2 | | //! |
3 | | //! LCP-05: Balance211 Work Distribution (Intel MKL pattern) |
4 | | //! LCP-09: Batch Splitting Strategies |
5 | | |
6 | | // ============================================================================ |
7 | | // LCP-05: Balance211 Work Distribution (Intel MKL pattern) |
8 | | // ============================================================================ |
9 | | |
10 | | /// Balance211 work distribution (Intel MKL pattern). |
11 | | /// |
12 | | /// Distributes N items across T threads such that no thread |
13 | | /// has more than 1 extra item compared to any other. |
14 | | /// |
15 | | /// # Example |
16 | | /// ```rust |
17 | | /// use trueno::brick::balance211; |
18 | | /// |
19 | | /// let ranges = balance211(10, 3); |
20 | | /// // Thread 0: (0, 4) - 4 items |
21 | | /// // Thread 1: (4, 3) - 3 items |
22 | | /// // Thread 2: (7, 3) - 3 items |
23 | | /// assert_eq!(ranges.len(), 3); |
24 | | /// ``` |
25 | | #[must_use] |
26 | 0 | pub fn balance211(n: usize, nthreads: usize) -> Vec<(usize, usize)> { |
27 | 0 | if nthreads == 0 { |
28 | 0 | return vec![]; |
29 | 0 | } |
30 | 0 | let div = n / nthreads; |
31 | 0 | let rem = n % nthreads; |
32 | | |
33 | 0 | (0..nthreads) |
34 | 0 | .map(|i| { |
35 | 0 | let offset = if i < rem { |
36 | 0 | (div + 1) * i |
37 | | } else { |
38 | 0 | div * i + rem |
39 | | }; |
40 | 0 | let count = if i < rem { div + 1 } else { div }; |
41 | 0 | (offset, count) |
42 | 0 | }) |
43 | 0 | .collect() |
44 | 0 | } |
45 | | |
46 | | /// Iterator adapter for balanced work distribution. |
47 | | pub struct Balance211Iter { |
48 | | ranges: Vec<(usize, usize)>, |
49 | | current: usize, |
50 | | } |
51 | | |
52 | | impl Balance211Iter { |
53 | | /// Create a new balanced work iterator. |
54 | 0 | pub fn new(n: usize, nthreads: usize) -> Self { |
55 | 0 | Self { |
56 | 0 | ranges: balance211(n, nthreads), |
57 | 0 | current: 0, |
58 | 0 | } |
59 | 0 | } |
60 | | } |
61 | | |
62 | | impl Iterator for Balance211Iter { |
63 | | type Item = std::ops::Range<usize>; |
64 | | |
65 | 0 | fn next(&mut self) -> Option<Self::Item> { |
66 | 0 | if self.current >= self.ranges.len() { |
67 | 0 | return None; |
68 | 0 | } |
69 | 0 | let (offset, count) = self.ranges[self.current]; |
70 | 0 | self.current += 1; |
71 | 0 | Some(offset..offset + count) |
72 | 0 | } |
73 | | } |
74 | | |
75 | | impl ExactSizeIterator for Balance211Iter { |
76 | 0 | fn len(&self) -> usize { |
77 | 0 | self.ranges.len() - self.current |
78 | 0 | } |
79 | | } |
80 | | |
81 | | // ============================================================================ |
82 | | // LCP-09: Batch Splitting Strategies |
83 | | // ============================================================================ |
84 | | |
85 | | /// Strategy for splitting batches across workers. |
86 | | #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] |
87 | | pub enum BatchSplitStrategy { |
88 | | /// Simple equal division (may leave remainder) |
89 | | #[default] |
90 | | Simple, |
91 | | /// Equal distribution using Balance211 |
92 | | Equal, |
93 | | /// Sequence-aware (keeps sequences together) |
94 | | SequenceAware, |
95 | | } |
96 | | |
97 | | /// Split a batch into chunks according to strategy. |
98 | | /// |
99 | | /// # Example |
100 | | /// ```rust |
101 | | /// use trueno::brick::{split_batch, BatchSplitStrategy}; |
102 | | /// |
103 | | /// let chunks = split_batch(100, 4, BatchSplitStrategy::Equal); |
104 | | /// assert_eq!(chunks.len(), 4); |
105 | | /// assert_eq!(chunks.iter().sum::<usize>(), 100); |
106 | | /// ``` |
107 | | #[must_use] |
108 | 0 | pub fn split_batch(total: usize, num_workers: usize, strategy: BatchSplitStrategy) -> Vec<usize> { |
109 | 0 | if num_workers == 0 || total == 0 { |
110 | 0 | return vec![]; |
111 | 0 | } |
112 | | |
113 | 0 | match strategy { |
114 | | BatchSplitStrategy::Simple => { |
115 | 0 | let chunk_size = total / num_workers; |
116 | 0 | let mut chunks = vec![chunk_size; num_workers]; |
117 | | // Last worker gets remainder |
118 | 0 | if let Some(last) = chunks.last_mut() { |
119 | 0 | *last += total % num_workers; |
120 | 0 | } |
121 | 0 | chunks |
122 | | } |
123 | | BatchSplitStrategy::Equal => { |
124 | | // Use Balance211 for even distribution |
125 | 0 | balance211(total, num_workers) |
126 | 0 | .iter() |
127 | 0 | .map(|(_, count)| *count) |
128 | 0 | .collect() |
129 | | } |
130 | | BatchSplitStrategy::SequenceAware => { |
131 | | // For now, same as Equal (sequence boundaries would need external info) |
132 | 0 | balance211(total, num_workers) |
133 | 0 | .iter() |
134 | 0 | .map(|(_, count)| *count) |
135 | 0 | .collect() |
136 | | } |
137 | | } |
138 | 0 | } |
139 | | |
140 | | #[cfg(test)] |
141 | | mod tests { |
142 | | use super::*; |
143 | | |
144 | | #[test] |
145 | | fn test_balance211_basic() { |
146 | | let ranges = balance211(10, 3); |
147 | | assert_eq!(ranges.len(), 3); |
148 | | // First thread gets 4 items, others get 3 |
149 | | assert_eq!(ranges[0], (0, 4)); |
150 | | assert_eq!(ranges[1], (4, 3)); |
151 | | assert_eq!(ranges[2], (7, 3)); |
152 | | } |
153 | | |
154 | | #[test] |
155 | | fn test_balance211_even_division() { |
156 | | let ranges = balance211(12, 4); |
157 | | // 12 / 4 = 3 items each, no remainder |
158 | | for (i, &(offset, count)) in ranges.iter().enumerate() { |
159 | | assert_eq!(count, 3); |
160 | | assert_eq!(offset, i * 3); |
161 | | } |
162 | | } |
163 | | |
164 | | #[test] |
165 | | fn test_balance211_empty() { |
166 | | assert!(balance211(0, 4).iter().all(|&(_, c)| c == 0)); |
167 | | assert!(balance211(10, 0).is_empty()); |
168 | | } |
169 | | |
170 | | #[test] |
171 | | fn test_balance211_single_thread() { |
172 | | let ranges = balance211(100, 1); |
173 | | assert_eq!(ranges.len(), 1); |
174 | | assert_eq!(ranges[0], (0, 100)); |
175 | | } |
176 | | |
177 | | #[test] |
178 | | fn test_balance211_more_threads_than_items() { |
179 | | let ranges = balance211(3, 5); |
180 | | assert_eq!(ranges.len(), 5); |
181 | | // First 3 threads get 1 item each, last 2 get 0 |
182 | | let items: Vec<_> = ranges.iter().map(|(_, c)| *c).collect(); |
183 | | assert_eq!(items, vec![1, 1, 1, 0, 0]); |
184 | | } |
185 | | |
186 | | #[test] |
187 | | fn test_balance211_iter_basic() { |
188 | | let mut iter = Balance211Iter::new(10, 3); |
189 | | assert_eq!(iter.len(), 3); |
190 | | |
191 | | assert_eq!(iter.next(), Some(0..4)); |
192 | | assert_eq!(iter.next(), Some(4..7)); |
193 | | assert_eq!(iter.next(), Some(7..10)); |
194 | | assert_eq!(iter.next(), None); |
195 | | } |
196 | | |
197 | | #[test] |
198 | | fn test_balance211_iter_exact_size() { |
199 | | let iter = Balance211Iter::new(10, 3); |
200 | | assert_eq!(iter.len(), 3); |
201 | | |
202 | | let mut iter2 = Balance211Iter::new(10, 3); |
203 | | iter2.next(); |
204 | | assert_eq!(iter2.len(), 2); |
205 | | } |
206 | | |
207 | | #[test] |
208 | | fn test_batch_split_strategy_default() { |
209 | | assert_eq!(BatchSplitStrategy::default(), BatchSplitStrategy::Simple); |
210 | | } |
211 | | |
212 | | #[test] |
213 | | fn test_split_batch_simple() { |
214 | | let chunks = split_batch(100, 4, BatchSplitStrategy::Simple); |
215 | | assert_eq!(chunks.len(), 4); |
216 | | // First 3 get 25, last gets 25 + 0 = 25 |
217 | | assert_eq!(chunks, vec![25, 25, 25, 25]); |
218 | | } |
219 | | |
220 | | #[test] |
221 | | fn test_split_batch_simple_with_remainder() { |
222 | | let chunks = split_batch(10, 3, BatchSplitStrategy::Simple); |
223 | | assert_eq!(chunks.len(), 3); |
224 | | // 10 / 3 = 3, remainder = 1 goes to last |
225 | | assert_eq!(chunks, vec![3, 3, 4]); |
226 | | assert_eq!(chunks.iter().sum::<usize>(), 10); |
227 | | } |
228 | | |
229 | | #[test] |
230 | | fn test_split_batch_equal() { |
231 | | let chunks = split_batch(10, 3, BatchSplitStrategy::Equal); |
232 | | assert_eq!(chunks.len(), 3); |
233 | | // Balance211 gives first worker more: 4, 3, 3 |
234 | | assert_eq!(chunks, vec![4, 3, 3]); |
235 | | assert_eq!(chunks.iter().sum::<usize>(), 10); |
236 | | } |
237 | | |
238 | | #[test] |
239 | | fn test_split_batch_sequence_aware() { |
240 | | let chunks = split_batch(10, 3, BatchSplitStrategy::SequenceAware); |
241 | | // Currently same as Equal |
242 | | assert_eq!(chunks, vec![4, 3, 3]); |
243 | | } |
244 | | |
245 | | #[test] |
246 | | fn test_split_batch_empty() { |
247 | | assert!(split_batch(0, 4, BatchSplitStrategy::Simple).is_empty()); |
248 | | assert!(split_batch(100, 0, BatchSplitStrategy::Simple).is_empty()); |
249 | | } |
250 | | |
251 | | #[test] |
252 | | fn test_split_batch_single_worker() { |
253 | | let chunks = split_batch(100, 1, BatchSplitStrategy::Simple); |
254 | | assert_eq!(chunks, vec![100]); |
255 | | } |
256 | | |
257 | | /// FALSIFICATION TEST: Verify total items preserved after split |
258 | | /// |
259 | | /// The sum of all chunks must equal the original total. |
260 | | #[test] |
261 | | fn test_falsify_split_batch_preserves_total() { |
262 | | for total in [1, 10, 100, 997, 1000, 10000] { |
263 | | for workers in [1, 2, 3, 4, 7, 16, 100] { |
264 | | for strategy in [ |
265 | | BatchSplitStrategy::Simple, |
266 | | BatchSplitStrategy::Equal, |
267 | | BatchSplitStrategy::SequenceAware, |
268 | | ] { |
269 | | let chunks = split_batch(total, workers, strategy); |
270 | | let sum: usize = chunks.iter().sum(); |
271 | | assert_eq!( |
272 | | sum, total, |
273 | | "FALSIFICATION FAILED: split_batch({}, {}, {:?}) sum {} != {}", |
274 | | total, workers, strategy, sum, total |
275 | | ); |
276 | | } |
277 | | } |
278 | | } |
279 | | } |
280 | | |
281 | | /// FALSIFICATION TEST: Balance211 never gives more than 1 extra item |
282 | | /// |
283 | | /// The maximum difference between any two thread counts must be <= 1. |
284 | | #[test] |
285 | | fn test_falsify_balance211_max_diff_one() { |
286 | | for n in [1, 10, 100, 997, 1000] { |
287 | | for nthreads in [1, 2, 3, 4, 7, 16, 100] { |
288 | | let ranges = balance211(n, nthreads); |
289 | | if ranges.is_empty() { |
290 | | continue; |
291 | | } |
292 | | let counts: Vec<_> = ranges.iter().map(|(_, c)| *c).collect(); |
293 | | let max_count = *counts.iter().max().unwrap_or(&0); |
294 | | let min_count = *counts.iter().min().unwrap_or(&0); |
295 | | assert!( |
296 | | max_count - min_count <= 1, |
297 | | "FALSIFICATION FAILED: balance211({}, {}) has diff {} (max={}, min={})", |
298 | | n, |
299 | | nthreads, |
300 | | max_count - min_count, |
301 | | max_count, |
302 | | min_count |
303 | | ); |
304 | | } |
305 | | } |
306 | | } |
307 | | |
308 | | /// FALSIFICATION TEST: Balance211 ranges are contiguous and non-overlapping |
309 | | #[test] |
310 | | fn test_falsify_balance211_contiguous() { |
311 | | for n in [10, 100, 1000] { |
312 | | for nthreads in [2, 3, 4, 7] { |
313 | | let ranges = balance211(n, nthreads); |
314 | | let mut expected_offset = 0; |
315 | | for (i, &(offset, count)) in ranges.iter().enumerate() { |
316 | | assert_eq!( |
317 | | offset, expected_offset, |
318 | | "FALSIFICATION FAILED: balance211({}, {}) range {} offset {} != expected {}", |
319 | | n, nthreads, i, offset, expected_offset |
320 | | ); |
321 | | expected_offset += count; |
322 | | } |
323 | | assert_eq!( |
324 | | expected_offset, n, |
325 | | "FALSIFICATION FAILED: balance211({}, {}) total {} != {}", |
326 | | n, nthreads, expected_offset, n |
327 | | ); |
328 | | } |
329 | | } |
330 | | } |
331 | | } |