spatialrust_ai/
onnxruntime.rs

1//! ONNX Runtime CPU execution provider adapter.
2
3use std::sync::Arc;
4
5use ort::{
6    ep::CPU,
7    memory::MemoryInfo,
8    session::{builder::GraphOptimizationLevel, Session},
9    value::{DynValue, Tensor, TensorElementType, TensorRef, TensorValueType, ValueType},
10};
11use spatialrust_tensor::{DataType, Device, HostTensorStorage, TensorBuffer, TensorDescriptor};
12
13use crate::{
14    AiError, AiResult, CopyPolicy, Dimension, GraphOptimization, InferenceBackend,
15    IoBinding as SpatialIoBinding, ModelInfo, ModelSession, ModelSource, NamedTensors,
16    OutputBinding, PreallocatedStorage, RunOptions, SessionOptions, TensorSpec,
17};
18
19const BACKEND_NAME: &str = "onnxruntime-cpu";
20
21/// ONNX Runtime backend configured explicitly for the CPU execution provider.
22#[derive(Clone, Copy, Debug, Default)]
23pub struct OnnxRuntimeBackend;
24
25impl InferenceBackend for OnnxRuntimeBackend {
26    fn name(&self) -> &str {
27        BACKEND_NAME
28    }
29
30    fn create_session(
31        &self,
32        source: &ModelSource,
33        options: &SessionOptions,
34    ) -> AiResult<Box<dyn ModelSession>> {
35        options.validate()?;
36        let mut builder = Session::builder().map_err(ort_error)?;
37        builder = builder.with_execution_providers([CPU::default().build()]).map_err(ort_error)?;
38        if let Some(threads) = options.intra_threads {
39            builder = builder.with_intra_threads(threads).map_err(ort_error)?;
40        }
41        if let Some(threads) = options.inter_threads {
42            builder = builder.with_inter_threads(threads).map_err(ort_error)?;
43        }
44        builder = builder
45            .with_optimization_level(match options.graph_optimization {
46                GraphOptimization::Disabled => GraphOptimizationLevel::Disable,
47                GraphOptimization::Basic => GraphOptimizationLevel::Level1,
48                GraphOptimization::Extended => GraphOptimizationLevel::Level2,
49                GraphOptimization::All => GraphOptimizationLevel::Level3,
50            })
51            .map_err(ort_error)?;
52        builder = builder.with_deterministic_compute(options.deterministic).map_err(ort_error)?;
53        let session = match source {
54            ModelSource::Path(path) => builder.commit_from_file(path).map_err(ort_error)?,
55            ModelSource::Bytes(bytes) => builder.commit_from_memory(bytes).map_err(ort_error)?,
56            ModelSource::Mock(_) => {
57                return Err(AiError::Unsupported {
58                    backend: BACKEND_NAME.into(),
59                    operation: "ModelSource::Mock (use MockInferenceBackend)".into(),
60                });
61            }
62        };
63        let info = model_info(&session)?;
64        Ok(Box::new(OnnxRuntimeSession { session, info }))
65    }
66}
67
68/// Loaded ONNX Runtime CPU session.
69#[derive(Debug)]
70pub struct OnnxRuntimeSession {
71    session: Session,
72    info: ModelInfo,
73}
74
75impl ModelSession for OnnxRuntimeSession {
76    fn backend_name(&self) -> &str {
77        BACKEND_NAME
78    }
79
80    fn model_info(&self) -> &ModelInfo {
81        &self.info
82    }
83
84    fn run_with_options(
85        &mut self,
86        inputs: NamedTensors,
87        options: RunOptions,
88    ) -> AiResult<NamedTensors> {
89        self.info.validate_inputs(&inputs)?;
90        if options.input_copy == CopyPolicy::Forbid {
91            return Err(AiError::CopyRequired {
92                direction: "input host",
93                name: self.info.inputs.first().map_or_else(String::new, |spec| spec.name.clone()),
94            });
95        }
96        if options.output_copy == CopyPolicy::Forbid {
97            return Err(AiError::CopyRequired {
98                direction: "output host",
99                name: self.info.outputs.first().map_or_else(String::new, |spec| spec.name.clone()),
100            });
101        }
102        let values = inputs
103            .iter()
104            .map(|(name, tensor)| Ok((name.to_owned(), to_ort_tensor(name, tensor)?)))
105            .collect::<AiResult<Vec<_>>>()?;
106        let outputs = self.session.run(values).map_err(ort_error)?;
107        let mut named = NamedTensors::new();
108        for spec in &self.info.outputs {
109            let value = outputs.get(&spec.name).ok_or_else(|| AiError::Backend {
110                backend: BACKEND_NAME.into(),
111                message: format!("runtime omitted output `{}`", spec.name),
112            })?;
113            named.insert(spec.name.clone(), copy_ort_output(spec, value)?)?;
114        }
115        Ok(named)
116    }
117
118    fn run_with_binding(&mut self, binding: &mut SpatialIoBinding) -> AiResult<()> {
119        binding.clear_results();
120        self.info.validate_inputs(binding.inputs())?;
121        let mut ort_binding = self.session.create_binding().map_err(ort_error)?;
122        for (name, tensor) in binding.inputs().iter() {
123            bind_zero_copy_input(&mut ort_binding, name, tensor)?;
124        }
125
126        let requested = binding
127            .outputs()
128            .iter()
129            .map(|output| match output {
130                OutputBinding::Allocate { name, .. } => name.clone(),
131                OutputBinding::PreallocatedCpu(output) => output.name().to_owned(),
132            })
133            .collect::<Vec<_>>();
134        for output in binding.outputs_mut() {
135            match output {
136                OutputBinding::Allocate { name, device } => {
137                    output_spec(&self.info, name)?;
138                    if *device != Device::CPU {
139                        return Err(AiError::Unsupported {
140                            backend: BACKEND_NAME.into(),
141                            operation: format!("bound output device {device:?}"),
142                        });
143                    }
144                    ort_binding
145                        .bind_output_to_device(name.clone(), &MemoryInfo::default())
146                        .map_err(ort_error)?;
147                }
148                OutputBinding::PreallocatedCpu(output) => {
149                    let spec = output_spec(&self.info, output.name())?;
150                    spec.validate(output.descriptor())?;
151                    if !output.descriptor().is_c_contiguous()
152                        || output.descriptor().byte_offset() != 0
153                    {
154                        return Err(AiError::Unsupported {
155                            backend: BACKEND_NAME.into(),
156                            operation: format!(
157                                "strided preallocated output `{}`; allocate a compact buffer",
158                                output.name()
159                            ),
160                        });
161                    }
162                    let name = output.name().to_owned();
163                    let shape = output.descriptor().shape().to_vec();
164                    bind_preallocated_output(
165                        &mut ort_binding,
166                        &name,
167                        shape,
168                        output.take_storage()?,
169                        spec.dtype,
170                    )?;
171                }
172            }
173        }
174
175        let mut outputs = self.session.run_binding(&ort_binding).map_err(ort_error)?;
176        let mut results = NamedTensors::new();
177        for name in requested {
178            let spec = output_spec(&self.info, &name)?;
179            let value = outputs.remove(&name).ok_or_else(|| AiError::Backend {
180                backend: BACKEND_NAME.into(),
181                message: format!("runtime omitted bound output `{name}`"),
182            })?;
183            results.insert(name, retain_ort_output(spec, value)?)?;
184        }
185        drop(outputs);
186        binding.set_results(results);
187        Ok(())
188    }
189}
190
191fn output_spec<'a>(info: &'a ModelInfo, name: &str) -> AiResult<&'a TensorSpec> {
192    info.outputs
193        .iter()
194        .find(|spec| spec.name == name)
195        .ok_or_else(|| AiError::UnexpectedTensor(name.to_owned()))
196}
197
198fn bind_zero_copy_input(
199    binding: &mut ort::session::IoBinding,
200    name: &str,
201    tensor: &TensorBuffer,
202) -> AiResult<()> {
203    let descriptor = tensor.descriptor();
204    if !descriptor.is_c_contiguous() || descriptor.byte_offset() != 0 {
205        return Err(AiError::Unsupported {
206            backend: BACKEND_NAME.into(),
207            operation: format!("strided bound input `{name}`; explicitly pack it first"),
208        });
209    }
210    if let Some(storage) = tensor.host_storage() {
211        if let Some(storage) = storage.as_any().downcast_ref::<OrtHostStorage>() {
212            return storage.bind_input(binding, name);
213        }
214    }
215    let expected = descriptor
216        .element_count()
217        .map_err(|error| AiError::InvalidConfiguration(error.to_string()))?;
218    let shape = descriptor.shape().to_vec();
219    macro_rules! bind_arc {
220        ($getter:ident, $type:ty) => {{
221            let values = tensor.$getter().ok_or_else(|| AiError::CopyRequired {
222                direction: "input alignment",
223                name: name.to_owned(),
224            })?;
225            if values.len() != expected {
226                return Err(AiError::InvalidConfiguration(format!(
227                    "bound input `{name}` allocation has {} elements, expected {expected}",
228                    values.len()
229                )));
230            }
231            let value = TensorRef::<$type>::from_array_view((shape, values)).map_err(ort_error)?;
232            binding.bind_input(name, &value).map_err(ort_error)
233        }};
234    }
235    match descriptor.dtype() {
236        DataType::U8 => bind_arc!(shared_bytes, u8),
237        DataType::U16 => bind_arc!(shared_u16, u16),
238        DataType::U32 => bind_arc!(shared_u32, u32),
239        DataType::I16 => bind_arc!(shared_i16, i16),
240        DataType::I32 => bind_arc!(shared_i32, i32),
241        DataType::I64 => bind_arc!(shared_i64, i64),
242        DataType::F32 => bind_arc!(shared_f32, f32),
243        DataType::F64 => bind_arc!(shared_f64, f64),
244        dtype => Err(AiError::Unsupported {
245            backend: BACKEND_NAME.into(),
246            operation: format!("zero-copy bound input dtype {dtype:?}"),
247        }),
248    }
249}
250
251fn bind_preallocated_output(
252    binding: &mut ort::session::IoBinding,
253    name: &str,
254    shape: Vec<usize>,
255    storage: PreallocatedStorage,
256    dtype: DataType,
257) -> AiResult<()> {
258    match (storage, dtype) {
259        (PreallocatedStorage::Bytes(values), DataType::U8) => binding
260            .bind_output(name, Tensor::<u8>::from_array((shape, values)).map_err(ort_error)?)
261            .map_err(ort_error),
262        (PreallocatedStorage::U16(values), DataType::U16) => binding
263            .bind_output(name, Tensor::<u16>::from_array((shape, values)).map_err(ort_error)?)
264            .map_err(ort_error),
265        (PreallocatedStorage::F32(values), DataType::F32) => binding
266            .bind_output(name, Tensor::<f32>::from_array((shape, values)).map_err(ort_error)?)
267            .map_err(ort_error),
268        _ => Err(AiError::CopyRequired { direction: "output alignment", name: name.to_owned() }),
269    }
270}
271
272#[derive(Debug)]
273enum OrtHostStorage {
274    U8(Tensor<u8>),
275    U16(Tensor<u16>),
276    U32(Tensor<u32>),
277    I8(Tensor<i8>),
278    I16(Tensor<i16>),
279    I32(Tensor<i32>),
280    I64(Tensor<i64>),
281    F32(Tensor<f32>),
282    F64(Tensor<f64>),
283}
284
285impl OrtHostStorage {
286    fn bind_input(&self, binding: &mut ort::session::IoBinding, name: &str) -> AiResult<()> {
287        macro_rules! bind {
288            ($value:expr) => {
289                binding.bind_input(name, $value).map_err(ort_error)
290            };
291        }
292        match self {
293            Self::U8(value) => bind!(value),
294            Self::U16(value) => bind!(value),
295            Self::U32(value) => bind!(value),
296            Self::I8(value) => bind!(value),
297            Self::I16(value) => bind!(value),
298            Self::I32(value) => bind!(value),
299            Self::I64(value) => bind!(value),
300            Self::F32(value) => bind!(value),
301            Self::F64(value) => bind!(value),
302        }
303    }
304}
305
306impl HostTensorStorage for OrtHostStorage {
307    fn dtype(&self) -> DataType {
308        match self {
309            Self::U8(_) => DataType::U8,
310            Self::U16(_) => DataType::U16,
311            Self::U32(_) => DataType::U32,
312            Self::I8(_) => DataType::I8,
313            Self::I16(_) => DataType::I16,
314            Self::I32(_) => DataType::I32,
315            Self::I64(_) => DataType::I64,
316            Self::F32(_) => DataType::F32,
317            Self::F64(_) => DataType::F64,
318        }
319    }
320
321    fn allocation_bytes(&self) -> &[u8] {
322        macro_rules! bytes {
323            ($value:expr) => {
324                bytemuck::cast_slice($value.extract_tensor().1)
325            };
326        }
327        match self {
328            Self::U8(value) => value.extract_tensor().1,
329            Self::U16(value) => bytes!(value),
330            Self::U32(value) => bytes!(value),
331            Self::I8(value) => bytes!(value),
332            Self::I16(value) => bytes!(value),
333            Self::I32(value) => bytes!(value),
334            Self::I64(value) => bytes!(value),
335            Self::F32(value) => bytes!(value),
336            Self::F64(value) => bytes!(value),
337        }
338    }
339
340    fn as_any(&self) -> &dyn std::any::Any {
341        self
342    }
343}
344
345fn retain_ort_output(spec: &TensorSpec, value: DynValue) -> AiResult<TensorBuffer> {
346    macro_rules! retain {
347        ($type:ty, $variant:ident) => {{
348            let tensor = value.downcast::<TensorValueType<$type>>().map_err(ort_error)?;
349            let shape = tensor
350                .extract_tensor()
351                .0
352                .iter()
353                .map(|&dimension| usize::try_from(dimension).expect("runtime shape is concrete"))
354                .collect::<Vec<_>>();
355            let storage: Arc<dyn HostTensorStorage> = Arc::new(OrtHostStorage::$variant(tensor));
356            (shape, storage)
357        }};
358    }
359    let (shape, storage) = match spec.dtype {
360        DataType::U8 => retain!(u8, U8),
361        DataType::U16 => retain!(u16, U16),
362        DataType::U32 => retain!(u32, U32),
363        DataType::I8 => retain!(i8, I8),
364        DataType::I16 => retain!(i16, I16),
365        DataType::I32 => retain!(i32, I32),
366        DataType::I64 => retain!(i64, I64),
367        DataType::F32 => retain!(f32, F32),
368        DataType::F64 => retain!(f64, F64),
369        dtype => {
370            return Err(AiError::Unsupported {
371                backend: BACKEND_NAME.into(),
372                operation: format!("retaining bound output dtype {dtype:?}"),
373            })
374        }
375    };
376    TensorBuffer::try_from_host_storage(
377        storage,
378        TensorDescriptor::contiguous(spec.dtype, shape, Device::CPU),
379    )
380    .map_err(|error| AiError::Backend { backend: BACKEND_NAME.into(), message: error.to_string() })
381}
382
383fn model_info(session: &Session) -> AiResult<ModelInfo> {
384    let inputs = session.inputs().iter().map(outlet_spec).collect::<AiResult<Vec<_>>>()?;
385    let outputs = session.outputs().iter().map(outlet_spec).collect::<AiResult<Vec<_>>>()?;
386    let info = ModelInfo { name: None, inputs, outputs };
387    info.validate()?;
388    Ok(info)
389}
390
391fn outlet_spec(outlet: &ort::value::Outlet) -> AiResult<TensorSpec> {
392    let ValueType::Tensor { ty, shape, dimension_symbols } = outlet.dtype() else {
393        return Err(AiError::Unsupported {
394            backend: BACKEND_NAME.into(),
395            operation: format!("non-tensor ONNX value `{}`", outlet.name()),
396        });
397    };
398    let dtype = decode_ort_dtype(*ty).ok_or_else(|| AiError::Unsupported {
399        backend: BACKEND_NAME.into(),
400        operation: format!("ONNX dtype {ty} for `{}`", outlet.name()),
401    })?;
402    let dimensions = shape
403        .iter()
404        .enumerate()
405        .map(|(index, &dimension)| {
406            if dimension >= 0 {
407                Dimension::Fixed(dimension as usize)
408            } else {
409                let symbol = &dimension_symbols[index];
410                if symbol.is_empty() {
411                    Dimension::Dynamic
412                } else {
413                    Dimension::Symbol(symbol.clone())
414                }
415            }
416        })
417        .collect();
418    Ok(TensorSpec::new(outlet.name(), dtype, dimensions))
419}
420
421fn decode_ort_dtype(dtype: TensorElementType) -> Option<DataType> {
422    Some(match dtype {
423        TensorElementType::Float32 => DataType::F32,
424        TensorElementType::Float64 => DataType::F64,
425        TensorElementType::Uint8 => DataType::U8,
426        TensorElementType::Uint16 => DataType::U16,
427        TensorElementType::Uint32 => DataType::U32,
428        TensorElementType::Int8 => DataType::I8,
429        TensorElementType::Int16 => DataType::I16,
430        TensorElementType::Int32 => DataType::I32,
431        TensorElementType::Int64 => DataType::I64,
432        TensorElementType::Float16 => DataType::F16,
433        TensorElementType::Bfloat16 => DataType::BF16,
434        TensorElementType::Bool => DataType::BOOL,
435        _ => return None,
436    })
437}
438
439fn to_ort_tensor(name: &str, tensor: &TensorBuffer) -> AiResult<DynValue> {
440    let descriptor = tensor.descriptor();
441    if !descriptor.is_c_contiguous() {
442        return Err(AiError::Unsupported {
443            backend: BACKEND_NAME.into(),
444            operation: format!("strided input `{name}`; explicitly pack it first"),
445        });
446    }
447    let range = descriptor.required_byte_range().map_err(|error| AiError::Backend {
448        backend: BACKEND_NAME.into(),
449        message: error.to_string(),
450    })?;
451    let bytes = &tensor.allocation_bytes()[range];
452    let shape = descriptor.shape().to_vec();
453    macro_rules! typed {
454        ($type:ty, $width:expr) => {{
455            let values = decode_ne::<$type>(bytes, $width, |chunk| {
456                <$type>::from_ne_bytes(chunk.try_into().expect("fixed chunk width"))
457            });
458            Tensor::<$type>::from_array((shape, values)).map(|value| value.into_dyn())
459        }};
460    }
461    let value = match descriptor.dtype() {
462        DataType::U8 => Tensor::<u8>::from_array((shape, bytes.to_vec())).map(|v| v.into_dyn()),
463        DataType::U16 => typed!(u16, 2),
464        DataType::U32 => typed!(u32, 4),
465        DataType::I8 => Tensor::<i8>::from_array((
466            shape,
467            bytes.iter().map(|&value| value as i8).collect::<Vec<_>>(),
468        ))
469        .map(|v| v.into_dyn()),
470        DataType::I16 => typed!(i16, 2),
471        DataType::I32 => typed!(i32, 4),
472        DataType::I64 => typed!(i64, 8),
473        DataType::F32 => typed!(f32, 4),
474        DataType::F64 => typed!(f64, 8),
475        dtype => {
476            return Err(AiError::Unsupported {
477                backend: BACKEND_NAME.into(),
478                operation: format!("input dtype {dtype:?}"),
479            })
480        }
481    };
482    value.map_err(ort_error)
483}
484
485fn decode_ne<T>(bytes: &[u8], width: usize, decode: impl Fn(&[u8]) -> T) -> Vec<T> {
486    debug_assert_eq!(bytes.len() % width, 0);
487    bytes.chunks_exact(width).map(decode).collect()
488}
489
490fn copy_ort_output(spec: &TensorSpec, value: &DynValue) -> AiResult<TensorBuffer> {
491    macro_rules! extract {
492        ($type:ty, $constructor:ident) => {{
493            let (shape, values) = value.try_extract_tensor::<$type>().map_err(ort_error)?;
494            let descriptor = TensorDescriptor::contiguous(
495                spec.dtype,
496                shape.iter().map(|&dimension| dimension as usize).collect(),
497                Device::CPU,
498            );
499            TensorBuffer::$constructor(values.to_vec(), descriptor)
500        }};
501    }
502    let result = match spec.dtype {
503        DataType::U8 => extract!(u8, try_new),
504        DataType::U16 => extract!(u16, try_from_u16),
505        DataType::U32 => extract!(u32, try_from_u32),
506        DataType::I8 => {
507            let (shape, values) = value.try_extract_tensor::<i8>().map_err(ort_error)?;
508            let descriptor = TensorDescriptor::contiguous(
509                DataType::I8,
510                shape.iter().map(|&dimension| dimension as usize).collect(),
511                Device::CPU,
512            );
513            TensorBuffer::try_new(values.iter().map(|&item| item as u8).collect(), descriptor)
514        }
515        DataType::I16 => extract!(i16, try_from_i16),
516        DataType::I32 => extract!(i32, try_from_i32),
517        DataType::I64 => extract!(i64, try_from_i64),
518        DataType::F32 => extract!(f32, try_from_f32),
519        DataType::F64 => extract!(f64, try_from_f64),
520        dtype => {
521            return Err(AiError::Unsupported {
522                backend: BACKEND_NAME.into(),
523                operation: format!("output dtype {dtype:?}"),
524            })
525        }
526    };
527    result.map_err(|error| AiError::Backend {
528        backend: BACKEND_NAME.into(),
529        message: error.to_string(),
530    })
531}
532
533fn ort_error(error: impl std::fmt::Display) -> AiError {
534    AiError::Backend { backend: BACKEND_NAME.into(), message: error.to_string() }
535}
536
537#[cfg(test)]
538mod tests {
539    use std::sync::Arc;
540
541    use super::OnnxRuntimeBackend;
542    use crate::{
543        AiError, CopyPolicy, Dimension, InferenceBackend, IoBinding, ModelSource, NamedTensors,
544        OutputBinding, PreallocatedOutput, RunOptions, SessionOptions,
545    };
546    use spatialrust_tensor::{DataType, Device, TensorBuffer, TensorDescriptor};
547
548    const DOUBLE_DYNAMIC: &[u8] = &[
549        8, 8, 18, 16, 115, 112, 97, 116, 105, 97, 108, 114, 117, 115, 116, 45, 116, 101, 115, 116,
550        58, 106, 10, 27, 10, 5, 105, 110, 112, 117, 116, 10, 5, 105, 110, 112, 117, 116, 18, 6,
551        111, 117, 116, 112, 117, 116, 34, 3, 65, 100, 100, 18, 14, 100, 111, 117, 98, 108, 101, 95,
552        100, 121, 110, 97, 109, 105, 99, 90, 28, 10, 5, 105, 110, 112, 117, 116, 18, 19, 10, 17, 8,
553        1, 18, 13, 10, 7, 18, 5, 98, 97, 116, 99, 104, 10, 2, 8, 3, 98, 29, 10, 6, 111, 117, 116,
554        112, 117, 116, 18, 19, 10, 17, 8, 1, 18, 13, 10, 7, 18, 5, 98, 97, 116, 99, 104, 10, 2, 8,
555        3, 66, 4, 10, 0, 16, 13,
556    ];
557
558    fn f32_tensor(shape: Vec<usize>, values: &[f32]) -> TensorBuffer {
559        TensorBuffer::try_from_f32(
560            values.to_vec(),
561            TensorDescriptor::contiguous(DataType::F32, shape, Device::CPU),
562        )
563        .unwrap()
564    }
565
566    fn f32_values(tensor: &TensorBuffer) -> Vec<f32> {
567        tensor
568            .allocation_bytes()
569            .chunks_exact(4)
570            .map(|chunk| f32::from_ne_bytes(chunk.try_into().unwrap()))
571            .collect()
572    }
573
574    fn dynamic_identity_model_with_onnx_dtype(dtype: u8) -> Arc<[u8]> {
575        let mut model = DOUBLE_DYNAMIC.to_vec();
576        let graph = model
577            .windows(4)
578            .position(|window| window == [58, 106, 10, 27])
579            .expect("embedded graph and node lengths");
580        model[graph + 1] = 104;
581        let node = graph + 2;
582        let identity_node = [
583            10, 25, 10, 5, 105, 110, 112, 117, 116, 18, 6, 111, 117, 116, 112, 117, 116, 34, 8, 73,
584            100, 101, 110, 116, 105, 116, 121,
585        ];
586        model.splice(node..node + 29, identity_node);
587        let mut replacements = 0;
588        for index in 0..model.len().saturating_sub(3) {
589            if model[index..index + 4] == [8, 1, 18, 13] {
590                model[index + 1] = dtype;
591                replacements += 1;
592            }
593        }
594        assert_eq!(replacements, 2, "input and output tensor types must both be replaced");
595        Arc::from(model)
596    }
597
598    #[test]
599    fn cpu_session_preserves_named_dynamic_contract_and_requires_copy_opt_in() {
600        let backend = OnnxRuntimeBackend;
601        let mut session = backend
602            .create_session(
603                &ModelSource::Bytes(Arc::from(DOUBLE_DYNAMIC)),
604                &SessionOptions::default(),
605            )
606            .unwrap();
607        assert_eq!(session.model_info().inputs[0].name, "input");
608        assert_eq!(
609            session.model_info().inputs[0].shape,
610            vec![Dimension::Symbol("batch".into()), Dimension::Fixed(3)]
611        );
612
613        let mut inputs = NamedTensors::new();
614        inputs.insert("input", f32_tensor(vec![2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])).unwrap();
615        assert!(matches!(session.run(inputs.clone()), Err(AiError::CopyRequired { .. })));
616        let outputs = session
617            .run_with_options(
618                inputs,
619                RunOptions { input_copy: CopyPolicy::Allow, output_copy: CopyPolicy::Allow },
620            )
621            .unwrap();
622        let output = outputs.get("output").unwrap();
623        assert_eq!(output.descriptor().shape(), &[2, 3]);
624        assert_eq!(f32_values(output), &[2.0, 4.0, 6.0, 8.0, 10.0, 12.0]);
625    }
626
627    #[test]
628    fn io_binding_allocates_dynamic_output_and_reuses_it_as_zero_copy_input() {
629        let backend = OnnxRuntimeBackend;
630        let mut session = backend
631            .create_session(
632                &ModelSource::Bytes(Arc::from(DOUBLE_DYNAMIC)),
633                &SessionOptions::default(),
634            )
635            .unwrap();
636        let mut inputs = NamedTensors::new();
637        inputs.insert("input", f32_tensor(vec![1, 3], &[1.0, 2.0, 3.0])).unwrap();
638        let mut binding = IoBinding::try_new(
639            inputs,
640            vec![OutputBinding::Allocate { name: "output".into(), device: Device::CPU }],
641        )
642        .unwrap();
643        session.run_with_binding(&mut binding).unwrap();
644        let first = binding.results().unwrap().get("output").unwrap();
645        assert!(first.host_storage().is_some());
646        assert_eq!(f32_values(first), &[2.0, 4.0, 6.0]);
647
648        let mut inputs = NamedTensors::new();
649        inputs.insert("input", first.clone()).unwrap();
650        let mut chained = IoBinding::try_new(
651            inputs,
652            vec![OutputBinding::Allocate { name: "output".into(), device: Device::CPU }],
653        )
654        .unwrap();
655        session.run_with_binding(&mut chained).unwrap();
656        assert_eq!(
657            f32_values(chained.results().unwrap().get("output").unwrap()),
658            &[4.0, 8.0, 12.0]
659        );
660    }
661
662    #[test]
663    fn io_binding_writes_directly_into_caller_preallocated_f32_storage() {
664        let backend = OnnxRuntimeBackend;
665        let mut session = backend
666            .create_session(
667                &ModelSource::Bytes(Arc::from(DOUBLE_DYNAMIC)),
668                &SessionOptions::default(),
669            )
670            .unwrap();
671        let mut inputs = NamedTensors::new();
672        inputs.insert("input", f32_tensor(vec![2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])).unwrap();
673        let descriptor = TensorDescriptor::contiguous(DataType::F32, vec![2, 3], Device::CPU);
674        let mut output = PreallocatedOutput::allocate("output", descriptor).unwrap();
675        let allocation = output.allocation_bytes_mut().unwrap().as_ptr();
676        let mut binding =
677            IoBinding::try_new(inputs, vec![OutputBinding::PreallocatedCpu(output)]).unwrap();
678        session.run_with_binding(&mut binding).unwrap();
679        let result = binding.results().unwrap().get("output").unwrap();
680        assert_eq!(result.allocation_bytes().as_ptr(), allocation);
681        assert_eq!(f32_values(result), &[2.0, 4.0, 6.0, 8.0, 10.0, 12.0]);
682    }
683
684    #[test]
685    fn io_binding_rejects_unaligned_raw_f32_storage_instead_of_copying() {
686        let backend = OnnxRuntimeBackend;
687        let mut session = backend
688            .create_session(
689                &ModelSource::Bytes(Arc::from(DOUBLE_DYNAMIC)),
690                &SessionOptions::default(),
691            )
692            .unwrap();
693        let descriptor = TensorDescriptor::contiguous(DataType::F32, vec![1, 3], Device::CPU);
694        let input = TensorBuffer::try_new(vec![0; 12], descriptor).unwrap();
695        let mut inputs = NamedTensors::new();
696        inputs.insert("input", input).unwrap();
697        let mut binding = IoBinding::try_new(
698            inputs,
699            vec![OutputBinding::Allocate { name: "output".into(), device: Device::CPU }],
700        )
701        .unwrap();
702        assert!(matches!(
703            session.run_with_binding(&mut binding),
704            Err(AiError::CopyRequired { direction: "input alignment", .. })
705        ));
706    }
707
708    #[test]
709    fn io_binding_supports_aligned_u8_and_u16_storage() {
710        let backend = OnnxRuntimeBackend;
711
712        let mut u8_session = backend
713            .create_session(
714                &ModelSource::Bytes(dynamic_identity_model_with_onnx_dtype(2)),
715                &SessionOptions::default(),
716            )
717            .unwrap();
718        let mut u8_inputs = NamedTensors::new();
719        u8_inputs
720            .insert(
721                "input",
722                TensorBuffer::try_new(
723                    vec![1, 2, 3],
724                    TensorDescriptor::contiguous(DataType::U8, vec![1, 3], Device::CPU),
725                )
726                .unwrap(),
727            )
728            .unwrap();
729        let mut u8_binding = IoBinding::try_new(
730            u8_inputs,
731            vec![OutputBinding::Allocate { name: "output".into(), device: Device::CPU }],
732        )
733        .unwrap();
734        u8_session.run_with_binding(&mut u8_binding).unwrap();
735        assert_eq!(
736            u8_binding.results().unwrap().get("output").unwrap().allocation_bytes(),
737            &[1, 2, 3]
738        );
739
740        let mut u16_session = backend
741            .create_session(
742                &ModelSource::Bytes(dynamic_identity_model_with_onnx_dtype(4)),
743                &SessionOptions::default(),
744            )
745            .unwrap();
746        let descriptor = TensorDescriptor::contiguous(DataType::U16, vec![1, 3], Device::CPU);
747        let mut u16_inputs = NamedTensors::new();
748        u16_inputs
749            .insert(
750                "input",
751                TensorBuffer::try_from_u16(vec![10, 20, 30], descriptor.clone()).unwrap(),
752            )
753            .unwrap();
754        let output = PreallocatedOutput::allocate("output", descriptor).unwrap();
755        let mut u16_binding =
756            IoBinding::try_new(u16_inputs, vec![OutputBinding::PreallocatedCpu(output)]).unwrap();
757        u16_session.run_with_binding(&mut u16_binding).unwrap();
758        assert_eq!(
759            bytemuck::cast_slice::<u8, u16>(
760                u16_binding.results().unwrap().get("output").unwrap().allocation_bytes()
761            ),
762            &[10, 20, 30]
763        );
764    }
765}
spatialrust_ai/onnxruntime.rs

spatialrust_ai/
onnxruntime.rs
