rightkit_embed/
coreml_units.rs

1//! Core ML compute-unit choice per sequence length. Plain data and parsing with no Core ML
2//! dependency, so the type exists on every platform (an `EmbedderConfig` carries it everywhere)
3//! and its tests run anywhere. Only a Core ML bundle on macOS acts on it.
4
5use std::fmt;
6use std::str::FromStr;
7
8use crate::e40::BUCKETS;
9
10/// Longest graph (in tokens) the default map sends to the Neural Engine; longer graphs use the GPU.
11pub const NEURAL_ENGINE_MAX_TOKENS: usize = 128;
12
13/// Which Core ML compute units run one function (`MLComputeUnits`).
14#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
15pub enum CoreMlComputeUnit {
16    /// `MLComputeUnitsAll`: Core ML chooses among CPU, GPU and Neural Engine.
17    All,
18    /// `MLComputeUnitsCPUAndNeuralEngine`.
19    CpuAndNe,
20    /// `MLComputeUnitsCPUAndGPU`.
21    CpuAndGpu,
22    /// `MLComputeUnitsCPUOnly`: the parity reference.
23    CpuOnly,
24}
25
26impl CoreMlComputeUnit {
27    /// Canonical spelling accepted by [`FromStr`]: `all`, `cpu_and_ne`, `cpu_and_gpu`, `cpu_only`.
28    pub const fn name(self) -> &'static str {
29        match self {
30            Self::All => "all",
31            Self::CpuAndNe => "cpu_and_ne",
32            Self::CpuAndGpu => "cpu_and_gpu",
33            Self::CpuOnly => "cpu_only",
34        }
35    }
36}
37
38impl fmt::Display for CoreMlComputeUnit {
39    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
40        f.write_str(self.name())
41    }
42}
43
44/// An unreadable compute-unit or unit-map specification.
45#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
46#[error("invalid Core ML compute units: {0}")]
47pub struct CoreMlUnitsParseError(String);
48
49impl FromStr for CoreMlComputeUnit {
50    type Err = CoreMlUnitsParseError;
51
52    /// Case-insensitive; `-`, `_` and spaces are ignored, so `CPU_AND_NE`, `cpu-and-ne` and
53    /// `cpuandne` are the same. Accepts `all`, `cpu_and_ne` (`ne`), `cpu_and_gpu` (`gpu`) and
54    /// `cpu_only` (`cpu`).
55    fn from_str(text: &str) -> Result<Self, Self::Err> {
56        let key: String = text
57            .chars()
58            .filter(|c| !matches!(c, '-' | '_' | ' '))
59            .map(|c| c.to_ascii_lowercase())
60            .collect();
61        match key.as_str() {
62            "all" => Ok(Self::All),
63            "cpuandne" | "cpuandneuralengine" | "ne" => Ok(Self::CpuAndNe),
64            "cpuandgpu" | "gpu" => Ok(Self::CpuAndGpu),
65            "cpuonly" | "cpu" => Ok(Self::CpuOnly),
66            _ => Err(CoreMlUnitsParseError(format!(
67                "`{}` is not one of all, cpu_and_ne, cpu_and_gpu, cpu_only",
68                text.trim()
69            ))),
70        }
71    }
72}
73
74/// The compute units of each Core ML function (`embed_32` .. `embed_512`).
75///
76/// [`Self::per_length_default`] is what a Core ML bundle uses unless
77/// [`EmbedderConfig::with_coreml_units`](crate::EmbedderConfig::with_coreml_units) says
78/// otherwise. Each function is loaded once on its own unit and kept resident.
79#[derive(Debug, Clone, PartialEq, Eq)]
80pub struct CoreMlUnits {
81    by_length: [CoreMlComputeUnit; BUCKETS.len()],
82}
83
84impl CoreMlUnits {
85    /// The measured default (Mac matrix, RightKit issue #81): `cpu_and_ne` for graphs of up to
86    /// 128 tokens (32, 64, 128; p50 about 2.3-3.0 ms) and `cpu_and_gpu` for 256 and 512
87    /// (p50 about 5.1 ms and 9.8-18.7 ms). The first `cpu_and_gpu` call per length can take
88    /// up to 51 s cold (shader build), which is why load pre-warms every function.
89    pub fn per_length_default() -> Self {
90        let mut by_length = [CoreMlComputeUnit::CpuAndNe; BUCKETS.len()];
91        for (unit, length) in by_length.iter_mut().zip(BUCKETS) {
92            if length > NEURAL_ENGINE_MAX_TOKENS {
93                *unit = CoreMlComputeUnit::CpuAndGpu;
94            }
95        }
96        Self { by_length }
97    }
98
99    /// Every function on `unit` (forces `all`, `cpu_and_ne`, `cpu_and_gpu` or `cpu_only`).
100    pub fn uniform(unit: CoreMlComputeUnit) -> Self {
101        Self {
102            by_length: [unit; BUCKETS.len()],
103        }
104    }
105
106    /// This map with the function of `length` tokens moved to `unit`. `None` when `length` is
107    /// not one of 32, 64, 128, 256, 512.
108    pub fn with_length(mut self, length: usize, unit: CoreMlComputeUnit) -> Option<Self> {
109        let index = BUCKETS.iter().position(|bucket| *bucket == length)?;
110        self.by_length[index] = unit;
111        Some(self)
112    }
113
114    /// The unit of the function of `length` tokens; `None` for a length that is not a graph.
115    pub fn unit_for(&self, length: usize) -> Option<CoreMlComputeUnit> {
116        BUCKETS
117            .iter()
118            .position(|bucket| *bucket == length)
119            .map(|index| self.by_length[index])
120    }
121
122    /// `(length, unit)` for every function, shortest first.
123    pub fn entries(&self) -> impl Iterator<Item = (usize, CoreMlComputeUnit)> + '_ {
124        BUCKETS.iter().copied().zip(self.by_length.iter().copied())
125    }
126
127    /// True when every function runs on `cpu_only`.
128    pub fn is_cpu_only(&self) -> bool {
129        self.by_length
130            .iter()
131            .all(|unit| *unit == CoreMlComputeUnit::CpuOnly)
132    }
133}
134
135impl From<CoreMlComputeUnit> for CoreMlUnits {
136    /// Every function on `unit`, like [`CoreMlUnits::uniform`].
137    fn from(unit: CoreMlComputeUnit) -> Self {
138        Self::uniform(unit)
139    }
140}
141
142impl Default for CoreMlUnits {
143    fn default() -> Self {
144        Self::per_length_default()
145    }
146}
147
148impl fmt::Display for CoreMlUnits {
149    /// `32=cpu_and_ne,64=cpu_and_ne,...`: parses back to the same map.
150    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
151        for (position, (length, unit)) in self.entries().enumerate() {
152            if position > 0 {
153                f.write_str(",")?;
154            }
155            write!(f, "{length}={unit}")?;
156        }
157        Ok(())
158    }
159}
160
161impl FromStr for CoreMlUnits {
162    type Err = CoreMlUnitsParseError;
163
164    /// A caller-facing override string (a setting or environment variable):
165    ///
166    /// - `default` or `per_length`: [`Self::per_length_default`];
167    /// - a bare unit (`all`, `cpu_and_ne`, `cpu_and_gpu`, `cpu_only`): every function on it;
168    /// - `length=unit` pairs separated by commas (`256=cpu_only,512=cpu_only`): those lengths
169    ///   override the per-length default, the others keep it. A length may appear once.
170    fn from_str(text: &str) -> Result<Self, Self::Err> {
171        let text = text.trim();
172        let fail = |why: String| CoreMlUnitsParseError(why);
173        if text.is_empty() {
174            return Err(fail("empty specification".into()));
175        }
176        let key = text.to_ascii_lowercase().replace(['-', ' '], "_");
177        if key == "default" || key == "per_length" {
178            return Ok(Self::per_length_default());
179        }
180        if !text.contains('=') {
181            return text.parse().map(Self::uniform);
182        }
183        let mut units = Self::per_length_default();
184        let mut seen = Vec::new();
185        for pair in text.split(',') {
186            let (length, unit) = pair
187                .split_once('=')
188                .ok_or_else(|| fail(format!("`{}` is not length=unit", pair.trim())))?;
189            let length: usize = length
190                .trim()
191                .parse()
192                .map_err(|_| fail(format!("`{}` is not a token length", length.trim())))?;
193            if seen.contains(&length) {
194                return Err(fail(format!("length {length} is given twice")));
195            }
196            seen.push(length);
197            units = units.with_length(length, unit.parse()?).ok_or_else(|| {
198                fail(format!(
199                    "{length} is not a function length (one of {BUCKETS:?})"
200                ))
201            })?;
202        }
203        Ok(units)
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210    use CoreMlComputeUnit::{All, CpuAndGpu, CpuAndNe, CpuOnly};
211
212    #[test]
213    fn the_default_map_is_ne_up_to_128_and_gpu_above() {
214        let units = CoreMlUnits::per_length_default();
215        let entries: Vec<_> = units.entries().collect();
216        assert_eq!(
217            entries,
218            vec![
219                (32, CpuAndNe),
220                (64, CpuAndNe),
221                (128, CpuAndNe),
222                (256, CpuAndGpu),
223                (512, CpuAndGpu),
224            ]
225        );
226        assert_eq!(units, CoreMlUnits::default());
227        assert_eq!(units.unit_for(128), Some(CpuAndNe));
228        assert_eq!(units.unit_for(256), Some(CpuAndGpu));
229        assert_eq!(units.unit_for(48), None, "embed_48 is not served");
230        assert!(!units.is_cpu_only());
231    }
232
233    #[test]
234    fn a_uniform_map_forces_one_unit() {
235        for unit in [All, CpuAndNe, CpuAndGpu, CpuOnly] {
236            let units = CoreMlUnits::uniform(unit);
237            assert!(units.entries().all(|(_, u)| u == unit));
238        }
239        assert!(CoreMlUnits::uniform(CpuOnly).is_cpu_only());
240        assert!(!CoreMlUnits::uniform(All).is_cpu_only());
241    }
242
243    #[test]
244    fn unit_names_parse_loosely_and_round_trip() {
245        for unit in [All, CpuAndNe, CpuAndGpu, CpuOnly] {
246            assert_eq!(unit.name().parse::<CoreMlComputeUnit>(), Ok(unit));
247            assert_eq!(unit.to_string(), unit.name());
248        }
249        assert_eq!("CPU-and-NE".parse::<CoreMlComputeUnit>(), Ok(CpuAndNe));
250        assert_eq!(" cpuAndGpu ".parse::<CoreMlComputeUnit>(), Ok(CpuAndGpu));
251        assert_eq!("cpu".parse::<CoreMlComputeUnit>(), Ok(CpuOnly));
252        assert!("ane".parse::<CoreMlComputeUnit>().is_err());
253        assert!("".parse::<CoreMlComputeUnit>().is_err());
254    }
255
256    #[test]
257    fn override_strings_parse() {
258        assert_eq!(
259            "default".parse::<CoreMlUnits>().unwrap(),
260            CoreMlUnits::per_length_default()
261        );
262        assert_eq!(
263            "cpu_only".parse::<CoreMlUnits>().unwrap(),
264            CoreMlUnits::uniform(CpuOnly)
265        );
266        assert_eq!(
267            "ALL".parse::<CoreMlUnits>().unwrap(),
268            CoreMlUnits::uniform(All)
269        );
270        let mixed: CoreMlUnits = "256=cpu_only, 32=all".parse().unwrap();
271        assert_eq!(mixed.unit_for(32), Some(All));
272        assert_eq!(
273            mixed.unit_for(64),
274            Some(CpuAndNe),
275            "unlisted keeps the default"
276        );
277        assert_eq!(mixed.unit_for(256), Some(CpuOnly));
278        assert_eq!(mixed.unit_for(512), Some(CpuAndGpu));
279    }
280
281    #[test]
282    fn bad_override_strings_are_refused() {
283        for bad in [
284            "",
285            "  ",
286            "fast",
287            "48=all",
288            "256=fast",
289            "x=all",
290            "256",
291            "256=all,256=cpu_only",
292            "256=all,,",
293            "all,512=all",
294        ] {
295            assert!(
296                bad.parse::<CoreMlUnits>().is_err(),
297                "`{bad}` must not parse"
298            );
299        }
300    }
301
302    #[test]
303    fn display_parses_back_to_the_same_map() {
304        let units = CoreMlUnits::per_length_default();
305        assert_eq!(
306            units.to_string(),
307            "32=cpu_and_ne,64=cpu_and_ne,128=cpu_and_ne,256=cpu_and_gpu,512=cpu_and_gpu"
308        );
309        assert_eq!(units.to_string().parse::<CoreMlUnits>().unwrap(), units);
310        let forced = CoreMlUnits::uniform(All);
311        assert_eq!(forced.to_string().parse::<CoreMlUnits>().unwrap(), forced);
312    }
313}
rightkit_embed/coreml_units.rs

rightkit_embed/
coreml_units.rs
