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