rlx_ir/
phase.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Streaming inference phases attached to LIR nodes (plan #16 / #28).
6
7use std::collections::HashMap;
8
9use crate::{Graph, NodeId, Op};
10
11/// Where in a streaming forward pass a node belongs.
12#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
14pub enum Phase {
15    /// One-time setup: embedding lookup, KV init, prompt prefill.
16    Prologue,
17    /// Per-token decode loop body.
18    SteadyState,
19    /// Final projection, sampling, detokenization.
20    Epilogue,
21}
22
23impl Phase {
24    pub fn order(self) -> u8 {
25        match self {
26            Self::Prologue => 0,
27            Self::SteadyState => 1,
28            Self::Epilogue => 2,
29        }
30    }
31}
32
33/// Per-node phase assignment for a streaming graph.
34#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
35#[derive(Debug, Clone, Default, PartialEq, Eq)]
36pub struct PhaseSchedule {
37    map: HashMap<NodeId, Phase>,
38}
39
40impl PhaseSchedule {
41    pub fn new() -> Self {
42        Self::default()
43    }
44
45    pub fn set(&mut self, node: NodeId, phase: Phase) {
46        self.map.insert(node, phase);
47    }
48
49    pub fn get(&self, node: NodeId) -> Option<Phase> {
50        self.map.get(&node).copied()
51    }
52
53    pub fn iter(&self) -> impl Iterator<Item = (NodeId, Phase)> + '_ {
54        self.map.iter().map(|(&id, &p)| (id, p))
55    }
56
57    pub fn len(&self) -> usize {
58        self.map.len()
59    }
60
61    pub fn is_empty(&self) -> bool {
62        self.map.is_empty()
63    }
64
65    /// Nodes in a given phase, in schedule order when `schedule` is provided,
66    /// otherwise sorted by [`NodeId`].
67    pub fn nodes_in(&self, phase: Phase) -> Vec<NodeId> {
68        self.nodes_in_ordered(phase, None)
69    }
70
71    pub fn nodes_in_ordered(&self, phase: Phase, schedule: Option<&[NodeId]>) -> Vec<NodeId> {
72        if let Some(order) = schedule {
73            return order
74                .iter()
75                .copied()
76                .filter(|id| self.get(*id) == Some(phase))
77                .collect();
78        }
79        let mut v: Vec<NodeId> = self
80            .map
81            .iter()
82            .filter_map(|(&id, &p)| if p == phase { Some(id) } else { None })
83            .collect();
84        v.sort();
85        v
86    }
87}
88
89/// Heuristic phase classifier for optimized MIR graphs.
90pub fn derive_phases(graph: &Graph) -> PhaseSchedule {
91    let mut sched = PhaseSchedule::new();
92    let n = graph.len();
93    if n == 0 {
94        return sched;
95    }
96
97    let mut last_compute_step: Option<usize> = None;
98    let mut last_sample_step: Option<usize> = None;
99    for (step, node) in graph.nodes().iter().enumerate() {
100        match &node.op {
101            Op::Sample { .. } | Op::TopK { .. } | Op::RngNormal { .. } | Op::RngUniform { .. } => {
102                last_sample_step = Some(step);
103            }
104            Op::MatMul
105            | Op::FusedMatMulBiasAct { .. }
106            | Op::Attention { .. }
107            | Op::FusedAttentionBlock { .. }
108            | Op::FusedTransformerLayer { .. }
109            | Op::DotGeneral { .. }
110            | Op::GroupedMatMul
111            | Op::DequantGroupedMatMul { .. }
112            | Op::DequantGroupedMatMulMlx { .. }
113            | Op::DequantMoEWeights { .. }
114            | Op::LoraMatMul { .. }
115            | Op::PartitionedConv { .. }
116            | Op::DequantMatMul { .. }
117            | Op::SynthMatMul { .. }
118            | Op::GatedDeltaNet { .. }
119            | Op::Lstm { .. }
120            | Op::Gru { .. }
121            | Op::Rnn { .. }
122            | Op::Mamba2 { .. } => {
123                last_compute_step = Some(step);
124            }
125            _ => {}
126        }
127    }
128
129    for (step, node) in graph.nodes().iter().enumerate() {
130        let phase = match &node.op {
131            Op::Input { .. } | Op::Param { .. } | Op::Constant { .. } => Phase::Prologue,
132            Op::Sample { .. } | Op::TopK { .. } | Op::RngNormal { .. } | Op::RngUniform { .. } => {
133                Phase::Epilogue
134            }
135            _ => {
136                if let Some(last) = last_sample_step {
137                    if step > last
138                        || (last_compute_step.is_some() && Some(step) > last_compute_step)
139                    {
140                        Phase::Epilogue
141                    } else {
142                        Phase::SteadyState
143                    }
144                } else if let Some(last) = last_compute_step {
145                    if step > last {
146                        Phase::Epilogue
147                    } else {
148                        Phase::SteadyState
149                    }
150                } else {
151                    Phase::SteadyState
152                }
153            }
154        };
155        sched.set(node.id, phase);
156    }
157    sched
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163    use crate::{DType, Shape};
164
165    #[test]
166    fn derive_phases_classifies_typical_graph() {
167        let f = DType::F32;
168        let mut g = Graph::new("derive");
169        let x = g.input("x", Shape::new(&[1, 8], f));
170        let w = g.param("w", Shape::new(&[8, 4], f));
171        let mm = g.matmul(x, w, Shape::new(&[1, 4], f));
172        let s = g.sample(mm, 0, 1.0, 1.0, 0, Shape::new(&[1], f));
173        g.set_outputs(vec![s]);
174
175        let sched = derive_phases(&g);
176        assert_eq!(sched.get(x), Some(Phase::Prologue));
177        assert_eq!(sched.get(w), Some(Phase::Prologue));
178        assert_eq!(sched.get(mm), Some(Phase::SteadyState));
179        assert_eq!(sched.get(s), Some(Phase::Epilogue));
180    }
181
182    #[test]
183    fn phase_ordering_is_deterministic() {
184        assert!(Phase::Prologue.order() < Phase::SteadyState.order());
185        assert!(Phase::SteadyState.order() < Phase::Epilogue.order());
186    }
187}
rlx_ir/phase.rs

rlx_ir/
phase.rs
