metrale_model_layers/circuit_exec/
routes.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The runtime routes the executor runs. A FUSIONS.toml `[[runtime]]` route
4//! (`metrale_circuit::runtime`) is compiled into its own program beside the primary one, at
5//! every mode and row count it applies to. Each step evaluates the route's condition exactly as
6//! the legacy dispatch does and runs the matching program.
7//!
8//! Owner: model-layers circuit executor.
9//! Invariants:
10//! - Every route id has an evaluator here; a rule set naming another is refused at build, never
11//!   run on its primary arm alone.
12//! - `gdn_state_slots_fragmented` holds when some row's GDN state is not at row 0's plus `i`
13//!   slots, the check `ssm_batched_recurrent.rs` makes per layer. Layers that disagree are an
14//!   error: the pools share one slot index per sequence, so they cannot.
15
16use anyhow::{Result, bail, ensure};
17use metrale_circuit::{FusionPlan, RuntimeRoute};
18
19use super::program::{GdnState, Program};
20
21/// 2026-09-30: The routes this executor can evaluate.
22pub const KNOWN: [&str; 1] = ["gdn_state_slots_fragmented"];
23
24/// 2026-09-30: A route's arm, compiled.
25pub struct RoutedProgram {
26    /// 2026-09-30: The route.
27    pub route: RuntimeRoute,
28    /// 2026-09-30: The arm's program.
29    pub program: Program,
30    /// 2026-09-30: The plan it was compiled from.
31    pub plan: FusionPlan,
32}
33
34/// 2026-09-30: Refuse a rule set with a route no evaluator here serves.
35pub fn check_known(runtime: &[RuntimeRoute]) -> Result<()> {
36    if let Some(r) = runtime.iter().find(|r| !KNOWN.contains(&r.id.as_str())) {
37        bail!(
38            "runtime route `{}` has no evaluator in the executor (it knows {KNOWN:?})",
39            r.id
40        );
41    }
42    Ok(())
43}
44
45/// 2026-09-30: Whether `route` holds for a step whose rows read `gdn` (per layer, per row).
46/// `pitch` is each layer's `(h, conv)` slot pitch, `None` for a layer without GDN state.
47pub fn holds(
48    route: &RuntimeRoute,
49    pitch: &[Option<(usize, usize)>],
50    gdn: &[Vec<GdnState>],
51) -> Result<bool> {
52    match route.id.as_str() {
53        "gdn_state_slots_fragmented" => fragmented(pitch, gdn),
54        other => bail!("runtime route `{other}` has no evaluator"),
55    }
56}
57
58fn fragmented(pitch: &[Option<(usize, usize)>], gdn: &[Vec<GdnState>]) -> Result<bool> {
59    let mut verdict: Option<(usize, bool)> = None;
60    for (layer, (p, rows)) in pitch.iter().zip(gdn).enumerate() {
61        let (Some((h, conv)), Some(base)) = (p, rows.first()) else {
62            continue;
63        };
64        let contiguous = rows
65            .iter()
66            .enumerate()
67            .all(|(i, s)| s.h == base.h.offset(i * h) && s.conv == base.conv.offset(i * conv));
68        match verdict {
69            None => verdict = Some((layer, contiguous)),
70            Some((first, c)) => ensure!(
71                c == contiguous,
72                "GDN layers {first} and {layer} disagree on slot contiguity"
73            ),
74        }
75    }
76    Ok(verdict.is_some_and(|(_, contiguous)| !contiguous))
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82    use metrale_gpu_runtime::gpu::DevicePtr;
83
84    fn state(h: u64, conv: u64) -> GdnState {
85        GdnState {
86            h: DevicePtr(h),
87            conv: DevicePtr(conv),
88            h_steps: [DevicePtr::NULL; super::super::program::MAX_VERIFY_STEPS],
89            conv_steps: [DevicePtr::NULL; super::super::program::MAX_VERIFY_STEPS],
90        }
91    }
92
93    fn route(id: &str) -> RuntimeRoute {
94        RuntimeRoute {
95            id: id.into(),
96            when: Default::default(),
97            modes: Default::default(),
98            rows: (2, 128),
99            plans_as: Default::default(),
100            why: "test".into(),
101            cite: "test".into(),
102        }
103    }
104
105    // 2026-09-30: The check is legacy's: each row at row 0 plus `i` pitches, on both the h and
106    // the conv pool. Mutations: comparing only h, or only the first two rows, misses a case.
107    #[test]
108    fn fragmented_holds_exactly_when_a_row_is_out_of_place() {
109        let r = route("gdn_state_slots_fragmented");
110        let pitch = [Some((0x100, 0x10)), None, Some((0x100, 0x10))];
111        let layer = |rows: &[(u64, u64)]| -> Vec<GdnState> {
112            rows.iter().map(|&(h, c)| state(h, c)).collect()
113        };
114        let ok: Vec<GdnState> = layer(&[(0x1000, 0x50), (0x1100, 0x60), (0x1200, 0x70)]);
115        let gdn = vec![ok.clone(), Vec::new(), ok.clone()];
116        assert!(!holds(&r, &pitch, &gdn).unwrap());
117        for bad in [
118            layer(&[(0x1000, 0x50), (0x1100, 0x60), (0x1300, 0x70)]),
119            layer(&[(0x1000, 0x50), (0x1100, 0x60), (0x1200, 0x80)]),
120            layer(&[(0x1000, 0x50), (0x1200, 0x60), (0x1300, 0x70)]),
121        ] {
122            let gdn = vec![bad.clone(), Vec::new(), bad];
123            assert!(holds(&r, &pitch, &gdn).unwrap());
124            let mixed = vec![ok.clone(), Vec::new(), gdn[0].clone()];
125            assert!(
126                holds(&r, &pitch, &mixed)
127                    .unwrap_err()
128                    .to_string()
129                    .contains("disagree")
130            );
131        }
132    }
133
134    #[test]
135    fn an_unknown_route_is_refused_at_build() {
136        assert!(check_known(&[route("gdn_state_slots_fragmented")]).is_ok());
137        let e = check_known(&[route("mystery")]).unwrap_err().to_string();
138        assert!(e.contains("mystery") && e.contains("no evaluator"), "{e}");
139    }
140}
metrale_model_layers/circuit_exec/routes.rs

metrale_model_layers/circuit_exec/
routes.rs
