nxpu_backend_mediatek/
lib.rs

1//! MediaTek APU backend for NxPU.
2//!
3//! Delegates compilation to the TFLite backend with MediaTek APU-specific
4//! validation and NeuroPilot compilation hints.
5
6use nxpu_analysis::analyze;
7use nxpu_backend_core::{
8    Backend, BackendError, BackendOptions, BackendOutput, Diagnostic, DiagnosticLevel, Precision,
9    PrecisionPolicy, validate_patterns,
10};
11use nxpu_backend_tflite::TfLiteBackend;
12use nxpu_ir::Module;
13
14mod support;
15
16use support::MediaTekApuSupport;
17
18/// MediaTek APU backend with NeuroPilot compilation hints.
19#[derive(Debug)]
20pub struct MediaTekBackend;
21
22impl Backend for MediaTekBackend {
23    fn name(&self) -> &str {
24        "MediaTek APU"
25    }
26
27    fn targets(&self) -> &[&str] {
28        &["mediatek", "neuropilot"]
29    }
30
31    fn preferred_precision(&self) -> Precision {
32        Precision::Int8
33    }
34
35    fn compile(
36        &self,
37        module: &Module,
38        opts: &BackendOptions,
39    ) -> Result<BackendOutput, BackendError> {
40        // 1. Classify entry points and validate.
41        let mut op_names = Vec::new();
42        for (i, ep) in module.entry_points.iter().enumerate() {
43            match analyze::classify_entry_point(module, i) {
44                Ok(pattern) => {
45                    op_names.extend(analyze::pattern_op_names(&pattern));
46                }
47                Err(e) => {
48                    return Err(BackendError::Unsupported(format!(
49                        "entry point '{}': {e}",
50                        ep.name
51                    )));
52                }
53            }
54        }
55
56        let precision = resolve_precision(opts, self.preferred_precision());
57        let op_refs: Vec<&str> = op_names.iter().map(|s| s.as_str()).collect();
58        let mut diagnostics = validate_patterns(&MediaTekApuSupport, &op_refs, precision);
59
60        // 2. Generate TFLite model.
61        let mut output = TfLiteBackend.compile(module, opts)?;
62
63        diagnostics.extend(output.diagnostics);
64
65        // 3. NeuroPilot compilation hints.
66        for file in &output.files {
67            diagnostics.push(Diagnostic {
68                level: DiagnosticLevel::Info,
69                message: format!("To compile for MediaTek APU: ncc-tflite {}", file.name),
70            });
71        }
72        diagnostics.push(Diagnostic {
73            level: DiagnosticLevel::Info,
74            message: "NeuroPilot SDK: ncc-tflite --arch mdla3.0 output.tflite -o output.dla".into(),
75        });
76        diagnostics.push(Diagnostic {
77            level: DiagnosticLevel::Info,
78            message: "For quantization: use TFLite quantization-aware training or \
79                      post-training quantization"
80                .into(),
81        });
82
83        output.diagnostics = diagnostics;
84        Ok(output)
85    }
86}
87
88fn resolve_precision(opts: &BackendOptions, preferred: Precision) -> Precision {
89    match opts.precision {
90        PrecisionPolicy::Explicit(p) => p,
91        PrecisionPolicy::Auto => preferred,
92        PrecisionPolicy::Keep => Precision::F32,
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use nxpu_backend_core::{BackendOptions, OutputContent};
100
101    #[test]
102    fn backend_metadata() {
103        let backend = MediaTekBackend;
104        assert_eq!(backend.name(), "MediaTek APU");
105        assert!(backend.targets().contains(&"mediatek"));
106        assert!(backend.targets().contains(&"neuropilot"));
107        assert_eq!(backend.preferred_precision(), Precision::Int8);
108    }
109
110    #[test]
111    fn compile_matmul_with_neuropilot_hints() {
112        let source = std::fs::read_to_string(concat!(
113            env!("CARGO_MANIFEST_DIR"),
114            "/../../examples/matmul.wgsl"
115        ))
116        .unwrap();
117        let module = nxpu_parser::parse(&source).unwrap();
118
119        let output = MediaTekBackend
120            .compile(&module, &BackendOptions::default())
121            .unwrap();
122        assert_eq!(output.files.len(), 1);
123        assert_eq!(output.files[0].name, "output.tflite");
124        assert!(matches!(output.files[0].content, OutputContent::Binary(_)));
125
126        let messages: Vec<&str> = output
127            .diagnostics
128            .iter()
129            .map(|d| d.message.as_str())
130            .collect();
131        assert!(messages.iter().any(|m| m.contains("ncc-tflite")));
132    }
133
134    fn load_and_compile(example: &str, opts: &BackendOptions) -> BackendOutput {
135        let source = std::fs::read_to_string(format!(
136            "{}/../../examples/{example}.wgsl",
137            env!("CARGO_MANIFEST_DIR")
138        ))
139        .unwrap();
140        let module = nxpu_parser::parse(&source).unwrap();
141        MediaTekBackend.compile(&module, opts).unwrap()
142    }
143
144    #[test]
145    fn compile_conv2d() {
146        let output = load_and_compile("conv2d", &BackendOptions::default());
147        assert_ne!(output.files.len(), 0);
148        let has_tflite = output.files.iter().any(|f| f.name.ends_with(".tflite"));
149        assert!(has_tflite);
150    }
151
152    #[test]
153    fn compile_relu() {
154        let output = load_and_compile("relu", &BackendOptions::default());
155        assert_ne!(output.files.len(), 0);
156    }
157
158    #[test]
159    fn compile_attention() {
160        let output = load_and_compile("attention", &BackendOptions::default());
161        assert_ne!(output.files.len(), 0);
162    }
163
164    #[test]
165    fn resolve_precision_explicit_and_keep() {
166        let explicit_opts = BackendOptions {
167            precision: PrecisionPolicy::Explicit(Precision::F16),
168            ..BackendOptions::default()
169        };
170        assert_eq!(
171            resolve_precision(&explicit_opts, Precision::Int8),
172            Precision::F16
173        );
174
175        let keep_opts = BackendOptions {
176            precision: PrecisionPolicy::Keep,
177            ..BackendOptions::default()
178        };
179        assert_eq!(
180            resolve_precision(&keep_opts, Precision::Int8),
181            Precision::F32
182        );
183    }
184
185    #[test]
186    fn neuropilot_arch_diagnostic() {
187        let output = load_and_compile("matmul", &BackendOptions::default());
188        let messages: Vec<&str> = output
189            .diagnostics
190            .iter()
191            .map(|d| d.message.as_str())
192            .collect();
193        assert!(messages.iter().any(|m| m.contains("mdla3.0")));
194    }
195}
nxpu_backend_mediatek/lib.rs

nxpu_backend_mediatek/
lib.rs
