nxpu_backend_arm_ethos/
lib.rs

1//! Arm Ethos NPU backend for NxPU.
2//!
3//! Compiles NxPU IR via the TFLite backend, then optionally invokes the
4//! Arm Vela compiler to produce Ethos-U optimized `.tflite` binaries.
5//! Validates operator patterns against the Ethos-U55/U65 support matrix
6//! and emits Vela accelerator configuration hints.
7//!
8//! When `vela` is not found on `$PATH`, the backend falls back to emitting
9//! a standard `.tflite` file with a diagnostic hint.
10
11use std::process::Command;
12
13use nxpu_analysis::analyze;
14use nxpu_backend_core::{
15    Backend, BackendError, BackendOptions, BackendOutput, Diagnostic, DiagnosticLevel,
16    OutputContent, OutputFile, Precision, PrecisionPolicy, validate_patterns,
17};
18use nxpu_backend_tflite::TfLiteBackend;
19use nxpu_ir::Module;
20
21pub mod support;
22
23use support::EthosU55Support;
24
25/// Arm Ethos NPU backend.
26///
27/// Compilation pipeline:
28/// 1. Classify and validate patterns against Ethos-U support matrix.
29/// 2. Lower IR → TFLite FlatBuffer via [`TfLiteBackend`].
30/// 3. If `vela` is available, invoke it on the TFLite file to produce
31///    an Ethos-U optimized binary.
32/// 4. Otherwise, emit the unoptimized TFLite with a diagnostic.
33#[derive(Debug)]
34pub struct ArmEthosBackend;
35
36/// Check whether the `vela` CLI tool is available on PATH.
37fn vela_available() -> bool {
38    Command::new("vela")
39        .arg("--version")
40        .output()
41        .is_ok_and(|o| o.status.success())
42}
43
44/// Run the Vela compiler on a TFLite file, returning the optimized bytes.
45///
46/// Uses a unique temporary directory per invocation to avoid race conditions
47/// during parallel compilation.
48fn run_vela(tflite_bytes: &[u8]) -> Result<Vec<u8>, BackendError> {
49    let temp_dir = std::env::temp_dir().join(format!("nxpu-vela-{}", std::process::id()));
50    std::fs::create_dir_all(&temp_dir)
51        .map_err(|e| BackendError::Other(format!("failed to create temp dir: {e}")))?;
52
53    // Use a closure to ensure cleanup on all exit paths.
54    let result = (|| {
55        let input_name = "input";
56        let input_path = temp_dir.join(format!("{input_name}.tflite"));
57        std::fs::write(&input_path, tflite_bytes)
58            .map_err(|e| BackendError::Other(format!("failed to write temp tflite: {e}")))?;
59
60        let output = Command::new("vela")
61            .arg(&input_path)
62            .arg("--output-dir")
63            .arg(&temp_dir)
64            .output()
65            .map_err(|e| BackendError::Other(format!("failed to run vela: {e}")))?;
66
67        if !output.status.success() {
68            let stderr = String::from_utf8_lossy(&output.stderr);
69            return Err(BackendError::Other(format!("vela failed: {stderr}")));
70        }
71
72        // Vela outputs to <output-dir>/<input_name>_vela.tflite
73        let vela_output = temp_dir.join(format!("{input_name}_vela.tflite"));
74        let optimized = std::fs::read(&vela_output)
75            .map_err(|e| BackendError::Other(format!("failed to read vela output: {e}")))?;
76
77        Ok(optimized)
78    })();
79
80    // Always clean up temp directory, regardless of success or failure.
81    let _ = std::fs::remove_dir_all(&temp_dir);
82
83    result
84}
85
86impl Backend for ArmEthosBackend {
87    fn name(&self) -> &str {
88        "Arm Ethos NPU"
89    }
90
91    fn targets(&self) -> &[&str] {
92        &["arm-ethos", "ethos-u"]
93    }
94
95    fn preferred_precision(&self) -> Precision {
96        Precision::Int8
97    }
98
99    fn compile(
100        &self,
101        module: &Module,
102        opts: &BackendOptions,
103    ) -> Result<BackendOutput, BackendError> {
104        // 1. Classify entry points and validate against Ethos-U support matrix.
105        let mut op_names = Vec::new();
106        for (i, ep) in module.entry_points.iter().enumerate() {
107            match analyze::classify_entry_point(module, i) {
108                Ok(pattern) => {
109                    op_names.extend(analyze::pattern_op_names(&pattern));
110                }
111                Err(e) => {
112                    return Err(BackendError::Unsupported(format!(
113                        "entry point '{}': {e}",
114                        ep.name
115                    )));
116                }
117            }
118        }
119
120        let precision = resolve_precision(opts, self.preferred_precision());
121        let op_refs: Vec<&str> = op_names.iter().map(|s| s.as_str()).collect();
122        let mut diagnostics = validate_patterns(&EthosU55Support, &op_refs, precision);
123
124        // 2. Generate TFLite model
125        let tflite_output = TfLiteBackend.compile(module, opts)?;
126        diagnostics.extend(tflite_output.diagnostics);
127
128        let mut files = Vec::new();
129
130        // Check Vela availability once, not per file.
131        let has_vela = vela_available();
132
133        for file in &tflite_output.files {
134            let tflite_bytes = match &file.content {
135                OutputContent::Binary(b) => b,
136                OutputContent::Text(_) => {
137                    files.push(file.clone());
138                    continue;
139                }
140            };
141
142            // 3. Try to invoke Vela
143            if has_vela {
144                match run_vela(tflite_bytes) {
145                    Ok(optimized) => {
146                        diagnostics.push(Diagnostic {
147                            level: DiagnosticLevel::Info,
148                            message: format!(
149                                "Vela compilation successful ({} -> {} bytes)",
150                                tflite_bytes.len(),
151                                optimized.len()
152                            ),
153                        });
154                        files.push(OutputFile {
155                            name: file.name.replace(".tflite", "_vela.tflite"),
156                            content: OutputContent::Binary(optimized),
157                        });
158                        // Also emit original for reference
159                        files.push(file.clone());
160                    }
161                    Err(e) => {
162                        diagnostics.push(Diagnostic {
163                            level: DiagnosticLevel::Warning,
164                            message: format!(
165                                "Vela compilation failed, emitting unoptimized TFLite: {e}"
166                            ),
167                        });
168                        files.push(file.clone());
169                    }
170                }
171            } else {
172                // Vela not available — fall back with hint
173                diagnostics.push(Diagnostic {
174                    level: DiagnosticLevel::Info,
175                    message: "vela not found on PATH; emitting unoptimized TFLite. \
176                              Install: pip install ethos-u-vela"
177                        .into(),
178                });
179                diagnostics.push(Diagnostic {
180                    level: DiagnosticLevel::Info,
181                    message: format!(
182                        "To optimize for Ethos-U55: vela {} --accelerator-config ethos-u55-128",
183                        file.name
184                    ),
185                });
186                diagnostics.push(Diagnostic {
187                    level: DiagnosticLevel::Info,
188                    message: format!(
189                        "To optimize for Ethos-U65: vela {} --accelerator-config ethos-u65-512",
190                        file.name
191                    ),
192                });
193                files.push(file.clone());
194            }
195        }
196
197        // Quantization calibration hint
198        diagnostics.push(Diagnostic {
199            level: DiagnosticLevel::Info,
200            message: "For Int8 quantization: use TFLite post-training quantization with \
201                      representative dataset"
202                .into(),
203        });
204
205        Ok(BackendOutput { files, diagnostics })
206    }
207}
208
209fn resolve_precision(opts: &BackendOptions, preferred: Precision) -> Precision {
210    match opts.precision {
211        PrecisionPolicy::Explicit(p) => p,
212        PrecisionPolicy::Auto => preferred,
213        PrecisionPolicy::Keep => Precision::F32,
214    }
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220    use nxpu_backend_core::BackendOptions;
221
222    #[test]
223    fn backend_metadata() {
224        let backend = ArmEthosBackend;
225        assert_eq!(backend.name(), "Arm Ethos NPU");
226        assert!(backend.targets().contains(&"arm-ethos"));
227        assert!(backend.targets().contains(&"ethos-u"));
228        assert_eq!(backend.preferred_precision(), Precision::Int8);
229    }
230
231    #[test]
232    fn compile_produces_tflite_with_validation() {
233        let source = std::fs::read_to_string(concat!(
234            env!("CARGO_MANIFEST_DIR"),
235            "/../../examples/matmul.wgsl"
236        ))
237        .unwrap();
238        let module = nxpu_parser::parse(&source).unwrap();
239
240        let output = ArmEthosBackend
241            .compile(&module, &BackendOptions::default())
242            .unwrap();
243
244        // Should have at least one output file
245        assert_ne!(output.files.len(), 0);
246
247        // At least one file should be a .tflite
248        let has_tflite = output.files.iter().any(|f| f.name.ends_with(".tflite"));
249        assert!(has_tflite);
250
251        // Should have diagnostics about Ethos-U and vela
252        assert_ne!(output.diagnostics.len(), 0);
253        let messages: Vec<&str> = output
254            .diagnostics
255            .iter()
256            .map(|d| d.message.as_str())
257            .collect();
258        assert!(
259            messages
260                .iter()
261                .any(|m| m.contains("ethos-u") || m.contains("Ethos"))
262        );
263    }
264
265    #[test]
266    fn vela_availability_check() {
267        // This test just verifies the function doesn't panic.
268        // On most CI systems vela won't be installed.
269        let _available = vela_available();
270    }
271
272    fn load_and_compile(example: &str, opts: &BackendOptions) -> BackendOutput {
273        let source = std::fs::read_to_string(format!(
274            "{}/../../examples/{example}.wgsl",
275            env!("CARGO_MANIFEST_DIR")
276        ))
277        .unwrap();
278        let module = nxpu_parser::parse(&source).unwrap();
279        ArmEthosBackend.compile(&module, opts).unwrap()
280    }
281
282    #[test]
283    fn compile_conv2d() {
284        let output = load_and_compile("conv2d", &BackendOptions::default());
285        assert_ne!(output.files.len(), 0);
286        let has_tflite = output.files.iter().any(|f| f.name.ends_with(".tflite"));
287        assert!(has_tflite);
288    }
289
290    #[test]
291    fn compile_relu() {
292        let output = load_and_compile("relu", &BackendOptions::default());
293        assert_ne!(output.files.len(), 0);
294    }
295
296    #[test]
297    fn compile_attention() {
298        let output = load_and_compile("attention", &BackendOptions::default());
299        assert_ne!(output.files.len(), 0);
300    }
301
302    #[test]
303    fn resolve_precision_explicit_and_keep() {
304        let explicit_opts = BackendOptions {
305            precision: PrecisionPolicy::Explicit(Precision::F16),
306            ..BackendOptions::default()
307        };
308        assert_eq!(
309            resolve_precision(&explicit_opts, Precision::Int8),
310            Precision::F16
311        );
312
313        let keep_opts = BackendOptions {
314            precision: PrecisionPolicy::Keep,
315            ..BackendOptions::default()
316        };
317        assert_eq!(
318            resolve_precision(&keep_opts, Precision::Int8),
319            Precision::F32
320        );
321    }
322}
nxpu_backend_arm_ethos/lib.rs

nxpu_backend_arm_ethos/
lib.rs
