Skip to main content

virtio_accel_coreml/
lib.rs

1//! Core ML host backend for Apple Neural Engine capable Macs.
2//!
3//! The production path accepts device-neutral TOSA 1.0 FlatBuffers, validates and analyzes them
4//! with `virtio-accel-tosa`, and lowers supported static floating-point and exact INT8 graphs
5//! inside this host-native crate. Core ML models are configured with `CPUAndNeuralEngine`: supported
6//! operations may execute on the ANE, while Core ML remains free to place unsupported operations
7//! on the CPU. Program buffers are page-aligned allocations wrapped directly by `MLMultiArray`;
8//! output execution is accepted only when Core ML uses the same allocation as its output backing.
9
10#![cfg_attr(not(target_os = "macos"), forbid(unsafe_code))]
11
12mod artifact;
13mod lower;
14mod mlprogram;
15
16pub use artifact::{ArtifactBuildError, CoreMlArtifact, FeatureRole};
17pub use lower::{COREML_TOSA_TARGET, LoweringError, supports_tosa_dtype, supports_tosa_operator};
18pub use mlprogram::COREML_TOSA_INTEGER_TARGET;
19
20use virtio_accel_core::{ArtifactFormat, TargetIdentity};
21
22/// Provider artifact format for [`CoreMlArtifact`].
23pub const ARTIFACT_FORMAT: ArtifactFormat = match ArtifactFormat::new(0x434d_4c50) {
24    Some(format) => format,
25    None => panic!("Core ML artifact format must be nonzero"),
26};
27
28/// Core ML path-artifact ABI v1 targeting CPU plus Apple Neural Engine execution.
29pub const TARGET_IDENTITY: TargetIdentity = TargetIdentity([
30    0x434f_5245,
31    0x4d4c_0001,
32    0x414e_4503,
33    0x4d41_434f,
34    0x000e_0000,
35    0,
36    0,
37    0,
38    0,
39    0,
40    0,
41    0,
42]);
43
44/// The Core ML runtime does not publish a finite upper bound for model residency.
45///
46/// Requiring the maximal charge makes the provider promise truthful: a process cannot retain
47/// `u64::MAX` bytes for one model. Device integrations must set their aggregate program-residency
48/// policy accordingly when admitting a Core ML program.
49pub const REQUIRED_RESIDENT_BYTES: u64 = u64::MAX;
50
51/// Failure to initialize a Core ML backend instance.
52#[derive(Clone, Copy, Debug, PartialEq, Eq)]
53pub enum InitError {
54    /// The backend is only executable on macOS 14 or newer.
55    UnsupportedPlatform,
56    /// The configured model root is missing, not a directory, or not representable as UTF-8.
57    InvalidModelRoot,
58    /// Core ML does not report an accessible Apple Neural Engine.
59    NeuralEngineUnavailable,
60}
61
62impl std::fmt::Display for InitError {
63    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64        write!(formatter, "{self:?}")
65    }
66}
67
68impl std::error::Error for InitError {}
69
70#[cfg(target_os = "macos")]
71mod macos;
72#[cfg(target_os = "macos")]
73pub use macos::{
74    CoreMlAccelerator, CoreMlBuffer, CoreMlContext, CoreMlEvent, CoreMlProgram, CoreMlQueue,
75};
76
77/// Non-macOS placeholder that keeps workspace consumers portable.
78#[cfg(not(target_os = "macos"))]
79#[derive(Clone, Copy, Debug, Default)]
80pub struct CoreMlAccelerator;
81
82#[cfg(not(target_os = "macos"))]
83impl CoreMlAccelerator {
84    /// Report that Core ML is unavailable on this target.
85    pub fn new(_model_root: impl AsRef<std::path::Path>) -> Result<Self, InitError> {
86        Err(InitError::UnsupportedPlatform)
87    }
88
89    /// Report that the native TOSA-to-Core ML execution path is unavailable on this target.
90    pub fn new_tosa() -> Result<Self, InitError> {
91        Err(InitError::UnsupportedPlatform)
92    }
93}