beat_this/runtime/
ort.rs

1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3
4use anyhow::Result;
5use ndarray::ArrayD;
6use ort::ep::{CoreML, ExecutionProvider};
7use ort::session::builder::GraphOptimizationLevel;
8use ort::session::Session;
9use ort::value::{DynValue, Value};
10
11use super::{Model, Runtime, Tensor};
12
13/// Ort-based ONNX inference runtime.
14///
15/// Automatically tries CoreML on macOS (falls back to CPU if unavailable).
16pub struct OrtRuntime {
17    pub optimization_level: GraphOptimizationLevel,
18    pub intra_threads: usize,
19    pub profiling_path: Option<PathBuf>,
20}
21
22impl Default for OrtRuntime {
23    fn default() -> Self {
24        Self {
25            optimization_level: GraphOptimizationLevel::Level3,
26            // 0 = let ORT pick the thread count automatically.
27            // This is critical for performance: on Apple Silicon M1, intra_threads=1
28            // gives ~15.7s for beat inference on a 4.5-min track, while auto (0) gives
29            // ~3.4s — a 4.6x speedup. The bottleneck is batched MatMul in the attention
30            // layers (73% of inference time), which parallelizes well across cores.
31            intra_threads: 0,
32            profiling_path: None,
33        }
34    }
35}
36
37impl OrtRuntime {
38    /// Check if CoreML is available in the loaded ORT runtime.
39    pub fn is_coreml_available(&self) -> bool {
40        CoreML::default().is_available().unwrap_or(false)
41    }
42}
43
44impl Runtime for OrtRuntime {
45    type Model = OrtModel;
46
47    #[allow(clippy::needless_match)]
48    fn load_model(&self, path: &Path) -> Result<OrtModel> {
49        // Match is needed because GraphOptimizationLevel doesn't implement Copy or Clone.
50        let optimization_level = match self.optimization_level {
51            GraphOptimizationLevel::Disable => GraphOptimizationLevel::Disable,
52            GraphOptimizationLevel::Level1 => GraphOptimizationLevel::Level1,
53            GraphOptimizationLevel::Level2 => GraphOptimizationLevel::Level2,
54            GraphOptimizationLevel::Level3 => GraphOptimizationLevel::Level3,
55            GraphOptimizationLevel::All => GraphOptimizationLevel::All,
56        };
57        let mut builder = Session::builder()?
58            .with_optimization_level(optimization_level)?
59            .with_intra_threads(self.intra_threads)?
60            .with_execution_providers([CoreML::default().build()])?;
61        if let Some(ref profile_path) = self.profiling_path {
62            builder = builder.with_profiling(profile_path)?;
63        }
64        let session = builder.commit_from_file(path)?;
65        Ok(OrtModel { session })
66    }
67}
68
69/// An ort-backed model wrapping `ort::Session`.
70pub struct OrtModel {
71    session: Session,
72}
73
74impl OrtModel {
75    /// End profiling and flush the trace JSON file. Returns the profile file path.
76    pub fn end_profiling(&mut self) -> Result<String> {
77        Ok(self.session.end_profiling()?)
78    }
79}
80
81impl Model for OrtModel {
82    fn run(&mut self, inputs: &[(&str, &Tensor)]) -> Result<HashMap<String, Tensor>> {
83        // Convert Tensor inputs to ort DynValues
84        let ort_inputs: Vec<(String, DynValue)> = inputs
85            .iter()
86            .map(|(name, tensor)| {
87                let shape: Vec<usize> = tensor.shape.clone();
88                let array = ArrayD::from_shape_vec(shape, tensor.data.clone())?;
89                let value: DynValue = Value::from_array(array)?.into_dyn();
90                Ok((name.to_string(), value))
91            })
92            .collect::<Result<Vec<_>>>()?;
93
94        // Build input refs for session.run()
95        let input_refs: Vec<(&str, &DynValue)> = ort_inputs
96            .iter()
97            .map(|(name, value)| (name.as_str(), value))
98            .collect();
99
100        let outputs = self.session.run(input_refs)?;
101
102        // Convert outputs to Tensor map
103        let mut result = HashMap::new();
104        for (name, value) in outputs.iter() {
105            let (shape, data) = value.try_extract_tensor::<f32>()?;
106            let tensor = Tensor {
107                shape: shape.iter().map(|&d| d as usize).collect(),
108                data: data.to_vec(),
109            };
110            result.insert(name.to_string(), tensor);
111        }
112
113        Ok(result)
114    }
115}
beat_this/runtime/ort.rs

beat_this/runtime/
ort.rs
