tcontract_exec/batch.rs
1//! Batch lane policy: the shared vocabulary and the one decision function.
2//!
3//! Three libraries decide how a batch of independent problems is split over a
4//! thread budget (tcontract's packed driver and faer strategy, cpueinsum's
5//! grouped GEMM, tlinalg's lane plan). This module gives them one vocabulary
6//! and one pure function, [`batch_decision`], so a caller can force a route and
7//! two crates cannot drift on the same decision. The types are plain data with
8//! no scheduler dependency; they live apart from [`Exec`](crate::Exec) so moving
9//! them later is a re-export, not a redesign.
10//!
11//! # Strategies
12//!
13//! [`BatchStrategy::Auto`] is the only heuristic. Every other variant forces a
14//! route and never falls back: the route either honors it (clamped to the
15//! items and the budget) or the library reports
16//! [`crate::ExecError::StrategyUnavailable`].
17//!
18//! Forced strategies are evaluated against the **original host budget**. Only
19//! `Auto` narrows, because the packed driver narrows its `Exec` twice (see
20//! below).
21//!
22//! # What `Auto` computes
23//!
24//! `Auto` has two forms, chosen by the query. This is how one formula serves
25//! both the shared work model and the packed driver's legacy rule:
26//!
27//! * Without `balanced_lanes`, the issue's generic formula: estimate
28//! `item_ns` and `total_ns`, cap the lanes at `total_ns /
29//! min_work_per_lane_ns`, and fan out only when `lanes >= 2`, `items >=
30//! min_items` and `items >= lanes * min_items_per_lane`. Each lane runs
31//! items serially; with one lane, `item_threads` is the width of one item, so
32//! a route that permits it (faer) can parallelize inside the item.
33//! * With `balanced_lanes`, the packed two-stage rule below, which the default
34//! [`WorkModel`] reproduces to the lane.
35//!
36//! # The legacy rules and how `Auto` expresses them
37//!
38//! This step lands the vocabulary with defaults that reproduce today's
39//! tcontract packed decision; the sweep that fits one shared default to all
40//! three libraries is a later step. The three legacy rules and their mapping:
41//!
42//! | Legacy rule | Source | Encoding here |
43//! |---|---|---|
44//! | packed balanced lanes | `crates/tcontract/src/driver/batch.rs:38`, `plan/mod.rs:641`, `driver/route.rs:49` | `Auto` + `balanced_lanes` + the default [`WorkModel`] (its `width` is the existing [`WidthPolicy`], its item cost is `NS_PER_FLOP`) |
45//! | faer item-width-1 | `crates/tcontract/src/strategy/faer.rs:243` | `Auto` + `item_parallel`; a decision with `item_threads > 1` is the inner route, `lanes > 1` with item threads one is the outer route |
46//! | `contract_batched` item count | `crates/tcontract/src/batch.rs:140` | forced strategies, not `Auto`: [`BatchStrategy::Outer`] when the items fill the budget; below that one lane, the caller running each item on the whole budget (`Sequential`/`Inner`). It has no work threshold. |
47//!
48//! The packed rule is two-stage. `plan/mod.rs:641` first narrows the `Exec` to
49//! `width_for(total_ns)` over the whole batch, and `driver/batch.rs:38`
50//! recomputes the width against that narrowed budget and only then applies its
51//! *tiny* (one item already below the serial threshold) and *balanced* (every
52//! lane non-empty, the split near-even, the batch alone fills the budget)
53//! tests. `Auto` reproduces exactly that: it narrows `budget` with
54//! `model.width`, then applies the tiny/balanced tests when the query says
55//! `balanced_lanes`, and reports the narrowed width as `item_threads` when it
56//! keeps one lane. The test `auto_matches_legacy_packed_table` pins this on a
57//! table.
58//!
59//! # Examples
60//!
61//! ```
62//! use std::num::NonZeroUsize;
63//! use tcontract_exec::batch::{batch_decision, BatchPolicy, BatchQuery, BatchStrategy};
64//!
65//! let budget = NonZeroUsize::new(8).unwrap();
66//! let query = BatchQuery::new(16, 64.0);
67//!
68//! // A forced route is clamped, never a fallback.
69//! let outer = BatchPolicy::new(BatchStrategy::Outer);
70//! let d = batch_decision(&outer, budget, &query);
71//! assert_eq!((d.lanes, d.item_threads), (8, 1));
72//! ```
73
74use core::num::NonZeroUsize;
75
76use crate::WidthPolicy;
77
78/// The default policy [`BatchStrategy::Auto`] applies when no explicit policy is
79/// set on an [`Exec`](crate::Exec).
80///
81/// It is the literal form of [`BatchPolicy::default`], so a `Serial` context can
82/// hand out a `'static` policy without carrying one. A unit test pins the two
83/// together.
84pub(crate) const DEFAULT_BATCH_POLICY: BatchPolicy = BatchPolicy {
85 strategy: BatchStrategy::Auto,
86 model: WorkModel {
87 // Reproduce `driver/batch.rs`: `item_ns = 2 * m * n * k * NS_PER_FLOP`.
88 item_overhead_ns: 0.0,
89 muladds_per_ns: 10.0,
90 min_work_per_lane_ns: 0.0,
91 min_items: 2,
92 min_items_per_lane: 1,
93 width: WidthPolicy {
94 serial_below_ns: 50_000.0,
95 entry_base_ns: 5_000.0,
96 entry_per_thread_ns: 5_000.0,
97 },
98 },
99};
100
101/// How the items of a batch are split over a thread budget.
102///
103/// # Examples
104///
105/// ```
106/// use tcontract_exec::batch::BatchStrategy;
107/// assert_eq!(BatchStrategy::default(), BatchStrategy::Auto);
108/// ```
109#[non_exhaustive]
110#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
111pub enum BatchStrategy {
112 /// Choose from the [`WorkModel`] and the query's route capabilities.
113 #[default]
114 Auto,
115 /// One lane, one item thread: no parallelism at all.
116 Sequential,
117 /// `lanes = min(budget, items)`, each item serial.
118 Outer,
119 /// One lane; the item kernel gets the whole budget.
120 Inner,
121 /// Exactly `lanes` lanes (clamped to the items and the budget), the rest
122 /// of the budget split evenly as item threads. A route that cannot run item
123 /// threads inside a lane refuses the resulting decision with
124 /// [`crate::ExecError::StrategyUnavailable`] before any write; it never silently
125 /// runs the items serially (see [`BatchDecision`]'s route contract).
126 Split {
127 /// Requested lane count.
128 lanes: NonZeroUsize,
129 },
130}
131
132/// The cost and fan-out model [`BatchStrategy::Auto`] applies.
133///
134/// The defaults are the values tcontract used before this vocabulary existed,
135/// not a tuned result: they reproduce the packed decision together with the
136/// existing [`WidthPolicy`]. The shared, swept defaults are a later step.
137///
138/// The three item-cost fields follow tenferro's restored
139/// `CpuBatchThresholds`: one item costs roughly `item_overhead_ns` plus one
140/// nanosecond per `muladds_per_ns` multiply-adds, a lane needs
141/// `min_work_per_lane_ns` of estimated work, and a fan-out needs `min_items`
142/// items, at least `min_items_per_lane` per lane.
143///
144/// # Examples
145///
146/// ```
147/// use tcontract_exec::batch::WorkModel;
148///
149/// let model = WorkModel::default().with_min_items(4).with_muladds_per_ns(32.0);
150/// assert_eq!(model.min_items, 4);
151/// assert_eq!(model.muladds_per_ns, 32.0);
152/// ```
153#[non_exhaustive]
154#[derive(Clone, Copy, Debug, PartialEq)]
155pub struct WorkModel {
156 /// Estimated fixed cost of one item, in nanoseconds.
157 ///
158 /// The default is zero so that [`Default`] reproduces the packed
159 /// `NS_PER_FLOP` estimate exactly.
160 pub item_overhead_ns: f64,
161 /// Estimated multiply-adds per nanosecond on one lane.
162 ///
163 /// `1.0 / NS_PER_FLOP / 2.0 == 10.0` reproduces
164 /// `2 * macs * NS_PER_FLOP`: the packed driver counts two flops per
165 /// multiply-add. A value `<= 0` is treated as one multiply-add per
166 /// nanosecond.
167 pub muladds_per_ns: f64,
168 /// Least estimated work a lane must receive, in nanoseconds; `0` disables
169 /// the work gate (the packed default).
170 pub min_work_per_lane_ns: f64,
171 /// Fewest items for which `Auto` fans out at all.
172 pub min_items: usize,
173 /// Fewest items each lane must receive.
174 pub min_items_per_lane: usize,
175 /// Team-width cost model, used to narrow the budget over the whole batch.
176 pub width: WidthPolicy,
177}
178
179impl Default for WorkModel {
180 fn default() -> Self {
181 Self {
182 // Reproduce `driver/batch.rs`: `item_ns = 2 * m * n * k * NS_PER_FLOP`.
183 item_overhead_ns: 0.0,
184 muladds_per_ns: 10.0,
185 min_work_per_lane_ns: 0.0,
186 min_items: 2,
187 min_items_per_lane: 1,
188 width: WidthPolicy::default(),
189 }
190 }
191}
192
193impl WorkModel {
194 /// Return this model with a different per-item overhead.
195 #[must_use]
196 pub fn with_item_overhead_ns(mut self, nanoseconds: f64) -> Self {
197 self.item_overhead_ns = nanoseconds;
198 self
199 }
200
201 /// Return this model with a different lane throughput.
202 ///
203 /// # Examples
204 ///
205 /// ```
206 /// use tcontract_exec::batch::WorkModel;
207 /// assert_eq!(WorkModel::default().with_muladds_per_ns(16.0).muladds_per_ns, 16.0);
208 /// ```
209 #[must_use]
210 pub fn with_muladds_per_ns(mut self, muladds_per_ns: f64) -> Self {
211 self.muladds_per_ns = muladds_per_ns;
212 self
213 }
214
215 /// Return this model with a different minimum work per lane.
216 ///
217 /// # Examples
218 ///
219 /// ```
220 /// use tcontract_exec::batch::WorkModel;
221 /// assert_eq!(WorkModel::default().with_min_work_per_lane_ns(8_000.0).min_work_per_lane_ns, 8_000.0);
222 /// ```
223 #[must_use]
224 pub fn with_min_work_per_lane_ns(mut self, nanoseconds: f64) -> Self {
225 self.min_work_per_lane_ns = nanoseconds;
226 self
227 }
228
229 /// Return this model with a different minimum fan-out item count.
230 #[must_use]
231 pub fn with_min_items(mut self, items: usize) -> Self {
232 self.min_items = items;
233 self
234 }
235
236 /// Return this model with a different minimum item count per lane.
237 #[must_use]
238 pub fn with_min_items_per_lane(mut self, items: usize) -> Self {
239 self.min_items_per_lane = items;
240 self
241 }
242
243 /// Return this model with a different team-width cost model.
244 #[must_use]
245 pub fn with_width(mut self, width: WidthPolicy) -> Self {
246 self.width = width;
247 self
248 }
249}
250
251/// A strategy and the model its `Auto` arm applies.
252///
253/// # Examples
254///
255/// ```
256/// use tcontract_exec::batch::{BatchPolicy, BatchStrategy};
257/// assert_eq!(BatchPolicy::default().strategy, BatchStrategy::Auto);
258/// ```
259#[derive(Clone, Copy, Debug, Default, PartialEq)]
260pub struct BatchPolicy {
261 /// The route selector.
262 pub strategy: BatchStrategy,
263 /// The cost model `Auto` applies.
264 pub model: WorkModel,
265}
266
267impl BatchPolicy {
268 /// A policy with `strategy` and the default model.
269 #[must_use]
270 pub fn new(strategy: BatchStrategy) -> Self {
271 Self {
272 strategy,
273 model: WorkModel::default(),
274 }
275 }
276
277 /// Return this policy with a different strategy.
278 #[must_use]
279 pub fn with_strategy(mut self, strategy: BatchStrategy) -> Self {
280 self.strategy = strategy;
281 self
282 }
283
284 /// Return this policy with a different model.
285 #[must_use]
286 pub fn with_model(mut self, model: WorkModel) -> Self {
287 self.model = model;
288 self
289 }
290}
291
292/// What one batch of one route looks like to [`batch_decision`].
293///
294/// Work is described by an estimate, not by dimensions: each family supplies
295/// the multiply-adds of **one** item (GEMM `mnk`, LU `n^3/3`, solve `n^2 k`,
296/// QR `2mn^2/3`, with complex elements weighing four). The booleans and
297/// `balanced_lanes` say what the route permits; they are how the legacy rules
298/// above become inputs instead of hidden constants.
299///
300/// # Examples
301///
302/// ```
303/// use tcontract_exec::batch::BatchQuery;
304///
305/// let query = BatchQuery::new(32, 4096.0).with_item_parallel(false);
306/// assert_eq!(query.items, 32);
307/// assert!(!query.item_parallel);
308/// ```
309#[non_exhaustive]
310#[derive(Clone, Copy, Debug, PartialEq)]
311pub struct BatchQuery {
312 /// Number of independent items; `0` means there is no work.
313 pub items: usize,
314 /// Multiply-adds of one item.
315 pub item_muladds: f64,
316 /// The route permits item-parallel width greater than one, i.e.
317 /// `item_threads > 1` is meaningful. `Auto` keeps item threads one when
318 /// this is false.
319 pub item_parallel: bool,
320 /// `Auto` may only keep lanes when every lane is non-empty, the split is
321 /// near-even and the batch alone fills the narrowed budget: the packed
322 /// driver's balanced-lanes rule.
323 pub balanced_lanes: bool,
324 /// `Auto` prefers the inner route (one lane, the item kernel gets a team) and
325 /// fans out over items only when one item alone cannot use a team: the faer
326 /// batched-GEMM rule. The item width is then the width of one item, not the
327 /// whole budget.
328 pub outer_only_when_item_serial: bool,
329}
330
331impl BatchQuery {
332 /// A query for `items` items of `item_muladds` multiply-adds each, on a
333 /// route that permits item parallelism and has no vendor batch route.
334 #[must_use]
335 pub fn new(items: usize, item_muladds: f64) -> Self {
336 Self {
337 items,
338 item_muladds,
339 item_parallel: true,
340 balanced_lanes: false,
341 outer_only_when_item_serial: false,
342 }
343 }
344
345 /// Return this query with the item-parallel capability set.
346 #[must_use]
347 pub fn with_item_parallel(mut self, item_parallel: bool) -> Self {
348 self.item_parallel = item_parallel;
349 self
350 }
351
352 /// Return this query with the balanced-lanes capability set.
353 #[must_use]
354 pub fn with_balanced_lanes(mut self, balanced_lanes: bool) -> Self {
355 self.balanced_lanes = balanced_lanes;
356 self
357 }
358
359 /// Return this query with the "outer only when one item cannot use a team"
360 /// capability set (the faer batched-GEMM rule).
361 #[must_use]
362 pub fn with_outer_only_when_item_serial(mut self, value: bool) -> Self {
363 self.outer_only_when_item_serial = value;
364 self
365 }
366}
367
368/// The route [`batch_decision`] selected.
369///
370/// # The route contract
371///
372/// A decision is the **whole** concurrency of one call, and a route must never
373/// discard half of it: `lanes` outer lanes run a contiguous share of the items,
374/// each lane may run its items on `item_threads` threads, and `batch_decision`
375/// keeps `lanes * item_threads <= budget`. A route that cannot run an item on a
376/// sub-team inside a lane either honours both axes or refuses a decision that
377/// asks for both — only a forced [`BatchStrategy::Split`] produces `lanes > 1`
378/// together with `item_threads > 1` — with
379/// [`crate::ExecError::StrategyUnavailable`], before it writes any output. Silently
380/// running the items serially inside the lanes is not a route.
381///
382/// The routes in tcontract, and where each one stands:
383///
384/// * the faer strategy honours both axes: every item runs with
385/// `item_threads` slots as the width it hands faer through `Exec::external`
386/// inside its lane's share of the token, which nests there (`crates/tcontract/src/strategy/faer.rs`);
387/// * the packed driver refuses it
388/// (`crates/tcontract/src/driver/route.rs`): its inner parallelism is a
389/// co-scheduled, barrier-bearing team, and a lane worker cannot join one
390/// without blocking on an outer barrier;
391/// * `contract_batched` refuses it too (`crates/tcontract/src/batch.rs`): its
392/// schedule is one item per lane, or the items one after another on the whole
393/// context, as the decision log records.
394///
395/// # Examples
396///
397/// ```
398/// use tcontract_exec::batch::BatchDecision;
399/// let d = BatchDecision { lanes: 2, item_threads: 1 };
400/// assert_eq!(d.lanes * d.item_threads, 2);
401/// ```
402#[derive(Clone, Copy, Debug, PartialEq, Eq)]
403pub struct BatchDecision {
404 /// Outer lanes; each lane runs a contiguous share of the items.
405 pub lanes: usize,
406 /// Threads one item may use. `1` for a route that runs items serially inside
407 /// each lane; a [`BatchStrategy::Split`] asks for both lanes and item threads,
408 /// and a route that cannot run item threads inside lanes refuses it (see the
409 /// type's route contract).
410 pub item_threads: usize,
411}
412
413/// Decide the route for one batch on `budget` threads under `policy`.
414///
415/// Forced strategies are clamped (`lanes <= min(items, budget)`,
416/// `lanes * item_threads <= budget`) but never fall back. `Auto` narrows the
417/// budget first (see the module docs) and uses the original budget only for the
418/// forced strategies.
419///
420/// Whether a *route* can run the returned decision is the route's own
421/// preflight: a decision with both lanes and item threads is honoured or
422/// refused before any output write, never partially applied (see
423/// [`BatchDecision`]'s route contract).
424///
425/// # Examples
426///
427/// ```
428/// use std::num::NonZeroUsize;
429/// use tcontract_exec::batch::{batch_decision, BatchPolicy, BatchQuery, BatchStrategy};
430///
431/// let budget = NonZeroUsize::new(4).unwrap();
432/// let query = BatchQuery::new(2, 1024.0);
433/// let split = BatchPolicy::new(BatchStrategy::Split {
434/// lanes: NonZeroUsize::new(8).unwrap(),
435/// });
436/// let d = batch_decision(&split, budget, &query);
437/// // Clamped to the items and the budget.
438/// assert_eq!((d.lanes, d.item_threads), (2, 2));
439/// ```
440pub fn batch_decision(
441 policy: &BatchPolicy,
442 budget: NonZeroUsize,
443 query: &BatchQuery,
444) -> BatchDecision {
445 let budget = budget.get();
446 match policy.strategy {
447 BatchStrategy::Auto => auto_decision(&policy.model, budget, query),
448 BatchStrategy::Sequential => BatchDecision {
449 lanes: 1,
450 item_threads: 1,
451 },
452 BatchStrategy::Outer => BatchDecision {
453 lanes: clamp_lanes(budget, query.items),
454 item_threads: 1,
455 },
456 BatchStrategy::Inner => BatchDecision {
457 lanes: 1,
458 item_threads: budget.max(1),
459 },
460 BatchStrategy::Split { lanes } => {
461 let lanes = clamp_lanes(budget, query.items).min(lanes.get());
462 BatchDecision {
463 lanes,
464 item_threads: budget / lanes,
465 }
466 }
467 }
468}
469
470/// `lanes = min(budget, items)`, never below one.
471fn clamp_lanes(budget: usize, items: usize) -> usize {
472 items.clamp(1, budget.max(1))
473}
474
475/// The `Auto` arm.
476///
477/// Without `balanced_lanes` this is the issue's generic formula: estimate the
478/// work, cap the lanes by the work model, and fan out when the `min_items` and
479/// `min_items_per_lane` gates pass. With `balanced_lanes` it is the packed
480/// driver's two-stage narrowing, which is what the default [`WorkModel`]
481/// reproduces to the lane.
482fn auto_decision(model: &WorkModel, budget: usize, query: &BatchQuery) -> BatchDecision {
483 let items = query.items;
484 if items == 0 || budget < 2 {
485 return one_lane(1);
486 }
487 let item_ns = item_ns(model, query.item_muladds);
488 let total_ns = item_ns * items as f64;
489 if query.outer_only_when_item_serial {
490 faer_auto(model, budget, query, item_ns, total_ns)
491 } else if query.balanced_lanes {
492 packed_auto(model, budget, query, item_ns, total_ns)
493 } else {
494 generic_auto(model, budget, query, item_ns, total_ns)
495 }
496}
497
498/// The faer batched-GEMM rule: prefer the inner route, and fan out over items
499/// only when one item alone is too small for a team (`width_for(item_ns) == 1`)
500/// and there are at least two items with work enough to leave the caller.
501///
502/// This is the item-width-1 rule of `strategy/faer.rs`: a large item gets the
503/// whole team, a tiny one gets lanes so the team would not pay barriers for
504/// nothing.
505fn faer_auto(
506 model: &WorkModel,
507 budget: usize,
508 query: &BatchQuery,
509 item_ns: f64,
510 total_ns: f64,
511) -> BatchDecision {
512 let inner = width_for(item_ns, &model.width, budget);
513 if inner > 1 {
514 return one_lane(inner);
515 }
516 let total = width_for(total_ns, &model.width, budget);
517 if total > 1 && query.items > 1 {
518 let lanes = budget.min(query.items).min(total);
519 return BatchDecision {
520 lanes,
521 item_threads: 1,
522 };
523 }
524 one_lane(1)
525}
526
527/// The issue's generic `auto_lanes`: lanes from the work model alone.
528fn generic_auto(
529 model: &WorkModel,
530 budget: usize,
531 query: &BatchQuery,
532 item_ns: f64,
533 total_ns: f64,
534) -> BatchDecision {
535 let items = query.items;
536 let by_work = work_lane_cap(model, total_ns);
537 let lanes = budget.min(items).min(by_work);
538 if lanes >= 2
539 && items >= model.min_items
540 && items >= lanes.saturating_mul(model.min_items_per_lane)
541 {
542 return BatchDecision {
543 lanes,
544 item_threads: 1,
545 };
546 }
547 // One lane: the inner route is the width of one item, when the route
548 // permits item parallelism (the faer rule).
549 let threads = if query.item_parallel {
550 width_for(item_ns, &model.width, budget)
551 } else {
552 1
553 };
554 one_lane(threads)
555}
556
557/// The packed driver's two-stage narrowing: `plan/mod.rs:641` narrows the
558/// budget over the whole batch, `driver/batch.rs:38` then applies the
559/// tiny/balanced rule against the narrowed budget.
560fn packed_auto(
561 model: &WorkModel,
562 budget: usize,
563 query: &BatchQuery,
564 item_ns: f64,
565 total_ns: f64,
566) -> BatchDecision {
567 let items = query.items;
568 let narrowed = width_for(total_ns, &model.width, budget);
569 let one_lane_threads = if query.item_parallel { narrowed } else { 1 };
570 if narrowed < 2 {
571 return one_lane(one_lane_threads);
572 }
573 let mut lanes = items.min(narrowed).min(work_lane_cap(model, total_ns));
574 let gates_pass = lanes >= 2
575 && items >= model.min_items
576 && items >= lanes.saturating_mul(model.min_items_per_lane);
577 if !gates_pass {
578 return one_lane(one_lane_threads);
579 }
580 // Tiny items keep lanes (a team would pay barriers for nothing); otherwise
581 // the batch alone must fill the narrowed budget and split near-evenly.
582 let tiny = item_ns < model.width.serial_below_ns;
583 let balanced = items.saturating_mul(10)
584 >= 9usize
585 .saturating_mul(lanes)
586 .saturating_mul(items.div_ceil(lanes));
587 if !(tiny || (items >= narrowed && balanced)) {
588 lanes = 1;
589 }
590 BatchDecision {
591 lanes,
592 item_threads: if lanes > 1 { 1 } else { one_lane_threads },
593 }
594}
595
596/// Most lanes the estimated total work can fill, given
597/// [`WorkModel::min_work_per_lane_ns`]. Zero disables the gate.
598fn work_lane_cap(model: &WorkModel, total_ns: f64) -> usize {
599 if model.min_work_per_lane_ns > 0.0 {
600 let raw = total_ns / model.min_work_per_lane_ns;
601 if raw.is_finite() && raw >= 1.0 {
602 raw.floor() as usize
603 } else {
604 1
605 }
606 } else {
607 usize::MAX
608 }
609}
610
611/// One lane carrying `threads` item threads.
612fn one_lane(threads: usize) -> BatchDecision {
613 BatchDecision {
614 lanes: 1,
615 item_threads: threads.max(1),
616 }
617}
618
619/// Estimated cost of one item in nanoseconds.
620fn item_ns(model: &WorkModel, item_muladds: f64) -> f64 {
621 let rate = if model.muladds_per_ns > 0.0 {
622 model.muladds_per_ns
623 } else {
624 1.0
625 };
626 model.item_overhead_ns + item_muladds / rate
627}
628
629/// Team width for `serial_ns` at `budget`, minimizing `WidthPolicy`.
630///
631/// The one body lives in [`crate::width`], which [`Exec::width_for`](crate::Exec::width_for)
632/// also calls; the pure module keeps no copy.
633fn width_for(serial_ns: f64, policy: &WidthPolicy, budget: usize) -> usize {
634 crate::width::width_for(serial_ns, policy, budget)
635}
636
637#[cfg(test)]
638mod tests {
639 use super::*;
640
641 fn budget(n: usize) -> NonZeroUsize {
642 NonZeroUsize::new(n).unwrap()
643 }
644
645 /// Independent copy of the legacy packed computation.
646 ///
647 /// `plan/mod.rs:641` narrows the `Exec` with `width_for(total_ns)`, then
648 /// `driver/batch.rs:38` recomputes the width against that narrowed budget
649 /// and applies the tiny/balanced rule. Returns `(lanes, item_threads)` where
650 /// `item_threads` is the narrowed team width the route uses with one lane.
651 fn legacy_packed(items: usize, item_muladds: f64, host_budget: usize) -> (usize, usize) {
652 let policy = WidthPolicy::default();
653 // `NS_PER_FLOP` applies to flops, and the packed driver counts two flops
654 // per multiply-add.
655 let item_ns = 2.0 * item_muladds * 0.05;
656 let total_ns = item_ns * items as f64;
657 let narrowed = legacy_width_for(total_ns, &policy, host_budget);
658 if items < 2 || narrowed < 2 {
659 return (1, narrowed);
660 }
661 let width = legacy_width_for(total_ns, &policy, narrowed);
662 if width < 2 {
663 return (1, narrowed);
664 }
665 let tiny = item_ns < policy.serial_below_ns;
666 let lanes = items.min(width);
667 let balanced = 10 * items >= 9 * lanes * items.div_ceil(lanes);
668 let lanes = if tiny || (items >= narrowed && balanced) {
669 lanes
670 } else {
671 1
672 };
673 (lanes, if lanes > 1 { 1 } else { narrowed })
674 }
675
676 /// Byte-for-byte copy of `Exec::width_for`'s arithmetic.
677 fn legacy_width_for(serial_ns: f64, policy: &WidthPolicy, budget: usize) -> usize {
678 if budget == 1 || serial_ns.is_nan() || serial_ns < policy.serial_below_ns {
679 return 1;
680 }
681 let cost = |k: usize| {
682 if k == 1 {
683 serial_ns
684 } else {
685 policy.entry_base_ns + policy.entry_per_thread_ns * k as f64 + serial_ns / k as f64
686 }
687 };
688 (1..=budget)
689 .min_by(|&a, &b| cost(a).total_cmp(&cost(b)))
690 .unwrap_or(1)
691 }
692
693 #[test]
694 fn auto_matches_legacy_packed_table() {
695 let policy = BatchPolicy::default();
696 let items = [0usize, 1, 2, 3, 7, 8, 9, 16, 64, 512];
697 let muladds = [
698 1.0,
699 64.0,
700 4_096.0,
701 32_768.0,
702 262_144.0,
703 1_000_000.0,
704 8_000_000.0,
705 ];
706 let budgets = [1usize, 2, 3, 4, 8, 16];
707 for &h in &items {
708 for &m in &muladds {
709 for &b in &budgets {
710 let query = BatchQuery::new(h, m).with_balanced_lanes(true);
711 let got = batch_decision(&policy, budget(b), &query);
712 let want = legacy_packed(h, m, b);
713 assert_eq!(
714 (got.lanes, got.item_threads),
715 (want.0, want.1),
716 "items={h} item_muladds={m} budget={b}"
717 );
718 }
719 }
720 }
721 }
722
723 #[test]
724 fn generic_auto_matches_tenferro_auto_lanes_table() {
725 // tenferro's `CpuBatchThresholds` defaults: 50 ns per item, 16
726 // multiply-adds per ns, 8 us per lane. This is the table at
727 // tenferro-rs `dot_runtime/tests.rs:3025` (`auto_lane_count_needs_enough_work_per_lane`).
728 let model = WorkModel::default()
729 .with_item_overhead_ns(50.0)
730 .with_muladds_per_ns(16.0)
731 .with_min_work_per_lane_ns(8_000.0);
732 let policy = BatchPolicy::new(BatchStrategy::Auto).with_model(model);
733 let lanes = |item_muladds: f64, batch: usize, threads: usize| {
734 batch_decision(
735 &policy,
736 budget(threads),
737 &BatchQuery::new(batch, item_muladds),
738 )
739 .lanes
740 };
741 // 4^3 = 64 multiply-adds: per-call overhead dominates, so many items
742 // fill a lane.
743 assert_eq!(lanes(64.0, 1024, 4), 4);
744 assert_eq!(lanes(64.0, 64, 4), 1);
745 // Larger items need fewer per lane (16^3 is 306 ns, so 27 per lane);
746 // lanes never exceed the threads.
747 assert_eq!(lanes(32_768.0, 8, 16), 2);
748 assert_eq!(lanes(4_096.0, 128, 8), 4);
749 // One lane of work is no fan-out.
750 assert_eq!(lanes(4_096.0, 32, 8), 1);
751 // Two items of tiny work still pass the item gate but not the work
752 // gate: zero minimum work lets the item thresholds decide alone.
753 let no_work_gate =
754 BatchPolicy::new(BatchStrategy::Auto).with_model(model.with_min_work_per_lane_ns(0.0));
755 assert_eq!(
756 batch_decision(&no_work_gate, budget(4), &BatchQuery::new(64, 64.0)).lanes,
757 4
758 );
759 assert_eq!(
760 batch_decision(&no_work_gate, budget(4), &BatchQuery::new(1, 64.0)).lanes,
761 1
762 );
763 }
764
765 #[test]
766 fn default_matches_the_static_policy() {
767 // `Exec::batch_policy` hands a `Serial` context the static policy, so the
768 // two must not drift.
769 assert_eq!(BatchPolicy::default(), DEFAULT_BATCH_POLICY);
770 assert_eq!(WorkModel::default(), DEFAULT_BATCH_POLICY.model);
771 }
772
773 /// Independent copy of the pre-#123 faer schedule.
774 ///
775 /// `strategy/faer.rs` computed `inner = width_for(item_ns)` and
776 /// `total = width_for(item_ns * count)` and picked outer lanes only when
777 /// `inner == 1 && total > 1 && count > 1`. Returns `(lanes, item_threads)`.
778 fn legacy_faer(count: usize, item_muladds: f64, budget: usize) -> (usize, usize) {
779 let policy = WidthPolicy::default();
780 let item_ns = item_muladds / 10.0;
781 let total_ns = item_ns * count as f64;
782 let inner = legacy_width_for(item_ns, &policy, budget);
783 let total = legacy_width_for(total_ns, &policy, budget);
784 if inner == 1 && total > 1 && count > 1 {
785 (total.min(count), 1)
786 } else {
787 (1, inner)
788 }
789 }
790
791 #[test]
792 fn auto_matches_legacy_faer_table() {
793 let policy = BatchPolicy::default();
794 let items = [1usize, 2, 3, 8, 16, 64, 512];
795 let muladds = [
796 1.0,
797 64.0,
798 4_096.0,
799 32_768.0,
800 262_144.0,
801 1_000_000.0,
802 8_000_000.0,
803 ];
804 let budgets = [1usize, 2, 3, 4, 8, 16];
805 for &count in &items {
806 for &m in &muladds {
807 for &b in &budgets {
808 let query = BatchQuery::new(count, m).with_outer_only_when_item_serial(true);
809 let got = batch_decision(&policy, budget(b), &query);
810 let want = legacy_faer(count, m, b);
811 assert_eq!(
812 (got.lanes, got.item_threads),
813 (want.0, want.1),
814 "count={count} item_muladds={m} budget={b}"
815 );
816 }
817 }
818 }
819 }
820
821 #[test]
822 fn faer_auto_keeps_one_item_on_the_team() {
823 // A large item: the inner route, with the width of one item.
824 let query = BatchQuery::new(4, 1_000_000.0).with_outer_only_when_item_serial(true);
825 let d = batch_decision(&BatchPolicy::default(), budget(8), &query);
826 assert!(d.item_threads > 1, "{d:?}");
827 assert_eq!(d.lanes, 1);
828 }
829
830 #[test]
831 fn exec_carries_the_policy_and_budget_keeps_it() {
832 let sp = strided_exec::Pool::builder().threads(4).build().unwrap();
833 let pool = crate::Pool::borrow(&sp);
834 let exec = crate::Exec::new(&pool, 4)
835 .unwrap()
836 .with_batch_policy(BatchPolicy::new(BatchStrategy::Outer));
837 assert_eq!(exec.batch_policy().strategy, BatchStrategy::Outer);
838 let narrowed = exec.limit(2.try_into().unwrap());
839 assert_eq!(narrowed.budget(), 2);
840 assert_eq!(narrowed.batch_policy().strategy, BatchStrategy::Outer);
841 assert_eq!(
842 crate::Exec::serial().batch_policy().strategy,
843 BatchStrategy::Auto
844 );
845 }
846
847 #[test]
848 fn forced_strategies_are_clamped_not_fallen_back() {
849 let query = BatchQuery::new(3, 1_000.0);
850 // Sequential ignores the budget and the item count.
851 assert_eq!(
852 batch_decision(
853 &BatchPolicy::new(BatchStrategy::Sequential),
854 budget(8),
855 &query
856 ),
857 BatchDecision {
858 lanes: 1,
859 item_threads: 1
860 }
861 );
862 // Outer uses min(items, budget), each item serial.
863 assert_eq!(
864 batch_decision(&BatchPolicy::new(BatchStrategy::Outer), budget(8), &query),
865 BatchDecision {
866 lanes: 3,
867 item_threads: 1
868 }
869 );
870 // Inner is one lane with the whole budget.
871 assert_eq!(
872 batch_decision(&BatchPolicy::new(BatchStrategy::Inner), budget(8), &query),
873 BatchDecision {
874 lanes: 1,
875 item_threads: 8
876 }
877 );
878 // Split is clamped to the items and the budget, then shares the rest.
879 assert_eq!(
880 batch_decision(
881 &BatchPolicy::new(BatchStrategy::Split { lanes: budget(8) }),
882 budget(8),
883 &query
884 ),
885 BatchDecision {
886 lanes: 3,
887 item_threads: 2
888 }
889 );
890 // A split wider than the budget clamps to the budget.
891 let many = BatchQuery::new(64, 1_000.0);
892 assert_eq!(
893 batch_decision(
894 &BatchPolicy::new(BatchStrategy::Split { lanes: budget(64) }),
895 budget(4),
896 &many
897 ),
898 BatchDecision {
899 lanes: 4,
900 item_threads: 1
901 }
902 );
903 // Empty items clamp to one lane, never zero.
904 let empty = BatchQuery::new(0, 1.0);
905 assert_eq!(
906 batch_decision(&BatchPolicy::new(BatchStrategy::Outer), budget(4), &empty),
907 BatchDecision {
908 lanes: 1,
909 item_threads: 1
910 }
911 );
912 assert_eq!(
913 batch_decision(
914 &BatchPolicy::new(BatchStrategy::Split { lanes: budget(2) }),
915 budget(4),
916 &empty
917 ),
918 BatchDecision {
919 lanes: 1,
920 item_threads: 4
921 }
922 );
923 }
924
925 #[test]
926 fn every_decision_fits_the_budget() {
927 // A route that honours a decision runs `lanes` lanes of `item_threads`
928 // item threads (faer does), so the split must never exceed the budget the
929 // route was given: that is what keeps the call inside the host's pool.
930 let budgets = [1usize, 2, 3, 4, 8, 16];
931 let items = [0usize, 1, 2, 3, 7, 8, 16, 64];
932 let policies = [
933 BatchStrategy::Auto,
934 BatchStrategy::Sequential,
935 BatchStrategy::Outer,
936 BatchStrategy::Inner,
937 BatchStrategy::Split { lanes: budget(1) },
938 BatchStrategy::Split { lanes: budget(2) },
939 BatchStrategy::Split { lanes: budget(8) },
940 ];
941 for &b in &budgets {
942 for &n in &items {
943 for strategy in policies {
944 let d = batch_decision(
945 &BatchPolicy::new(strategy),
946 budget(b),
947 &BatchQuery::new(n, 512.0),
948 );
949 assert!(
950 d.lanes >= 1 && d.lanes * d.item_threads <= b,
951 "strategy {strategy:?}, budget {b}, items {n}: {d:?}"
952 );
953 }
954 }
955 }
956 }
957
958 #[test]
959 fn forced_strategies_use_the_original_budget() {
960 // Auto narrows the whole batch to six lanes, but Outer is evaluated
961 // against the original host budget of eight.
962 let query = BatchQuery::new(9, 200_000.0).with_balanced_lanes(true);
963 let auto = batch_decision(&BatchPolicy::default(), budget(8), &query);
964 assert_eq!((auto.lanes, auto.item_threads), (6, 1));
965 let outer = batch_decision(&BatchPolicy::new(BatchStrategy::Outer), budget(8), &query);
966 assert_eq!((outer.lanes, outer.item_threads), (8, 1));
967 }
968
969 #[test]
970 fn item_parallel_capability_keeps_item_threads_one() {
971 // A moderate single item widens the narrow route to the team budget; a
972 // route that forbids item parallelism keeps one thread.
973 let query = BatchQuery::new(1, 1_000_000.0);
974 let allowed = batch_decision(&BatchPolicy::default(), budget(8), &query);
975 assert!(allowed.item_threads > 1, "{allowed:?}");
976 let forbidden = query.with_item_parallel(false);
977 let d = batch_decision(&BatchPolicy::default(), budget(8), &forbidden);
978 assert_eq!((d.lanes, d.item_threads), (1, 1));
979 }
980
981 #[test]
982 fn work_gate_reduces_lanes_and_disabling_it_does_not() {
983 let query = BatchQuery::new(8, 8_000_000.0).with_balanced_lanes(true);
984 let disabled = WorkModel::default();
985 let enabled = WorkModel::default().with_min_work_per_lane_ns(1_000_000.0);
986 let d0 = batch_decision(
987 &BatchPolicy::new(BatchStrategy::Auto).with_model(disabled),
988 budget(8),
989 &query,
990 );
991 let d1 = batch_decision(
992 &BatchPolicy::new(BatchStrategy::Auto).with_model(enabled),
993 budget(8),
994 &query,
995 );
996 assert!(d0.lanes > d1.lanes, "{d0:?} vs {d1:?}");
997 }
998
999 #[test]
1000 fn balanced_lanes_can_hold_a_lane_back() {
1001 // Nine items on eight lanes leave one lane with two: the packed rule
1002 // keeps the team instead of cutting an unbalanced batch.
1003 let query = BatchQuery::new(9, 1_000_000.0).with_item_parallel(true);
1004 let packed = batch_decision(
1005 &BatchPolicy::new(BatchStrategy::Auto).with_model(WorkModel::default()),
1006 budget(8),
1007 &query.with_balanced_lanes(true),
1008 );
1009 let generic = batch_decision(
1010 &BatchPolicy::new(BatchStrategy::Auto).with_model(WorkModel::default()),
1011 budget(8),
1012 &query.with_balanced_lanes(false),
1013 );
1014 assert_eq!(packed.lanes, 1);
1015 assert!(generic.lanes > 1, "{generic:?}");
1016 }
1017}