nxpu_backend_stablehlo/
lib.rs

1//! StableHLO/XLA backend emitter for NxPU (Google Cloud TPU).
2//!
3//! Emits StableHLO MLIR textual format (`.mlir`) from NxPU IR.
4//! StableHLO is the native entry point for Google Cloud TPU via OpenXLA.
5
6use nxpu_analysis::analyze;
7use nxpu_backend_core::{
8    Backend, BackendError, BackendOptions, BackendOutput, Diagnostic, DiagnosticLevel,
9    OutputContent, OutputFile,
10};
11use nxpu_ir::Module;
12
13mod lower;
14
15/// StableHLO backend targeting Google Cloud TPU via OpenXLA.
16#[derive(Debug)]
17pub struct StableHloBackend;
18
19impl Backend for StableHloBackend {
20    fn name(&self) -> &str {
21        "StableHLO"
22    }
23
24    fn targets(&self) -> &[&str] {
25        &["stablehlo", "xla"]
26    }
27
28    fn compile(
29        &self,
30        module: &Module,
31        _opts: &BackendOptions,
32    ) -> Result<BackendOutput, BackendError> {
33        if module.entry_points.is_empty() {
34            return Err(BackendError::Other("no entry points in module".into()));
35        }
36
37        let mut files = Vec::new();
38        let mut diagnostics = Vec::new();
39
40        for (i, ep) in module.entry_points.iter().enumerate() {
41            let pattern = analyze::classify_entry_point(module, i).map_err(|e| {
42                BackendError::Unsupported(format!("entry point '{}': {e}", ep.name))
43            })?;
44            if let analyze::KernelPattern::Unknown { reason } = &pattern {
45                return Err(BackendError::Unsupported(format!(
46                    "entry point '{}': unrecognized pattern: {reason}",
47                    ep.name
48                )));
49            }
50
51            let summary = match &pattern {
52                analyze::KernelPattern::MatMul { .. } => "dot_general",
53                analyze::KernelPattern::QuantizedMatMul { .. } => "quantized matmul",
54                analyze::KernelPattern::ElementWise { op, .. } => op.op_name(),
55                analyze::KernelPattern::Conv2D { .. } => "convolution",
56                analyze::KernelPattern::Pool { .. } => "reduce_window",
57                analyze::KernelPattern::Activation { op, .. } => op.op_name(),
58                analyze::KernelPattern::Reduce { .. } => "reduce",
59                analyze::KernelPattern::Transpose { .. } => "transpose",
60                analyze::KernelPattern::Reshape { .. } => "reshape",
61                analyze::KernelPattern::Normalization { .. } => "batch_norm_inference",
62                analyze::KernelPattern::Concat { .. } => "concatenate",
63                analyze::KernelPattern::Split { .. } => "slice",
64                analyze::KernelPattern::Attention { .. } => "attention",
65                analyze::KernelPattern::Gather { .. } => "gather",
66                analyze::KernelPattern::Scatter { .. } => "scatter",
67                analyze::KernelPattern::ElementWiseChain { .. } => "element-wise chain",
68                analyze::KernelPattern::Unknown { .. } => "Unknown",
69            };
70
71            diagnostics.push(Diagnostic {
72                level: DiagnosticLevel::Info,
73                message: format!("entry point '{}': StableHLO {}", ep.name, summary),
74            });
75
76            let mlir = lower::build_mlir(&pattern, &ep.name)?;
77
78            let filename = if module.entry_points.len() == 1 {
79                "output.mlir".into()
80            } else {
81                format!("{}.mlir", ep.name)
82            };
83
84            files.push(OutputFile {
85                name: filename,
86                content: OutputContent::Text(mlir),
87            });
88        }
89
90        Ok(BackendOutput { files, diagnostics })
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97    use nxpu_backend_core::BackendOptions;
98
99    #[test]
100    fn backend_metadata() {
101        let backend = StableHloBackend;
102        assert_eq!(backend.name(), "StableHLO");
103        assert!(backend.targets().contains(&"stablehlo"));
104        assert!(backend.targets().contains(&"xla"));
105    }
106
107    #[test]
108    fn compile_matmul_wgsl() {
109        let source = std::fs::read_to_string(concat!(
110            env!("CARGO_MANIFEST_DIR"),
111            "/../../examples/matmul.wgsl"
112        ))
113        .unwrap();
114        let module = nxpu_parser::parse(&source).unwrap();
115
116        let output = StableHloBackend
117            .compile(&module, &BackendOptions::default())
118            .unwrap();
119        assert_eq!(output.files.len(), 1);
120        assert_eq!(output.files[0].name, "output.mlir");
121
122        let text = match &output.files[0].content {
123            OutputContent::Text(t) => t,
124            _ => panic!("expected text output"),
125        };
126        assert!(text.contains("stablehlo.dot_general"));
127        assert!(text.contains("module @main"));
128    }
129
130    #[test]
131    fn compile_vecadd_wgsl() {
132        let source = r#"
133@group(0) @binding(0) var<storage, read> a: array<f32>;
134@group(0) @binding(1) var<storage, read> b: array<f32>;
135@group(0) @binding(2) var<storage, read_write> c: array<f32>;
136struct Params { N: u32 }
137@group(0) @binding(3) var<uniform> params: Params;
138@compute @workgroup_size(256)
139fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
140  let idx = gid.x;
141  if (idx >= params.N) { return; }
142  c[idx] = a[idx] + b[idx];
143}
144"#;
145        let module = nxpu_parser::parse(source).unwrap();
146        let output = StableHloBackend
147            .compile(&module, &BackendOptions::default())
148            .unwrap();
149
150        let text = match &output.files[0].content {
151            OutputContent::Text(t) => t,
152            _ => panic!("expected text"),
153        };
154        assert!(text.contains("stablehlo.add"));
155    }
156}
nxpu_backend_stablehlo/lib.rs

nxpu_backend_stablehlo/
lib.rs
