rlx_compile/
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//! Pipeline inspection — dump every HIR / MIR / LIR stage as text.
6
7use std::fmt::Write as _;
8
9use rlx_ir::hir::{HirModule, LowerError};
10use rlx_ir::{inspect_graph_diff, inspect_hir, inspect_lir, inspect_mir, inspect_mir_stats};
11
12use crate::compiler::{CompilePipeline, CompileResult};
13use rlx_fusion::fusion_report::FusionReport;
14
15/// Text dump of each compiler pipeline stage.
16#[derive(Debug, Clone)]
17pub struct PipelineInspect {
18    pub hir: String,
19    pub mir_lowered: String,
20    pub mir_diff: String,
21    pub mir_optimized: String,
22    pub lir: String,
23    pub fusion: String,
24}
25
26impl std::fmt::Display for PipelineInspect {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        writeln_section(f, "HIR", &self.hir)?;
29        writeln_section(f, "MIR (lowered)", &self.mir_lowered)?;
30        if !self.mir_diff.is_empty() {
31            writeln_section(f, "MIR (fusion diff)", &self.mir_diff)?;
32        }
33        writeln_section(f, "MIR (optimized)", &self.mir_optimized)?;
34        writeln_section(f, "FUSION", &self.fusion)?;
35        writeln_section(f, "LIR", &self.lir)
36    }
37}
38
39fn writeln_section(f: &mut std::fmt::Formatter<'_>, title: &str, body: &str) -> std::fmt::Result {
40    let mut header = String::new();
41    banner(&mut header, title);
42    write!(f, "{header}{body}")?;
43    if !body.ends_with('\n') {
44        writeln!(f)?;
45    }
46    Ok(())
47}
48
49/// Inspect every lowering stage for `hir` through `pipeline`.
50pub fn inspect_pipeline(
51    pipeline: &CompilePipeline,
52    hir: HirModule,
53) -> Result<PipelineInspect, LowerError> {
54    let hir_text = inspect_hir(&hir);
55    let mir_raw = CompilePipeline::lower_hir(hir)?;
56    let mir_lowered = inspect_mir(&mir_raw);
57    let mir_before = mir_raw.clone();
58    let (mir_opt, fusion) = pipeline.optimize_with_report(mir_raw);
59    let mir_diff = inspect_graph_diff(mir_before.as_graph(), mir_opt.as_graph());
60    let fusion_text = format!("{}\n{}", fusion, inspect_mir_stats(&mir_opt));
61    let lir = pipeline.plan_lir(mir_opt);
62    Ok(PipelineInspect {
63        hir: hir_text,
64        mir_lowered,
65        mir_diff,
66        mir_optimized: inspect_mir(&lir.mir),
67        lir: inspect_lir(&lir),
68        fusion: fusion_text,
69    })
70}
71
72/// Inspect a completed [`CompileResult`] plus the original HIR text.
73pub fn inspect_compiled(hir_text: &str, result: &CompileResult) -> PipelineInspect {
74    PipelineInspect {
75        hir: hir_text.to_string(),
76        mir_lowered: String::new(),
77        mir_diff: String::new(),
78        mir_optimized: inspect_mir(&result.lir.mir),
79        lir: inspect_lir(&result.lir),
80        fusion: format!("{}", result.fusion),
81    }
82}
83
84/// Write a full pipeline dump when `RLX_IR_DUMP` is set (path prefix or directory).
85pub fn maybe_dump_pipeline(dump: &PipelineInspect, module_name: &str) {
86    let Some(path) = rlx_ir::env::var("RLX_IR_DUMP") else {
87        return;
88    };
89    let target = if path.ends_with('/') || path.ends_with('\\') {
90        format!("{path}{module_name}.ir.txt")
91    } else {
92        path
93    };
94    if let Err(e) = std::fs::write(&target, dump.to_string()) {
95        eprintln!("[rlx] RLX_IR_DUMP write failed ({target}): {e}");
96    } else {
97        eprintln!("[rlx] wrote IR dump to {target}");
98    }
99}
100
101/// Fusion report only (post-optimize diagnostics).
102pub fn inspect_fusion(report: &FusionReport) -> String {
103    format!("{report}")
104}
105
106fn banner(out: &mut String, title: &str) {
107    let line = "═".repeat(title.len() + 4);
108    writeln!(out, "{line}").unwrap();
109    writeln!(out, "══ {title} ══").unwrap();
110    writeln!(out, "{line}").unwrap();
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use rlx_ir::DType;
117    use rlx_ir::Shape;
118
119    fn f32_shape(d: &[usize]) -> Shape {
120        Shape::new(d, DType::F32)
121    }
122
123    #[test]
124    fn inspect_pipeline_covers_all_stages() {
125        let mut hir = HirModule::new("probe");
126        let x = hir.input("x", f32_shape(&[2, 64]));
127        let w = hir.param("w", f32_shape(&[64, 64]));
128        let h = hir.linear(x, w, None, None, f32_shape(&[2, 64]));
129        hir.outputs = vec![h];
130
131        let pipe = CompilePipeline::default();
132        let dump = inspect_pipeline(&pipe, hir).expect("inspect");
133        assert!(dump.hir.contains("hir @probe"));
134        assert!(dump.mir_lowered.contains("mir @probe"));
135        assert!(dump.mir_optimized.contains("mir @probe"));
136        assert!(dump.lir.contains("lir @probe"));
137        assert!(dump.fusion.contains("nodes="));
138        let full = dump.to_string();
139        assert!(full.contains("══ HIR ══"));
140        assert!(full.contains("══ LIR ══"));
141    }
142}
rlx_compile/inspect.rs

rlx_compile/
inspect.rs
