rlx_ir/
inspect.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//! Text exporters for inspecting HIR / MIR / LIR during lowering.
6//!
7//! Use [`inspect_hir`], [`inspect_mir`], and [`inspect_lir`] to dump
8//! each pipeline stage as human-readable text (similar to LLVM `-print-*`
9//! flags). [`inspect_graph`] is the MIR body formatter shared by MIR and
10//! LIR dumps.
11
12use std::collections::BTreeMap;
13use std::fmt::Write as _;
14
15use crate::hir::{HirModule, HirNode, HirOp};
16use crate::lir::{LirBufferPlan, LirModule, LirViewAlias};
17use crate::mir::MirModule;
18use crate::phase::Phase;
19use crate::pretty::{header_line, op_kinds_line, pretty_print};
20use crate::{Graph, NodeId};
21
22/// Annotated HIR module dump.
23pub fn inspect_hir(hir: &HirModule) -> String {
24    let mut out = String::new();
25    writeln!(
26        out,
27        "hir @{} ({} nodes, {} outputs, fusion={:?})",
28        hir.name,
29        hir.len(),
30        hir.outputs.len(),
31        hir.fusion_policy,
32    )
33    .unwrap();
34    writeln!(out, "{}", hir_op_kinds_line(hir)).unwrap();
35    writeln!(out).unwrap();
36
37    let mut tag_w = 0usize;
38    for node in hir.nodes() {
39        let t = hir_node_tag(node);
40        tag_w = tag_w.max(t.len());
41    }
42
43    for node in hir.nodes() {
44        let tag = hir_node_tag(node);
45        write!(out, "  {tag:<width$} = ", width = tag_w).unwrap();
46        write!(out, "{}", format_hir_op(&node.op)).unwrap();
47        if !node.inputs.is_empty() {
48            write!(out, "(").unwrap();
49            for (i, inp) in node.inputs.iter().enumerate() {
50                if i > 0 {
51                    write!(out, ", ").unwrap();
52                }
53                write!(out, "{inp}").unwrap();
54            }
55            write!(out, ")").unwrap();
56        }
57        write!(out, " : {}", node.shape).unwrap();
58        if hir.outputs.contains(&node.id) {
59            write!(out, "  ← output").unwrap();
60        }
61        writeln!(out).unwrap();
62    }
63    if !hir.outputs.is_empty() {
64        write!(out, "  return ").unwrap();
65        for (i, o) in hir.outputs.iter().enumerate() {
66            if i > 0 {
67                write!(out, ", ").unwrap();
68            }
69            write!(out, "{o}").unwrap();
70        }
71        writeln!(out).unwrap();
72    }
73    out
74}
75
76/// Annotated MIR module dump (optimized tensor DAG).
77pub fn inspect_mir(mir: &MirModule) -> String {
78    inspect_mir_with_diff(mir, None)
79}
80
81/// MIR dump with optional fusion diff against a pre-optimize snapshot.
82pub fn inspect_mir_with_diff(mir: &MirModule, before: Option<&MirModule>) -> String {
83    let g = mir.as_graph();
84    let mut out = String::new();
85    writeln!(out, "mir @{} {{", mir.name()).unwrap();
86    if let Some(b) = before {
87        writeln!(out).unwrap();
88        out.push_str(&inspect_graph_diff(b.as_graph(), g));
89        writeln!(out).unwrap();
90        writeln!(out, "--- graph ---").unwrap();
91    }
92    writeln!(out).unwrap();
93    out.push_str(&pretty_print(g));
94    if !out.ends_with('\n') {
95        out.push('\n');
96    }
97    write!(out, "}}").unwrap();
98    out
99}
100
101/// Diff two MIR snapshots (typically pre/post fusion).
102pub fn inspect_mir_diff(before: &MirModule, after: &MirModule) -> String {
103    inspect_graph_diff(before.as_graph(), after.as_graph())
104}
105
106/// Summarize graph changes between pipeline stages.
107pub fn inspect_graph_diff(before: &Graph, after: &Graph) -> String {
108    use std::collections::BTreeMap;
109
110    let mut out = String::new();
111    writeln!(
112        out,
113        "  diff: {} → {} nodes ({} → {} outputs)",
114        before.len(),
115        after.len(),
116        before.outputs.len(),
117        after.outputs.len(),
118    )
119    .unwrap();
120
121    let count_kinds = |g: &Graph| {
122        let mut h: BTreeMap<String, i32> = BTreeMap::new();
123        for n in g.nodes() {
124            *h.entry(format!("{:?}", n.op.kind())).or_insert(0) += 1;
125        }
126        h
127    };
128    let b = count_kinds(before);
129    let a = count_kinds(after);
130    let mut keys: Vec<String> = b.keys().chain(a.keys()).cloned().collect();
131    keys.sort();
132    keys.dedup();
133    let mut changes = Vec::new();
134    for k in keys {
135        let d = a.get(&k).copied().unwrap_or(0) - b.get(&k).copied().unwrap_or(0);
136        if d != 0 {
137            changes.push(format!("{k}{d:+}"));
138        }
139    }
140    if !changes.is_empty() {
141        writeln!(out, "  op delta: {}", changes.join(", ")).unwrap();
142    }
143    out
144}
145
146/// Annotated LIR dump: optimized MIR + buffer plan + schedule.
147pub fn inspect_lir(lir: &LirModule) -> String {
148    let mut out = String::new();
149    writeln!(out, "lir @{} {{", lir.name()).unwrap();
150    writeln!(out, "  fingerprint: {:016x}", lir.fingerprint().0).unwrap();
151    writeln!(out).unwrap();
152    out.push_str(&inspect_buffer_plan(&lir.buffers));
153    if !lir.buffers.phases.is_empty() {
154        writeln!(out).unwrap();
155        out.push_str(&inspect_phases(&lir.buffers));
156    }
157    if !lir.buffers.io.inputs.is_empty() || !lir.buffers.io.params.is_empty() {
158        writeln!(out).unwrap();
159        out.push_str(&inspect_io_manifest(&lir.buffers));
160    }
161    writeln!(out).unwrap();
162    writeln!(out, "--- mir ---").unwrap();
163    out.push_str(&pretty_print(lir.as_graph()));
164    if !out.ends_with('\n') {
165        out.push('\n');
166    }
167    write!(out, "}}").unwrap();
168    out
169}
170
171/// Annotated graph dump (MIR body). Alias for [`pretty_print`].
172pub fn inspect_graph(g: &Graph) -> String {
173    pretty_print(g)
174}
175
176/// One-line HIR summary (header + op histogram).
177pub fn inspect_hir_stats(hir: &HirModule) -> String {
178    format!(
179        "hir @{} ({} nodes, {} outputs, fusion={:?})\n{}",
180        hir.name,
181        hir.len(),
182        hir.outputs.len(),
183        hir.fusion_policy,
184        hir_op_kinds_line(hir),
185    )
186}
187
188/// One-line MIR summary.
189pub fn inspect_mir_stats(mir: &MirModule) -> String {
190    let g = mir.as_graph();
191    format!(
192        "mir @{} — {}\n{}",
193        mir.name(),
194        header_line(g),
195        op_kinds_line(g),
196    )
197}
198
199/// Buffer plan section for LIR inspection.
200pub fn inspect_buffer_plan(plan: &LirBufferPlan) -> String {
201    let mut out = String::new();
202    let saved = plan.bytes_saved();
203    let naive = plan.total_unshared_bytes();
204    writeln!(
205        out,
206        "  arena: {} bytes (saved {} vs {} naive, align={})",
207        plan.arena_size, saved, naive, plan.alignment,
208    )
209    .unwrap();
210    writeln!(
211        out,
212        "  schedule: {} nodes, {} views",
213        plan.schedule.len(),
214        plan.view_aliases.len(),
215    )
216    .unwrap();
217    if !plan.dynamic_symbols.is_empty() {
218        let syms: Vec<String> = plan
219            .dynamic_symbols
220            .iter()
221            .map(|s| format!("?{s}"))
222            .collect();
223        writeln!(out, "  dynamic: {}", syms.join(", ")).unwrap();
224    }
225    writeln!(out).unwrap();
226    writeln!(out, "  # offset\tsize\tnode").unwrap();
227
228    let mut rows: Vec<(usize, usize, NodeId)> = plan
229        .assignments
230        .iter()
231        .map(|(id, slot)| (slot.offset, slot.size, *id))
232        .collect();
233    rows.sort_by_key(|(off, _, _)| *off);
234    for (off, sz, id) in rows {
235        let sched = plan
236            .schedule
237            .iter()
238            .position(|&n| n == id)
239            .map(|i| format!(" sched={i}"))
240            .unwrap_or_default();
241        let view = plan
242            .view_aliases
243            .get(&id)
244            .map(|LirViewAlias { root, byte_offset }| format!(" view→{root}+{byte_offset}"))
245            .unwrap_or_default();
246        let phase = plan
247            .phases
248            .get(id)
249            .map(|p| format!(" {p:?}"))
250            .unwrap_or_default();
251        writeln!(out, "  {off}\t{sz}\t{id}{sched}{view}{phase}").unwrap();
252    }
253    out
254}
255
256fn inspect_phases(plan: &LirBufferPlan) -> String {
257    let mut out = String::from("  phases:\n");
258    for phase in [Phase::Prologue, Phase::SteadyState, Phase::Epilogue] {
259        let nodes = plan.nodes_in_phase(phase);
260        if !nodes.is_empty() {
261            writeln!(out, "    {phase:?}: {nodes:?}").unwrap();
262        }
263    }
264    out
265}
266
267fn inspect_io_manifest(plan: &LirBufferPlan) -> String {
268    let mut out = String::from("  io:\n");
269    for (name, id) in &plan.io.inputs {
270        writeln!(out, "    input \"{name}\" → {id}").unwrap();
271    }
272    for (name, id) in &plan.io.params {
273        writeln!(out, "    param \"{name}\" → {id}").unwrap();
274    }
275    if !plan.io.outputs.is_empty() {
276        write!(out, "    outputs: {:?}", plan.io.outputs).unwrap();
277        out.push('\n');
278    }
279    out
280}
281
282fn hir_op_kinds_line(hir: &HirModule) -> String {
283    let mut hist: BTreeMap<String, usize> = BTreeMap::new();
284    for node in hir.nodes() {
285        *hist.entry(hir_op_kind(&node.op)).or_insert(0) += 1;
286    }
287    let parts: Vec<String> = hist.into_iter().map(|(k, c)| format!("{k}={c}")).collect();
288    format!("  block ops: {}", parts.join(", "))
289}
290
291fn hir_op_kind(op: &HirOp) -> String {
292    match op {
293        HirOp::Input { .. } => "Input".into(),
294        HirOp::Param { .. } => "Param".into(),
295        HirOp::Constant { .. } => "Constant".into(),
296        HirOp::Linear { .. } => "Linear".into(),
297        HirOp::LinearFused { .. } => "LinearFused".into(),
298        HirOp::SharedLinearPair { .. } => "SharedLinearPair".into(),
299        HirOp::SwiGLU => "SwiGLU".into(),
300        HirOp::ResidualRmsNorm { .. } => "ResidualRmsNorm".into(),
301        HirOp::Attention { .. } => "Attention".into(),
302        HirOp::DepthwiseConv1dCausal { .. } => "DepthwiseConv1dCausal".into(),
303        HirOp::DequantMatMul { .. } => "DequantMatMul".into(),
304        HirOp::GatedDeltaNet { .. } => "GatedDeltaNet".into(),
305        HirOp::Lstm { .. } => "Lstm".into(),
306        HirOp::Gru { .. } => "Gru".into(),
307        HirOp::RoPE { .. } => "RoPE".into(),
308        HirOp::RmsNorm { .. } => "RmsNorm".into(),
309        HirOp::Mir(_) => "Mir".into(),
310        HirOp::LlamaDecoderBlock { .. } => "LlamaDecoderBlock".into(),
311        HirOp::Qwen35MtpHead { .. } => "Qwen35MtpHead".into(),
312    }
313}
314
315fn hir_node_tag(node: &HirNode) -> String {
316    let label: Option<String> = match &node.op {
317        HirOp::Input { name } => Some(format!("input \"{name}\"")),
318        HirOp::Param { name } => Some(format!("param \"{name}\"")),
319        _ => node.name.as_deref().map(|s| format!("\"{s}\"")),
320    };
321    match label {
322        Some(s) => format!("{} [{s}]", node.id),
323        None => format!("{}", node.id),
324    }
325}
326
327fn format_hir_op(op: &HirOp) -> String {
328    match op {
329        HirOp::Input { name } => format!("input(\"{name}\")"),
330        HirOp::Param { name } => format!("param(\"{name}\")"),
331        HirOp::Constant { data } => format!("constant({} bytes)", data.len()),
332        HirOp::Linear {
333            activation,
334            has_bias,
335        } => {
336            let mut s = String::from("linear");
337            if *has_bias {
338                s.push_str("+bias");
339            }
340            if let Some(act) = activation {
341                write!(s, "+{act:?}").unwrap();
342            }
343            s
344        }
345        HirOp::LinearFused { activation } => match activation {
346            Some(act) => format!("linear_fused({act:?})"),
347            None => "linear_fused".into(),
348        },
349        HirOp::SharedLinearPair { slot } => format!("shared_linear_pair(out={slot})"),
350        HirOp::SwiGLU => "swiglu_ffn".into(),
351        HirOp::ResidualRmsNorm { eps } => format!("residual_rms_norm(eps={eps})"),
352        HirOp::Attention {
353            num_heads,
354            head_dim,
355            mask,
356        } => format!("attention(heads={num_heads}, dim={head_dim}, mask={mask:?})"),
357        HirOp::DepthwiseConv1dCausal { kernel_size } => {
358            format!("depthwise_conv1d_causal(k={kernel_size})")
359        }
360        HirOp::DequantMatMul { scheme } => format!("dequant_matmul({scheme})"),
361        HirOp::GatedDeltaNet {
362            state_size,
363            carry_state,
364            gate_per_channel,
365        } => {
366            let gpc = if *gate_per_channel { ",gpc" } else { "" };
367            if *carry_state {
368                format!("gated_delta_net(n={state_size},carry{gpc})")
369            } else {
370                format!("gated_delta_net(n={state_size}{gpc})")
371            }
372        }
373        HirOp::Lstm {
374            hidden_size,
375            num_layers,
376            bidirectional,
377            ..
378        } => {
379            let dir = if *bidirectional { "bi" } else { "uni" };
380            format!("lstm(h={hidden_size},layers={num_layers},{dir})")
381        }
382        HirOp::Gru {
383            hidden_size,
384            num_layers,
385            bidirectional,
386            ..
387        } => {
388            let dir = if *bidirectional { "bi" } else { "uni" };
389            format!("gru(h={hidden_size},layers={num_layers},{dir})")
390        }
391        HirOp::RoPE { head_dim, n_rot } => format!("rope(d={head_dim}, n_rot={n_rot})"),
392        HirOp::RmsNorm { eps } => format!("rms_norm(eps={eps})"),
393        HirOp::LlamaDecoderBlock {
394            num_heads,
395            head_dim,
396            num_kv_heads,
397            eps,
398            mask,
399            rope_style,
400        } => format!(
401            "llama_decoder_block(heads={num_heads}, dim={head_dim}, kv={num_kv_heads}, eps={eps}, mask={mask:?}, rope={rope_style:?})"
402        ),
403        HirOp::Qwen35MtpHead {
404            num_heads,
405            head_dim,
406            mtp_vocab,
407            ..
408        } => format!("qwen35_mtp_head(heads={num_heads}, dim={head_dim}, vocab={mtp_vocab})"),
409        HirOp::Mir(inner) => format!("mir({inner})"),
410    }
411}
412
413// ── convenience methods on pipeline types ───────────────────────────────
414
415impl HirModule {
416    /// Text dump for inspection. Alias for [`inspect_hir`].
417    pub fn inspect(&self) -> String {
418        inspect_hir(self)
419    }
420}
421
422impl MirModule {
423    /// Text dump for inspection. Alias for [`inspect_mir`].
424    pub fn inspect(&self) -> String {
425        inspect_mir(self)
426    }
427}
428
429impl LirModule {
430    /// Text dump for inspection. Alias for [`inspect_lir`].
431    pub fn inspect(&self) -> String {
432        inspect_lir(self)
433    }
434}
435
436#[cfg(test)]
437mod tests {
438    use super::*;
439    use crate::DType;
440    use crate::Shape;
441
442    fn f32_shape(d: &[usize]) -> Shape {
443        Shape::new(d, DType::F32)
444    }
445
446    #[test]
447    fn inspect_hir_includes_blocks_and_outputs() {
448        let mut hir = HirModule::new("layer");
449        let x = hir.input("x", f32_shape(&[2, 128]));
450        let w = hir.param("w", f32_shape(&[128, 128]));
451        let h = hir.linear(x, w, None, None, f32_shape(&[2, 128]));
452        hir.outputs = vec![h];
453
454        let text = inspect_hir(&hir);
455        assert!(text.contains("hir @layer"));
456        assert!(text.contains("linear"));
457        assert!(text.contains("← output"));
458        assert!(text.contains("fusion=Direct"));
459    }
460
461    #[test]
462    fn inspect_mir_wraps_pretty_print() {
463        let mut hir = HirModule::new("m");
464        let x = hir.input("x", f32_shape(&[4]));
465        hir.outputs = vec![x];
466        let mir = hir.lower_to_mir().expect("lower");
467
468        let text = inspect_mir(&mir);
469        assert!(text.contains("mir @m"));
470        assert!(text.contains("graph @m"));
471        assert!(text.contains("input(\"x\")"));
472    }
473
474    #[test]
475    fn named_block_appears_in_hir_dump() {
476        let mut hir = HirModule::new("layer");
477        let x = hir.input("x", f32_shape(&[2, 8]));
478        let w = hir.param("w", f32_shape(&[8, 8]));
479        let out = hir.named("layer0.ffn", |h| {
480            h.linear(x, w, None, None, f32_shape(&[2, 8]))
481        });
482        hir.outputs = vec![out];
483
484        let text = inspect_hir(&hir);
485        assert!(text.contains("layer0.ffn"));
486    }
487
488    #[test]
489    fn provenance_survives_lower() {
490        let mut hir = HirModule::new("m");
491        let x = hir.input("x", f32_shape(&[2, 8]));
492        let w = hir.param("w", f32_shape(&[8, 8]));
493        let out = hir.named("block", |h| h.linear(x, w, None, None, f32_shape(&[2, 8])));
494        hir.outputs = vec![out];
495
496        let mir = hir.lower_to_mir().expect("lower");
497        let text = inspect_mir(&mir);
498        assert!(text.contains("hir=h"));
499        assert!(text.contains("block"));
500    }
501
502    #[test]
503    fn inspect_lir_includes_buffer_plan() {
504        use crate::lir::{LirBufferPlan, LirBufferSlot, LirIoManifest};
505
506        let mut hir = HirModule::new("l");
507        let x = hir.input("x", f32_shape(&[4]));
508        hir.outputs = vec![x];
509        let mir = hir.lower_to_mir().expect("lower");
510        let plan = LirBufferPlan {
511            arena_size: 16,
512            assignments: [(
513                NodeId(0),
514                LirBufferSlot {
515                    offset: 0,
516                    size: 16,
517                },
518            )]
519            .into_iter()
520            .collect(),
521            schedule: vec![NodeId(0)],
522            io: LirIoManifest {
523                inputs: vec![("x".into(), NodeId(0))],
524                ..Default::default()
525            },
526            ..Default::default()
527        };
528        let lir = LirModule::new(mir, plan);
529
530        let text = inspect_lir(&lir);
531        assert!(text.contains("lir @l"));
532        assert!(text.contains("arena: 16 bytes"));
533        assert!(text.contains("fingerprint:"));
534        assert!(text.contains("--- mir ---"));
535    }
536}
rlx_ir/inspect.rs

rlx_ir/
inspect.rs
