spatialrust_ai/
lib.rs

1//! Backend-independent model, session, named-I/O, and explicit binding contracts.
2//!
3//! The default build contains no inference runtime. ONNX Runtime and hardware
4//! execution providers are additive features and must preserve the copy/device
5//! choices represented by these types.
6
7#![deny(unsafe_code)]
8#![warn(missing_docs)]
9
10use std::{path::PathBuf, sync::Arc};
11
12use spatialrust_tensor::{DataType, Device, TensorBuffer, TensorDescriptor};
13
14mod mock;
15pub use mock::{MockInferenceBackend, MockProfile};
16
17#[cfg(feature = "onnxruntime")]
18mod onnxruntime;
19#[cfg(feature = "onnxruntime")]
20pub use onnxruntime::{OnnxRuntimeBackend, OnnxRuntimeSession};
21
22/// Result type for inference operations.
23pub type AiResult<T> = Result<T, AiError>;
24
25/// Errors shared by inference contracts and backend adapters.
26#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
27pub enum AiError {
28    /// A named input or output appears more than once.
29    #[error("duplicate tensor name `{0}`")]
30    DuplicateName(String),
31    /// A required model input was not supplied.
32    #[error("missing required input `{0}`")]
33    MissingInput(String),
34    /// A tensor name is not declared by the model.
35    #[error("unexpected tensor `{0}`")]
36    UnexpectedTensor(String),
37    /// Actual dtype differs from the model contract.
38    #[error("tensor `{name}` requires {expected:?}, found {actual:?}")]
39    DataTypeMismatch {
40        /// Tensor name.
41        name: String,
42        /// Model dtype.
43        expected: DataType,
44        /// Supplied dtype.
45        actual: DataType,
46    },
47    /// Actual rank or dimension differs from the model contract.
48    #[error("tensor `{name}` shape {actual:?} does not match {expected:?}")]
49    ShapeMismatch {
50        /// Tensor name.
51        name: String,
52        /// Model dimensions.
53        expected: Vec<Dimension>,
54        /// Supplied dimensions.
55        actual: Vec<usize>,
56    },
57    /// A preallocated output does not have enough bytes.
58    #[error("preallocated output `{name}` needs {required} bytes, found {found}")]
59    OutputBufferTooSmall {
60        /// Output name.
61        name: String,
62        /// Required allocation bytes.
63        required: usize,
64        /// Available allocation bytes.
65        found: usize,
66    },
67    /// The selected backend cannot honor an operation or option.
68    #[error("backend `{backend}` does not support {operation}")]
69    Unsupported {
70        /// Backend identifier.
71        backend: String,
72        /// Unsupported operation.
73        operation: String,
74    },
75    /// Backend-specific failure with stable outer context.
76    #[error("backend `{backend}` failed: {message}")]
77    Backend {
78        /// Backend identifier.
79        backend: String,
80        /// Backend error message.
81        message: String,
82    },
83    /// Invalid public configuration.
84    #[error("invalid inference configuration: {0}")]
85    InvalidConfiguration(String),
86    /// A backend would need a copy that the caller did not authorize.
87    #[error(
88        "{direction} copy is required for tensor `{name}`; opt in explicitly or use I/O binding"
89    )]
90    CopyRequired {
91        /// Input or output transfer direction.
92        direction: &'static str,
93        /// Model-visible tensor name.
94        name: String,
95    },
96}
97
98/// One model dimension, fixed, unconstrained, or symbolically dynamic.
99#[derive(Clone, Debug, PartialEq, Eq, Hash)]
100pub enum Dimension {
101    /// Exact dimension size.
102    Fixed(usize),
103    /// Dynamic dimension without a model-provided symbol.
104    Dynamic,
105    /// Dynamic dimension sharing a model-provided symbolic name.
106    Symbol(String),
107}
108
109impl Dimension {
110    fn accepts(&self, actual: usize) -> bool {
111        match self {
112            Self::Fixed(expected) => *expected == actual,
113            Self::Dynamic | Self::Symbol(_) => true,
114        }
115    }
116}
117
118/// One named model input or output contract.
119#[derive(Clone, Debug, PartialEq, Eq)]
120pub struct TensorSpec {
121    /// ONNX/model-visible name.
122    pub name: String,
123    /// Required scalar/vector dtype.
124    pub dtype: DataType,
125    /// Ordered fixed or dynamic dimensions.
126    pub shape: Vec<Dimension>,
127}
128
129impl TensorSpec {
130    /// Creates a named tensor specification.
131    pub fn new(name: impl Into<String>, dtype: DataType, shape: Vec<Dimension>) -> Self {
132        Self { name: name.into(), dtype, shape }
133    }
134
135    /// Validates a concrete tensor descriptor against this specification.
136    pub fn validate(&self, descriptor: &TensorDescriptor) -> AiResult<()> {
137        if descriptor.dtype() != self.dtype {
138            return Err(AiError::DataTypeMismatch {
139                name: self.name.clone(),
140                expected: self.dtype,
141                actual: descriptor.dtype(),
142            });
143        }
144        if descriptor.shape().len() != self.shape.len()
145            || !self
146                .shape
147                .iter()
148                .zip(descriptor.shape())
149                .all(|(expected, &actual)| expected.accepts(actual))
150        {
151            return Err(AiError::ShapeMismatch {
152                name: self.name.clone(),
153                expected: self.shape.clone(),
154                actual: descriptor.shape().to_vec(),
155            });
156        }
157        Ok(())
158    }
159}
160
161/// Model-visible named input and output metadata.
162#[derive(Clone, Debug, Default, PartialEq, Eq)]
163pub struct ModelInfo {
164    /// Optional producer/model identifier.
165    pub name: Option<String>,
166    /// Ordered model inputs.
167    pub inputs: Vec<TensorSpec>,
168    /// Ordered model outputs.
169    pub outputs: Vec<TensorSpec>,
170}
171
172impl ModelInfo {
173    /// Validates uniqueness of all input and output names.
174    pub fn validate(&self) -> AiResult<()> {
175        ensure_unique(self.inputs.iter().map(|spec| spec.name.as_str()))?;
176        ensure_unique(self.outputs.iter().map(|spec| spec.name.as_str()))?;
177        Ok(())
178    }
179
180    /// Validates required names, dtype, and dynamic/fixed shapes for one request.
181    pub fn validate_inputs(&self, inputs: &NamedTensors) -> AiResult<()> {
182        for spec in &self.inputs {
183            let tensor =
184                inputs.get(&spec.name).ok_or_else(|| AiError::MissingInput(spec.name.clone()))?;
185            spec.validate(tensor.descriptor())?;
186        }
187        for (name, _) in inputs.iter() {
188            if !self.inputs.iter().any(|spec| spec.name == name) {
189                return Err(AiError::UnexpectedTensor(name.to_owned()));
190            }
191        }
192        Ok(())
193    }
194}
195
196fn ensure_unique<'a>(names: impl IntoIterator<Item = &'a str>) -> AiResult<()> {
197    let mut seen = Vec::<&str>::new();
198    for name in names {
199        if seen.contains(&name) {
200            return Err(AiError::DuplicateName(name.to_owned()));
201        }
202        seen.push(name);
203    }
204    Ok(())
205}
206
207/// Ordered, uniquely named tensor collection.
208#[derive(Clone, Debug, Default, PartialEq, Eq)]
209pub struct NamedTensors {
210    values: Vec<(String, TensorBuffer)>,
211}
212
213impl NamedTensors {
214    /// Creates an empty collection.
215    pub const fn new() -> Self {
216        Self { values: Vec::new() }
217    }
218
219    /// Inserts a unique named tensor while retaining insertion order.
220    pub fn insert(&mut self, name: impl Into<String>, tensor: TensorBuffer) -> AiResult<()> {
221        let name = name.into();
222        if self.values.iter().any(|(existing, _)| existing == &name) {
223            return Err(AiError::DuplicateName(name));
224        }
225        self.values.push((name, tensor));
226        Ok(())
227    }
228
229    /// Returns a tensor by model-visible name.
230    pub fn get(&self, name: &str) -> Option<&TensorBuffer> {
231        self.values.iter().find(|(candidate, _)| candidate == name).map(|(_, value)| value)
232    }
233
234    /// Returns ordered `(name, tensor)` pairs.
235    pub fn iter(&self) -> impl Iterator<Item = (&str, &TensorBuffer)> {
236        self.values.iter().map(|(name, tensor)| (name.as_str(), tensor))
237    }
238
239    /// Returns the tensor count.
240    pub fn len(&self) -> usize {
241        self.values.len()
242    }
243
244    /// Returns whether no tensors are present.
245    pub fn is_empty(&self) -> bool {
246        self.values.is_empty()
247    }
248
249    /// Consumes the collection into ordered pairs.
250    pub fn into_values(self) -> Vec<(String, TensorBuffer)> {
251        self.values
252    }
253}
254
255/// Model bytes, filesystem path, or built-in mock profile for session creation.
256#[derive(Clone, Debug)]
257pub enum ModelSource {
258    /// Read model bytes from this path during session creation.
259    Path(PathBuf),
260    /// Immutable in-memory model bytes.
261    Bytes(Arc<[u8]>),
262    /// Deterministic in-process mock profile (no ONNX / no weights).
263    Mock(MockProfile),
264}
265
266/// Graph optimization level requested from a backend.
267#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
268pub enum GraphOptimization {
269    /// Disable graph rewrites.
270    Disabled,
271    /// Apply safe basic rewrites.
272    Basic,
273    /// Apply extended rewrites.
274    Extended,
275    /// Apply all backend-supported rewrites.
276    #[default]
277    All,
278}
279
280/// Runtime-independent session configuration.
281#[derive(Clone, Debug, PartialEq, Eq)]
282pub struct SessionOptions {
283    /// Intra-operator worker count; `None` delegates to the backend.
284    pub intra_threads: Option<usize>,
285    /// Inter-operator worker count; `None` delegates to the backend.
286    pub inter_threads: Option<usize>,
287    /// Graph rewrite level.
288    pub graph_optimization: GraphOptimization,
289    /// Request deterministic kernels where supported.
290    pub deterministic: bool,
291}
292
293/// Whether a run may allocate and copy host tensor data.
294#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
295pub enum CopyPolicy {
296    /// Fail instead of silently copying.
297    #[default]
298    Forbid,
299    /// Permit a documented host-to-host copy for this run.
300    Allow,
301}
302
303/// Per-run host copy permissions.
304#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
305pub struct RunOptions {
306    /// Permission to repack/copy named inputs into backend values.
307    pub input_copy: CopyPolicy,
308    /// Permission to copy backend-owned outputs into `TensorBuffer`.
309    pub output_copy: CopyPolicy,
310}
311
312impl Default for SessionOptions {
313    fn default() -> Self {
314        Self {
315            intra_threads: None,
316            inter_threads: None,
317            graph_optimization: GraphOptimization::All,
318            deterministic: false,
319        }
320    }
321}
322
323impl SessionOptions {
324    /// Rejects zero thread counts.
325    pub fn validate(&self) -> AiResult<()> {
326        if self.intra_threads == Some(0) || self.inter_threads == Some(0) {
327            return Err(AiError::InvalidConfiguration("thread counts must be positive".into()));
328        }
329        Ok(())
330    }
331}
332
333/// Explicit output destination for an I/O-bound run.
334#[derive(Clone, Debug, PartialEq, Eq)]
335pub enum OutputBinding {
336    /// Ask the backend to allocate this output on the named device.
337    Allocate {
338        /// Model output name.
339        name: String,
340        /// Required allocation device.
341        device: Device,
342    },
343    /// Write into this caller-owned CPU allocation.
344    PreallocatedCpu(PreallocatedOutput),
345}
346
347/// Caller-owned mutable bytes for an explicitly bound CPU output.
348#[derive(Clone, Debug, PartialEq, Eq)]
349pub struct PreallocatedOutput {
350    name: String,
351    descriptor: TensorDescriptor,
352    storage: Option<PreallocatedStorage>,
353}
354
355#[derive(Clone, Debug)]
356pub(crate) enum PreallocatedStorage {
357    Bytes(Vec<u8>),
358    U16(Vec<u16>),
359    F32(Vec<f32>),
360}
361
362impl PreallocatedStorage {
363    fn allocation_bytes(&self) -> &[u8] {
364        match self {
365            Self::Bytes(values) => values,
366            Self::U16(values) => bytemuck::cast_slice(values),
367            Self::F32(values) => bytemuck::cast_slice(values),
368        }
369    }
370
371    fn allocation_bytes_mut(&mut self) -> &mut [u8] {
372        match self {
373            Self::Bytes(values) => values,
374            Self::U16(values) => bytemuck::cast_slice_mut(values),
375            Self::F32(values) => bytemuck::cast_slice_mut(values),
376        }
377    }
378}
379
380impl PartialEq for PreallocatedStorage {
381    fn eq(&self, other: &Self) -> bool {
382        self.allocation_bytes() == other.allocation_bytes()
383    }
384}
385
386impl Eq for PreallocatedStorage {}
387
388impl PreallocatedOutput {
389    /// Creates a checked preallocated CPU output.
390    pub fn try_new(
391        name: impl Into<String>,
392        descriptor: TensorDescriptor,
393        bytes: Vec<u8>,
394    ) -> AiResult<Self> {
395        let name = name.into();
396        let required = descriptor
397            .required_byte_range()
398            .map_err(|error| AiError::InvalidConfiguration(error.to_string()))?
399            .end;
400        if !descriptor.device().is_host_accessible() {
401            return Err(AiError::InvalidConfiguration(
402                "PreallocatedCpu requires host-accessible storage".into(),
403            ));
404        }
405        if bytes.len() < required {
406            return Err(AiError::OutputBufferTooSmall { name, required, found: bytes.len() });
407        }
408        Ok(Self { name, descriptor, storage: Some(PreallocatedStorage::Bytes(bytes)) })
409    }
410
411    /// Allocates aligned storage for a compact `u8`, `u16`, or `f32` CPU output.
412    pub fn allocate(name: impl Into<String>, descriptor: TensorDescriptor) -> AiResult<Self> {
413        if !descriptor.is_c_contiguous() || descriptor.byte_offset() != 0 {
414            return Err(AiError::InvalidConfiguration(
415                "preallocated output must be compact with byte_offset=0".into(),
416            ));
417        }
418        if descriptor.device() != Device::CPU {
419            return Err(AiError::InvalidConfiguration(
420                "preallocated output allocation currently requires Device::CPU".into(),
421            ));
422        }
423        let elements = descriptor
424            .element_count()
425            .map_err(|error| AiError::InvalidConfiguration(error.to_string()))?;
426        let storage = match descriptor.dtype() {
427            DataType::U8 => PreallocatedStorage::Bytes(vec![0; elements]),
428            DataType::U16 => PreallocatedStorage::U16(vec![0; elements]),
429            DataType::F32 => PreallocatedStorage::F32(vec![0.0; elements]),
430            dtype => {
431                return Err(AiError::InvalidConfiguration(format!(
432                    "preallocated output dtype {dtype:?} is not supported; use u8, u16, or f32"
433                )))
434            }
435        };
436        Ok(Self { name: name.into(), descriptor, storage: Some(storage) })
437    }
438
439    /// Returns the model output name.
440    pub fn name(&self) -> &str {
441        &self.name
442    }
443
444    /// Returns output metadata.
445    pub const fn descriptor(&self) -> &TensorDescriptor {
446        &self.descriptor
447    }
448
449    /// Returns mutable caller-owned output bytes for a backend binding.
450    pub fn allocation_bytes_mut(&mut self) -> AiResult<&mut [u8]> {
451        self.storage.as_mut().map(PreallocatedStorage::allocation_bytes_mut).ok_or_else(|| {
452            AiError::InvalidConfiguration(format!(
453                "preallocated output `{}` has already been consumed",
454                self.name
455            ))
456        })
457    }
458
459    /// Returns whether a bound run has taken ownership of this allocation.
460    pub fn is_consumed(&self) -> bool {
461        self.storage.is_none()
462    }
463
464    /// Converts a completed binding into generic owned tensor storage.
465    pub fn into_tensor(mut self) -> AiResult<TensorBuffer> {
466        let storage = self.take_storage()?;
467        tensor_from_preallocated(storage, self.descriptor)
468    }
469
470    pub(crate) fn take_storage(&mut self) -> AiResult<PreallocatedStorage> {
471        self.storage.take().ok_or_else(|| {
472            AiError::InvalidConfiguration(format!(
473                "preallocated output `{}` has already been consumed",
474                self.name
475            ))
476        })
477    }
478}
479
480fn tensor_from_preallocated(
481    storage: PreallocatedStorage,
482    descriptor: TensorDescriptor,
483) -> AiResult<TensorBuffer> {
484    let result = match storage {
485        PreallocatedStorage::Bytes(values) => TensorBuffer::try_new(values, descriptor),
486        PreallocatedStorage::U16(values) => TensorBuffer::try_from_u16(values, descriptor),
487        PreallocatedStorage::F32(values) => TensorBuffer::try_from_f32(values, descriptor),
488    };
489    result.map_err(|error| AiError::InvalidConfiguration(error.to_string()))
490}
491
492/// Inputs, requested output destinations, and completed named results.
493#[derive(Clone, Debug, PartialEq, Eq)]
494pub struct IoBinding {
495    inputs: NamedTensors,
496    outputs: Vec<OutputBinding>,
497    results: Option<NamedTensors>,
498}
499
500impl IoBinding {
501    /// Creates an explicit binding request.
502    pub fn try_new(inputs: NamedTensors, outputs: Vec<OutputBinding>) -> AiResult<Self> {
503        ensure_unique(outputs.iter().map(|output| match output {
504            OutputBinding::Allocate { name, .. } => name.as_str(),
505            OutputBinding::PreallocatedCpu(output) => output.name(),
506        }))?;
507        Ok(Self { inputs, outputs, results: None })
508    }
509
510    /// Returns bound named inputs.
511    pub const fn inputs(&self) -> &NamedTensors {
512        &self.inputs
513    }
514
515    /// Returns requested output destinations.
516    pub fn outputs(&self) -> &[OutputBinding] {
517        &self.outputs
518    }
519
520    /// Returns mutable requested output destinations.
521    pub fn outputs_mut(&mut self) -> &mut [OutputBinding] {
522        &mut self.outputs
523    }
524
525    /// Stores completed named results. Intended for backend implementers.
526    pub fn set_results(&mut self, results: NamedTensors) {
527        self.results = Some(results);
528    }
529
530    /// Clears results before a backend starts another run.
531    pub fn clear_results(&mut self) {
532        self.results = None;
533    }
534
535    /// Returns completed results after a successful bound run.
536    pub fn results(&self) -> Option<&NamedTensors> {
537        self.results.as_ref()
538    }
539
540    /// Consumes completed results, if present.
541    pub fn into_results(self) -> Option<NamedTensors> {
542        self.results
543    }
544}
545
546/// Stable interface implemented by inference engines.
547pub trait InferenceBackend: Send + Sync {
548    /// Stable backend identifier such as `onnxruntime-cpu`.
549    fn name(&self) -> &str;
550
551    /// Loads one model and returns an independent mutable session.
552    fn create_session(
553        &self,
554        source: &ModelSource,
555        options: &SessionOptions,
556    ) -> AiResult<Box<dyn ModelSession>>;
557}
558
559/// Loaded model session with named dynamic I/O.
560pub trait ModelSession: Send {
561    /// Stable backend identifier for diagnostics and capability checks.
562    fn backend_name(&self) -> &str;
563
564    /// Returns model input/output metadata captured at load time.
565    fn model_info(&self) -> &ModelInfo;
566
567    /// Runs without authorizing hidden host copies.
568    fn run(&mut self, inputs: NamedTensors) -> AiResult<NamedTensors> {
569        self.run_with_options(inputs, RunOptions::default())
570    }
571
572    /// Runs named I/O with explicit host copy permissions.
573    fn run_with_options(
574        &mut self,
575        inputs: NamedTensors,
576        options: RunOptions,
577    ) -> AiResult<NamedTensors>;
578
579    /// Runs with explicit input/output device or allocation bindings.
580    fn run_with_binding(&mut self, _binding: &mut IoBinding) -> AiResult<()> {
581        Err(AiError::Unsupported {
582            backend: self.backend_name().into(),
583            operation: "explicit I/O binding".into(),
584        })
585    }
586}
587
588#[cfg(test)]
589mod tests {
590    use super::{
591        AiError, Dimension, IoBinding, ModelInfo, NamedTensors, OutputBinding, PreallocatedOutput,
592        TensorSpec,
593    };
594    use spatialrust_tensor::{DataType, Device, TensorBuffer, TensorDescriptor};
595
596    fn tensor(shape: Vec<usize>) -> TensorBuffer {
597        let len = shape.iter().product::<usize>() * 4;
598        TensorBuffer::try_new(
599            vec![0; len],
600            TensorDescriptor::contiguous(DataType::F32, shape, Device::CPU),
601        )
602        .unwrap()
603    }
604
605    #[test]
606    fn validates_named_dynamic_inputs() {
607        let info = ModelInfo {
608            name: Some("dynamic-batch".into()),
609            inputs: vec![TensorSpec::new(
610                "images",
611                DataType::F32,
612                vec![Dimension::Symbol("batch".into()), Dimension::Fixed(3), Dimension::Dynamic],
613            )],
614            outputs: vec![],
615        };
616        let mut inputs = NamedTensors::new();
617        inputs.insert("images", tensor(vec![4, 3, 224])).unwrap();
618        info.validate_inputs(&inputs).unwrap();
619        let wrong = TensorDescriptor::contiguous(DataType::F32, vec![4, 1, 224], Device::CPU);
620        assert!(matches!(info.inputs[0].validate(&wrong), Err(AiError::ShapeMismatch { .. })));
621    }
622
623    #[test]
624    fn named_tensors_reject_duplicates_and_missing_inputs() {
625        let mut inputs = NamedTensors::new();
626        inputs.insert("x", tensor(vec![2])).unwrap();
627        assert!(matches!(inputs.insert("x", tensor(vec![2])), Err(AiError::DuplicateName(_))));
628        let info = ModelInfo {
629            name: None,
630            inputs: vec![TensorSpec::new("y", DataType::F32, vec![Dimension::Fixed(2)])],
631            outputs: vec![],
632        };
633        assert!(matches!(info.validate_inputs(&inputs), Err(AiError::MissingInput(_))));
634    }
635
636    #[test]
637    fn io_binding_requires_explicit_unique_outputs() {
638        let descriptor = TensorDescriptor::contiguous(DataType::F32, vec![2], Device::CPU);
639        let output = PreallocatedOutput::allocate("scores", descriptor).unwrap();
640        let binding =
641            IoBinding::try_new(NamedTensors::new(), vec![OutputBinding::PreallocatedCpu(output)])
642                .unwrap();
643        assert_eq!(binding.outputs().len(), 1);
644        assert!(IoBinding::try_new(
645            NamedTensors::new(),
646            vec![
647                OutputBinding::Allocate { name: "y".into(), device: Device::CPU },
648                OutputBinding::Allocate { name: "y".into(), device: Device::CPU },
649            ],
650        )
651        .is_err());
652    }
653}
spatialrust_ai/lib.rs

spatialrust_ai/
lib.rs
