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::{
18    COREML_TOSA_CAPABILITY, COREML_TOSA_TARGET, LoweringError, supports_tosa_dtype,
19    supports_tosa_operator,
20};
21pub use mlprogram::{COREML_TOSA_INTEGER_CAPABILITY, COREML_TOSA_INTEGER_TARGET};
22
23use virtio_accel_tosa::CapabilityDescriptor;
24#[cfg(not(target_os = "macos"))]
25use virtio_accel_tosa::TosaCapabilityProvider;
26
27#[cfg(target_os = "macos")]
28const FLOAT_CAPABILITIES: &[CapabilityDescriptor] = &[COREML_TOSA_CAPABILITY];
29#[cfg(target_os = "macos")]
30const ALL_CAPABILITIES: &[CapabilityDescriptor] =
31    &[COREML_TOSA_CAPABILITY, COREML_TOSA_INTEGER_CAPABILITY];
32#[cfg(not(target_os = "macos"))]
33const NO_CAPABILITIES: &[CapabilityDescriptor] = &[];
34
35use virtio_accel_core::{ArtifactFormat, TargetIdentity};
36
37/// Provider artifact format for [`CoreMlArtifact`].
38pub const ARTIFACT_FORMAT: ArtifactFormat = match ArtifactFormat::new(0x434d_4c50) {
39    Some(format) => format,
40    None => panic!("Core ML artifact format must be nonzero"),
41};
42
43/// Core ML path-artifact ABI v1 targeting CPU plus Apple Neural Engine execution.
44pub const TARGET_IDENTITY: TargetIdentity = TargetIdentity([
45    0x434f_5245,
46    0x4d4c_0001,
47    0x414e_4503,
48    0x4d41_434f,
49    0x000e_0000,
50    0,
51    0,
52    0,
53    0,
54    0,
55    0,
56    0,
57]);
58
59/// The Core ML runtime does not publish a finite upper bound for model residency.
60///
61/// Requiring the maximal charge makes the provider promise truthful: a process cannot retain
62/// `u64::MAX` bytes for one model. Device integrations must set their aggregate program-residency
63/// policy accordingly when admitting a Core ML program.
64pub const REQUIRED_RESIDENT_BYTES: u64 = u64::MAX;
65
66/// Failure to initialize a Core ML backend instance.
67#[derive(Clone, Copy, Debug, PartialEq, Eq)]
68pub enum InitError {
69    /// The backend is only executable on macOS 14 or newer.
70    UnsupportedPlatform,
71    /// The configured model root is missing, not a directory, or not representable as UTF-8.
72    InvalidModelRoot,
73    /// Core ML does not report an accessible Apple Neural Engine.
74    NeuralEngineUnavailable,
75}
76
77impl std::fmt::Display for InitError {
78    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        write!(formatter, "{self:?}")
80    }
81}
82
83impl std::error::Error for InitError {}
84
85#[cfg(target_os = "macos")]
86mod macos;
87#[cfg(target_os = "macos")]
88pub use macos::{
89    CoreMlAccelerator, CoreMlBuffer, CoreMlContext, CoreMlEvent, CoreMlProgram, CoreMlQueue,
90};
91
92/// Non-macOS placeholder that keeps workspace consumers portable.
93#[cfg(not(target_os = "macos"))]
94#[derive(Clone, Copy, Debug, Default)]
95pub struct CoreMlAccelerator;
96
97#[cfg(not(target_os = "macos"))]
98impl CoreMlAccelerator {
99    /// Report that Core ML is unavailable on this target.
100    pub fn new(_model_root: impl AsRef<std::path::Path>) -> Result<Self, InitError> {
101        Err(InitError::UnsupportedPlatform)
102    }
103
104    /// Report that the native TOSA-to-Core ML execution path is unavailable on this target.
105    pub fn new_tosa() -> Result<Self, InitError> {
106        Err(InitError::UnsupportedPlatform)
107    }
108}
109
110#[cfg(not(target_os = "macos"))]
111impl TosaCapabilityProvider for CoreMlAccelerator {
112    fn tosa_capabilities(&self) -> &'static [CapabilityDescriptor] {
113        NO_CAPABILITIES
114    }
115}
virtio_accel_coreml/lib.rs

virtio_accel_coreml/
lib.rs
