nxpu_backend_amd/
lib.rs

1//! AMD XDNA NPU backend for NxPU.
2//!
3//! Delegates compilation to the ONNX backend with AMD XDNA-specific metadata
4//! properties (target device, execution provider, quantization scheme).
5//! Validates operator patterns against the XDNA support matrix and emits
6//! Vitis AI conversion hints.
7
8use nxpu_analysis::analyze;
9use nxpu_backend_core::{
10    Backend, BackendError, BackendOptions, BackendOutput, Diagnostic, DiagnosticLevel,
11    OutputContent, Precision, PrecisionPolicy, validate_patterns,
12};
13use nxpu_backend_onnx::OnnxBackend;
14use nxpu_backend_onnx::proto::{ModelProto, StringStringEntryProto};
15use nxpu_ir::Module;
16use prost::Message;
17
18mod support;
19
20use support::AmdXdnaSupport;
21
22/// AMD XDNA NPU backend with enhanced ONNX metadata.
23#[derive(Debug)]
24pub struct AmdBackend;
25
26impl Backend for AmdBackend {
27    fn name(&self) -> &str {
28        "AMD XDNA NPU"
29    }
30
31    fn targets(&self) -> &[&str] {
32        &["amd-xdna", "amd-npu"]
33    }
34
35    fn preferred_precision(&self) -> Precision {
36        Precision::Int8
37    }
38
39    fn compile(
40        &self,
41        module: &Module,
42        opts: &BackendOptions,
43    ) -> Result<BackendOutput, BackendError> {
44        // 1. Classify entry points for validation.
45        let mut op_names = Vec::new();
46        for (i, ep) in module.entry_points.iter().enumerate() {
47            match analyze::classify_entry_point(module, i) {
48                Ok(pattern) => {
49                    op_names.extend(analyze::pattern_op_names(&pattern));
50                }
51                Err(e) => {
52                    return Err(BackendError::Unsupported(format!(
53                        "entry point '{}': {e}",
54                        ep.name
55                    )));
56                }
57            }
58        }
59
60        // 2. Validate against AMD XDNA support matrix.
61        let precision = resolve_precision(opts, self.preferred_precision());
62        let op_refs: Vec<&str> = op_names.iter().map(|s| s.as_str()).collect();
63        let mut diagnostics = validate_patterns(&AmdXdnaSupport, &op_refs, precision);
64
65        // 3. Generate ONNX model via base backend.
66        let mut output = OnnxBackend.compile(module, opts)?;
67
68        // 4. Inject XDNA metadata_props into each ONNX model file.
69        let metadata_props = vec![
70            StringStringEntryProto {
71                key: "xdna:target_device".into(),
72                value: "AMD Ryzen AI".into(),
73            },
74            StringStringEntryProto {
75                key: "xdna:execution_provider".into(),
76                value: "VitisAIExecutionProvider".into(),
77            },
78            StringStringEntryProto {
79                key: "xdna:quantization".into(),
80                value: format!("{precision}"),
81            },
82        ];
83
84        for file in &mut output.files {
85            let is_onnx = std::path::Path::new(&file.name)
86                .extension()
87                .is_some_and(|ext| ext.eq_ignore_ascii_case("onnx"));
88            if !is_onnx {
89                continue;
90            }
91            let OutputContent::Binary(bytes) = &file.content else {
92                continue;
93            };
94            let Ok(mut model) = ModelProto::decode(bytes.as_slice()) else {
95                continue;
96            };
97            model.metadata_props.extend(metadata_props.clone());
98            file.content = OutputContent::Binary(model.encode_to_vec());
99        }
100
101        diagnostics.extend(output.diagnostics);
102        diagnostics.push(Diagnostic {
103            level: DiagnosticLevel::Info,
104            message: "To compile for XDNA: use Vitis AI EP with ONNX Runtime".into(),
105        });
106        diagnostics.push(Diagnostic {
107            level: DiagnosticLevel::Info,
108            message:
109                "Vitis AI quantization: vai_q_onnx quantize_static --model output.onnx --calibration_data_reader <reader>"
110                    .into(),
111        });
112
113        output.diagnostics = diagnostics;
114        Ok(output)
115    }
116}
117
118fn resolve_precision(opts: &BackendOptions, preferred: Precision) -> Precision {
119    match opts.precision {
120        PrecisionPolicy::Explicit(p) => p,
121        PrecisionPolicy::Auto => preferred,
122        PrecisionPolicy::Keep => Precision::F32,
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129    use nxpu_backend_core::{BackendOptions, OutputContent};
130
131    #[test]
132    fn backend_metadata() {
133        let backend = AmdBackend;
134        assert_eq!(backend.name(), "AMD XDNA NPU");
135        assert!(backend.targets().contains(&"amd-xdna"));
136        assert!(backend.targets().contains(&"amd-npu"));
137        assert_eq!(backend.preferred_precision(), Precision::Int8);
138    }
139
140    #[test]
141    fn compile_matmul_has_metadata_props() {
142        let source = std::fs::read_to_string(concat!(
143            env!("CARGO_MANIFEST_DIR"),
144            "/../../examples/matmul.wgsl"
145        ))
146        .unwrap();
147        let module = nxpu_parser::parse(&source).unwrap();
148
149        let output = AmdBackend
150            .compile(&module, &BackendOptions::default())
151            .unwrap();
152        assert_eq!(output.files.len(), 1);
153        assert_eq!(output.files[0].name, "output.onnx");
154
155        // Decode and verify metadata_props
156        let bytes = match &output.files[0].content {
157            OutputContent::Binary(b) => b,
158            _ => panic!("expected binary"),
159        };
160        let model = ModelProto::decode(bytes.as_slice()).unwrap();
161        assert_ne!(model.metadata_props.len(), 0);
162
163        let keys: Vec<&str> = model
164            .metadata_props
165            .iter()
166            .map(|p| p.key.as_str())
167            .collect();
168        assert!(keys.contains(&"xdna:target_device"));
169        assert!(keys.contains(&"xdna:execution_provider"));
170        assert!(keys.contains(&"xdna:quantization"));
171    }
172
173    #[test]
174    fn diagnostics_include_vitis_hint() {
175        let source = std::fs::read_to_string(concat!(
176            env!("CARGO_MANIFEST_DIR"),
177            "/../../examples/matmul.wgsl"
178        ))
179        .unwrap();
180        let module = nxpu_parser::parse(&source).unwrap();
181
182        let output = AmdBackend
183            .compile(&module, &BackendOptions::default())
184            .unwrap();
185        let messages: Vec<&str> = output
186            .diagnostics
187            .iter()
188            .map(|d| d.message.as_str())
189            .collect();
190        assert!(messages.iter().any(|m| m.contains("Vitis AI")));
191    }
192
193    fn load_and_compile(example: &str, opts: &BackendOptions) -> BackendOutput {
194        let source = std::fs::read_to_string(format!(
195            "{}/../../examples/{example}.wgsl",
196            env!("CARGO_MANIFEST_DIR")
197        ))
198        .unwrap();
199        let module = nxpu_parser::parse(&source).unwrap();
200        AmdBackend.compile(&module, opts).unwrap()
201    }
202
203    #[test]
204    fn compile_conv2d() {
205        let output = load_and_compile("conv2d", &BackendOptions::default());
206        assert_ne!(output.files.len(), 0);
207        for file in &output.files {
208            assert_ne!(file.content.len(), 0);
209        }
210    }
211
212    #[test]
213    fn compile_relu() {
214        let output = load_and_compile("relu", &BackendOptions::default());
215        assert_ne!(output.files.len(), 0);
216        for file in &output.files {
217            assert_ne!(file.content.len(), 0);
218        }
219    }
220
221    #[test]
222    fn compile_attention() {
223        let output = load_and_compile("attention", &BackendOptions::default());
224        assert_ne!(output.files.len(), 0);
225        for file in &output.files {
226            assert_ne!(file.content.len(), 0);
227        }
228    }
229
230    #[test]
231    fn resolve_precision_explicit_and_keep() {
232        let explicit_opts = BackendOptions {
233            precision: PrecisionPolicy::Explicit(Precision::F16),
234            ..BackendOptions::default()
235        };
236        assert_eq!(
237            resolve_precision(&explicit_opts, Precision::Int8),
238            Precision::F16
239        );
240
241        let keep_opts = BackendOptions {
242            precision: PrecisionPolicy::Keep,
243            ..BackendOptions::default()
244        };
245        assert_eq!(
246            resolve_precision(&keep_opts, Precision::Int8),
247            Precision::F32
248        );
249    }
250
251    #[test]
252    fn metadata_reflects_explicit_precision() {
253        let opts = BackendOptions {
254            precision: PrecisionPolicy::Explicit(Precision::F16),
255            ..BackendOptions::default()
256        };
257        let output = load_and_compile("matmul", &opts);
258
259        let bytes = match &output.files[0].content {
260            OutputContent::Binary(b) => b,
261            _ => panic!("expected binary"),
262        };
263        let model = ModelProto::decode(bytes.as_slice()).unwrap();
264
265        let quant_prop = model
266            .metadata_props
267            .iter()
268            .find(|p| p.key == "xdna:quantization")
269            .expect("missing xdna:quantization metadata prop");
270        assert!(
271            quant_prop.value.to_lowercase().contains("f16"),
272            "expected quantization value to contain 'f16', got: {}",
273            quant_prop.value
274        );
275    }
276}
nxpu_backend_amd/lib.rs

nxpu_backend_amd/
lib.rs
