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