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 `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}