nxpu_backend_coreml/
lib.rs

1//! CoreML backend emitter for NxPU (Apple ANE).
2//!
3//! Emits CoreML ML Program (`.mlmodel`) protobuf from NxPU IR.
4//! The Apple Neural Engine operates at FP16 precision.
5
6use nxpu_analysis::analyze;
7use nxpu_backend_core::{
8    Backend, BackendError, BackendOptions, BackendOutput, Diagnostic, DiagnosticLevel,
9    OutputContent, OutputFile, Precision,
10};
11use nxpu_ir::Module;
12use prost::Message;
13
14mod lower;
15#[doc(hidden)]
16pub mod proto;
17
18/// CoreML backend targeting Apple Neural Engine.
19#[derive(Debug)]
20pub struct CoreMlBackend;
21
22impl Backend for CoreMlBackend {
23    fn name(&self) -> &str {
24        "CoreML"
25    }
26
27    fn targets(&self) -> &[&str] {
28        &["coreml", "apple-ane"]
29    }
30
31    fn preferred_precision(&self) -> Precision {
32        Precision::F16
33    }
34
35    fn compile(
36        &self,
37        module: &Module,
38        _opts: &BackendOptions,
39    ) -> Result<BackendOutput, BackendError> {
40        if module.entry_points.is_empty() {
41            return Err(BackendError::Other("no entry points in module".into()));
42        }
43
44        let mut files = Vec::new();
45        let mut diagnostics = Vec::new();
46
47        for (i, ep) in module.entry_points.iter().enumerate() {
48            let pattern = analyze::classify_entry_point(module, i).map_err(|e| {
49                BackendError::Unsupported(format!("entry point '{}': {e}", ep.name))
50            })?;
51            if let analyze::KernelPattern::Unknown { reason } = &pattern {
52                return Err(BackendError::Unsupported(format!(
53                    "entry point '{}': unrecognized pattern: {reason}",
54                    ep.name
55                )));
56            }
57
58            let summary = match &pattern {
59                analyze::KernelPattern::MatMul { .. } => "matmul",
60                analyze::KernelPattern::QuantizedMatMul { .. } => "quantized matmul",
61                analyze::KernelPattern::ElementWise { op, .. } => op.op_name(),
62                analyze::KernelPattern::Conv2D { .. } => "conv",
63                analyze::KernelPattern::Pool { kind, .. } => match kind {
64                    analyze::PoolKind::Max => "max_pool",
65                    analyze::PoolKind::Avg => "avg_pool",
66                },
67                analyze::KernelPattern::Activation { op, .. } => op.op_name(),
68                analyze::KernelPattern::Reduce { op, .. } => op.op_name(),
69                analyze::KernelPattern::Transpose { .. } => "transpose",
70                analyze::KernelPattern::Reshape { .. } => "reshape",
71                analyze::KernelPattern::Normalization { .. } => "batch_norm",
72                analyze::KernelPattern::Concat { .. } => "concat",
73                analyze::KernelPattern::Split { .. } => "split",
74                analyze::KernelPattern::Attention { .. } => "scaled_dot_product_attention",
75                analyze::KernelPattern::Gather { .. } => "gather",
76                analyze::KernelPattern::Scatter { .. } => "scatter",
77                analyze::KernelPattern::ElementWiseChain { .. } => "element-wise chain",
78                analyze::KernelPattern::Unknown { .. } => "Unknown",
79            };
80
81            diagnostics.push(Diagnostic {
82                level: DiagnosticLevel::Info,
83                message: format!("entry point '{}': MIL op {}", ep.name, summary),
84            });
85
86            let model = lower::build_model(&pattern, &ep.name)?;
87            let bytes = model.encode_to_vec();
88
89            let filename = if module.entry_points.len() == 1 {
90                "output.mlmodel".into()
91            } else {
92                format!("{}.mlmodel", ep.name)
93            };
94
95            files.push(OutputFile {
96                name: filename,
97                content: OutputContent::Binary(bytes),
98            });
99        }
100
101        Ok(BackendOutput { files, diagnostics })
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108    use nxpu_backend_core::BackendOptions;
109
110    #[test]
111    fn backend_metadata() {
112        let backend = CoreMlBackend;
113        assert_eq!(backend.name(), "CoreML");
114        assert!(backend.targets().contains(&"coreml"));
115        assert!(backend.targets().contains(&"apple-ane"));
116    }
117
118    #[test]
119    fn compile_matmul_wgsl() {
120        let source = std::fs::read_to_string(concat!(
121            env!("CARGO_MANIFEST_DIR"),
122            "/../../examples/matmul.wgsl"
123        ))
124        .unwrap();
125        let module = nxpu_parser::parse(&source).unwrap();
126
127        let output = CoreMlBackend
128            .compile(&module, &BackendOptions::default())
129            .unwrap();
130        assert_eq!(output.files.len(), 1);
131        assert_eq!(output.files[0].name, "output.mlmodel");
132
133        let bytes = match &output.files[0].content {
134            OutputContent::Binary(b) => b,
135            _ => panic!("expected binary output"),
136        };
137
138        let model = proto::Model::decode(bytes.as_slice()).unwrap();
139        assert_eq!(model.specification_version, proto::SPECIFICATION_VERSION);
140        let proto::model::Type::MlProgram(prog) = model.r#type.as_ref().unwrap();
141        assert_eq!(
142            prog.functions[0].block.as_ref().unwrap().operations[0].r#type,
143            "matmul"
144        );
145    }
146
147    #[test]
148    fn compile_vecadd_wgsl() {
149        let source = r#"
150@group(0) @binding(0) var<storage, read> a: array<f32>;
151@group(0) @binding(1) var<storage, read> b: array<f32>;
152@group(0) @binding(2) var<storage, read_write> c: array<f32>;
153struct Params { N: u32 }
154@group(0) @binding(3) var<uniform> params: Params;
155@compute @workgroup_size(256)
156fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
157  let idx = gid.x;
158  if (idx >= params.N) { return; }
159  c[idx] = a[idx] + b[idx];
160}
161"#;
162        let module = nxpu_parser::parse(source).unwrap();
163        let output = CoreMlBackend
164            .compile(&module, &BackendOptions::default())
165            .unwrap();
166
167        let bytes = match &output.files[0].content {
168            OutputContent::Binary(b) => b,
169            _ => panic!("expected binary"),
170        };
171        let model = proto::Model::decode(bytes.as_slice()).unwrap();
172        let proto::model::Type::MlProgram(prog) = model.r#type.as_ref().unwrap();
173        assert_eq!(
174            prog.functions[0].block.as_ref().unwrap().operations[0].r#type,
175            "add"
176        );
177    }
178}
nxpu_backend_coreml/lib.rs

nxpu_backend_coreml/
lib.rs
