Skip to main content

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}