nxpu_backend_onnx/
lib.rs

1//! ONNX backend emitter for NxPU.
2//!
3//! Analyzes NxPU IR to recognize high-level tensor operations (MatMul,
4//! element-wise ops) and emits ONNX (`.onnx`) model files using protobuf
5//! serialization via prost.
6
7use nxpu_analysis::{analyze, fusion};
8use nxpu_backend_core::{
9    Backend, BackendError, BackendOptions, BackendOutput, Diagnostic, DiagnosticLevel,
10    OutputContent, OutputFile,
11};
12use nxpu_ir::Module;
13use prost::Message;
14
15mod lower;
16#[doc(hidden)]
17pub mod proto;
18
19/// ONNX backend that compiles NxPU IR into `.onnx` model files.
20#[derive(Debug)]
21pub struct OnnxBackend;
22
23impl Backend for OnnxBackend {
24    fn name(&self) -> &str {
25        "ONNX"
26    }
27
28    fn targets(&self) -> &[&str] {
29        &["onnx"]
30    }
31
32    fn compile(
33        &self,
34        module: &Module,
35        opts: &BackendOptions,
36    ) -> Result<BackendOutput, BackendError> {
37        if module.entry_points.is_empty() {
38            return Err(BackendError::Other("no entry points in module".into()));
39        }
40
41        // 1. Classify all entry points.
42        let mut patterns = Vec::new();
43        for (i, ep) in module.entry_points.iter().enumerate() {
44            let pattern = analyze::classify_entry_point(module, i).map_err(|e| {
45                BackendError::Unsupported(format!("entry point '{}': {e}", ep.name))
46            })?;
47            if let analyze::KernelPattern::Unknown { reason } = &pattern {
48                return Err(BackendError::Unsupported(format!(
49                    "entry point '{}': unrecognized pattern: {reason}",
50                    ep.name
51                )));
52            }
53            patterns.push(pattern);
54        }
55
56        // 2. Extract embedded weights from module globals.
57        let weights = analyze::extract_embedded_weights(module);
58
59        // 3. Fuse adjacent patterns.
60        let fused = fusion::fuse_patterns(patterns);
61
62        // 4. Lower each fused pattern.
63        let mut files = Vec::new();
64        let mut diagnostics = Vec::new();
65
66        for (fp, ep_idx) in &fused {
67            let ep_name = if *ep_idx < module.entry_points.len() {
68                &module.entry_points[*ep_idx].name
69            } else {
70                "fused"
71            };
72
73            let summary = match fp {
74                fusion::FusedPattern::Single(p) => pattern_summary(p).to_string(),
75                fusion::FusedPattern::ConvBatchNorm { .. } => "Conv+BatchNorm (fused)".into(),
76                fusion::FusedPattern::WithActivation {
77                    base, activation, ..
78                } => {
79                    let base_name = match base.as_ref() {
80                        fusion::FusedPattern::Single(p) => pattern_summary(p).into_owned(),
81                        fusion::FusedPattern::ConvBatchNorm { .. } => "Conv+BatchNorm".to_string(),
82                        fusion::FusedPattern::MatMulBias { .. } => "Gemm".to_string(),
83                        _ => "fused".to_string(),
84                    };
85                    format!("{base_name}+{activation:?}")
86                }
87                fusion::FusedPattern::MatMulBias { .. } => "Gemm (fused)".into(),
88            };
89
90            diagnostics.push(Diagnostic {
91                level: DiagnosticLevel::Info,
92                message: format!("entry point '{ep_name}': classified as {summary}"),
93            });
94
95            let mut model = lower::build_fused_model(fp, ep_name, &weights)?;
96
97            // Inject per-channel QDQ nodes for weight tensors.
98            #[allow(clippy::collapsible_if)] // nested if-let for MSRV 1.87 compat (no let chains)
99            if !opts.per_channel_params.is_empty() {
100                if let Some(graph) = model.graph.as_mut() {
101                    lower::inject_per_channel_qdq(graph, &opts.per_channel_params);
102                }
103            }
104
105            // Embed quantization parameters in ONNX metadata_props.
106            for qp in &opts.quantization_params {
107                model.metadata_props.push(proto::StringStringEntryProto {
108                    key: format!("quant:{}:scale", qp.name),
109                    value: format!("{}", qp.scale),
110                });
111                model.metadata_props.push(proto::StringStringEntryProto {
112                    key: format!("quant:{}:zero_point", qp.name),
113                    value: format!("{}", qp.zero_point),
114                });
115            }
116
117            let bytes = model.encode_to_vec();
118
119            let filename = if fused.len() == 1 {
120                "output.onnx".into()
121            } else {
122                format!("{ep_name}.onnx")
123            };
124
125            files.push(OutputFile {
126                name: filename,
127                content: OutputContent::Binary(bytes),
128            });
129        }
130
131        Ok(BackendOutput { files, diagnostics })
132    }
133}
134
135/// Name the operators a pattern lowers to, for the `classified as …`
136/// diagnostic.
137///
138/// Owned rather than `&'static str` because a chain has no fixed name: it
139/// reports the operator sequence it emits — `Mul+Add`, `Cast+Mul+Mul` — since
140/// a category label would say nothing about what the kernel computes.
141fn pattern_summary(pattern: &analyze::KernelPattern) -> std::borrow::Cow<'static, str> {
142    use std::borrow::Cow;
143    Cow::Borrowed(match pattern {
144        analyze::KernelPattern::ElementWiseChain { cast, steps, .. } => {
145            return Cow::Owned(analyze::chain_summary(*cast, steps));
146        }
147        analyze::KernelPattern::MatMul { .. } => "MatMul",
148        // The whole graph, not its last node: there is no quantized matmul
149        // operator here, and "MatMul" alone would hide the dequantization
150        // that makes the answer right.
151        analyze::KernelPattern::QuantizedMatMul { bias, .. } => {
152            if bias.is_some() {
153                "Transpose+DequantizeLinear+MatMul+Add (int8 weights, per-channel scale)"
154            } else {
155                "Transpose+DequantizeLinear+MatMul (int8 weights, per-channel scale)"
156            }
157        }
158        analyze::KernelPattern::ElementWise { op, .. } => op.op_name(),
159        analyze::KernelPattern::Conv2D { .. } => "Conv",
160        analyze::KernelPattern::Pool { kind, .. } => kind.op_name(),
161        analyze::KernelPattern::Activation { op, .. } => op.op_name(),
162        analyze::KernelPattern::Reduce { op, .. } => op.op_name(),
163        analyze::KernelPattern::Transpose { .. } => "Transpose",
164        analyze::KernelPattern::Reshape { .. } => "Reshape",
165        analyze::KernelPattern::Normalization { .. } => "BatchNormalization",
166        analyze::KernelPattern::Concat { .. } => "Concat",
167        analyze::KernelPattern::Split { .. } => "Split",
168        analyze::KernelPattern::Attention { .. } => "Attention",
169        analyze::KernelPattern::Gather { .. } => "Gather",
170        analyze::KernelPattern::Scatter { .. } => "ScatterND",
171        analyze::KernelPattern::Unknown { .. } => "Unknown",
172    })
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178    use nxpu_backend_core::BackendOptions;
179
180    #[test]
181    fn backend_metadata() {
182        let backend = OnnxBackend;
183        assert_eq!(backend.name(), "ONNX");
184        assert!(backend.targets().contains(&"onnx"));
185    }
186
187    #[test]
188    fn compile_empty_module_fails() {
189        let backend = OnnxBackend;
190        let module = Module::default();
191        let result = backend.compile(&module, &BackendOptions::default());
192        assert!(result.is_err());
193    }
194
195    #[test]
196    fn compile_matmul_wgsl() {
197        let source = std::fs::read_to_string(concat!(
198            env!("CARGO_MANIFEST_DIR"),
199            "/../../examples/matmul.wgsl"
200        ))
201        .unwrap();
202        let module = nxpu_parser::parse(&source).unwrap();
203
204        let backend = OnnxBackend;
205        let output = backend
206            .compile(&module, &BackendOptions::default())
207            .unwrap();
208
209        assert_eq!(output.files.len(), 1);
210        assert_eq!(output.files[0].name, "output.onnx");
211
212        let bytes = match &output.files[0].content {
213            OutputContent::Binary(b) => b,
214            _ => panic!("expected binary output"),
215        };
216
217        // Decode and verify the ONNX model.
218        let model = proto::ModelProto::decode(bytes.as_slice()).unwrap();
219        assert_eq!(model.ir_version, 7);
220        assert_eq!(model.producer_name, "nxpu");
221        assert_eq!(model.opset_import[0].version, 13);
222
223        let graph = model.graph.as_ref().unwrap();
224        assert_eq!(graph.node.len(), 1);
225        assert_eq!(graph.node[0].op_type, "MatMul");
226        assert_eq!(graph.input.len(), 2);
227        assert_eq!(graph.output.len(), 1);
228    }
229
230    #[test]
231    fn compile_with_quantization_params_sets_metadata() {
232        let source = r#"
233@group(0) @binding(0) var<storage, read> a: array<f32>;
234@group(0) @binding(1) var<storage, read> b: array<f32>;
235@group(0) @binding(2) var<storage, read_write> c: array<f32>;
236
237struct Params { N: u32 }
238@group(0) @binding(3) var<uniform> params: Params;
239
240@compute @workgroup_size(256)
241fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
242  let idx = gid.x;
243  if (idx >= params.N) { return; }
244  c[idx] = a[idx] + b[idx];
245}
246"#;
247
248        let module = nxpu_parser::parse(source).unwrap();
249        let backend = OnnxBackend;
250        let opts = BackendOptions {
251            quantization_params: vec![
252                nxpu_backend_core::QuantParam {
253                    name: "x".into(),
254                    scale: 0.5,
255                    zero_point: 0,
256                },
257                nxpu_backend_core::QuantParam {
258                    name: "y".into(),
259                    scale: 0.25,
260                    zero_point: 128,
261                },
262            ],
263            ..Default::default()
264        };
265        let output = backend.compile(&module, &opts).unwrap();
266
267        assert_eq!(output.files.len(), 1);
268        let bytes = match &output.files[0].content {
269            OutputContent::Binary(b) => b,
270            _ => panic!("expected binary output"),
271        };
272
273        let model = proto::ModelProto::decode(bytes.as_slice()).unwrap();
274
275        // Should have 4 metadata_props: scale+zero_point for each of "x" and "y"
276        assert_eq!(model.metadata_props.len(), 4);
277
278        let keys: Vec<&str> = model
279            .metadata_props
280            .iter()
281            .map(|p| p.key.as_str())
282            .collect();
283        assert!(keys.contains(&"quant:x:scale"));
284        assert!(keys.contains(&"quant:x:zero_point"));
285        assert!(keys.contains(&"quant:y:scale"));
286        assert!(keys.contains(&"quant:y:zero_point"));
287
288        // Verify values
289        let x_scale = model
290            .metadata_props
291            .iter()
292            .find(|p| p.key == "quant:x:scale")
293            .unwrap();
294        assert_eq!(x_scale.value, "0.5");
295
296        let y_zp = model
297            .metadata_props
298            .iter()
299            .find(|p| p.key == "quant:y:zero_point")
300            .unwrap();
301        assert_eq!(y_zp.value, "128");
302    }
303
304    #[test]
305    fn compile_vecadd_wgsl() {
306        let source = r#"
307@group(0) @binding(0) var<storage, read> a: array<f32>;
308@group(0) @binding(1) var<storage, read> b: array<f32>;
309@group(0) @binding(2) var<storage, read_write> c: array<f32>;
310
311struct Params { N: u32 }
312@group(0) @binding(3) var<uniform> params: Params;
313
314@compute @workgroup_size(256)
315fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
316  let idx = gid.x;
317  if (idx >= params.N) { return; }
318  c[idx] = a[idx] + b[idx];
319}
320"#;
321
322        let module = nxpu_parser::parse(source).unwrap();
323
324        let backend = OnnxBackend;
325        let output = backend
326            .compile(&module, &BackendOptions::default())
327            .unwrap();
328
329        assert_eq!(output.files.len(), 1);
330
331        let bytes = match &output.files[0].content {
332            OutputContent::Binary(b) => b,
333            _ => panic!("expected binary output"),
334        };
335
336        let model = proto::ModelProto::decode(bytes.as_slice()).unwrap();
337        let graph = model.graph.as_ref().unwrap();
338        assert_eq!(graph.node.len(), 1);
339        assert_eq!(graph.node[0].op_type, "Add");
340        assert_eq!(graph.input.len(), 2);
341        assert_eq!(graph.output.len(), 1);
342    }
343
344    #[test]
345    fn compile_with_per_channel_params_injects_qdq() {
346        let source = r#"
347@group(0) @binding(0) var<storage, read> a: array<f32>;
348@group(0) @binding(1) var<storage, read> b: array<f32>;
349@group(0) @binding(2) var<storage, read_write> c: array<f32>;
350
351struct Params { N: u32 }
352@group(0) @binding(3) var<uniform> params: Params;
353
354@compute @workgroup_size(256)
355fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
356  let idx = gid.x;
357  if (idx >= params.N) { return; }
358  c[idx] = a[idx] + b[idx];
359}
360"#;
361
362        let module = nxpu_parser::parse(source).unwrap();
363        let backend = OnnxBackend;
364        let opts = BackendOptions {
365            per_channel_params: vec![nxpu_backend_core::PerChannelParam {
366                name: "a".into(),
367                scales: vec![0.1, 0.2],
368                zero_points: vec![0, 0],
369                channel_axis: 0,
370            }],
371            ..Default::default()
372        };
373        let output = backend.compile(&module, &opts).unwrap();
374
375        assert_eq!(output.files.len(), 1);
376        let bytes = match &output.files[0].content {
377            OutputContent::Binary(b) => b,
378            _ => panic!("expected binary output"),
379        };
380
381        let model = proto::ModelProto::decode(bytes.as_slice()).unwrap();
382        let graph = model.graph.as_ref().unwrap();
383
384        // Should have 3 nodes: QuantizeLinear, DequantizeLinear, Add
385        assert_eq!(graph.node.len(), 3);
386        assert_eq!(graph.node[0].op_type, "QuantizeLinear");
387        assert_eq!(graph.node[1].op_type, "DequantizeLinear");
388        assert_eq!(graph.node[2].op_type, "Add");
389
390        // Scale initializer should be present.
391        let scale_init = graph
392            .initializer
393            .iter()
394            .find(|i| i.name == "a_scale")
395            .expect("expected a_scale initializer");
396        assert_eq!(scale_init.float_data, vec![0.1, 0.2]);
397    }
398}
nxpu_backend_onnx/lib.rs

nxpu_backend_onnx/
lib.rs
