diff --git a/rust/onnxruntime/examples/issue22.rs b/rust/onnxruntime/examples/issue22.rs index 6c96e899fa774..1fb7fe28ff123 100644 --- a/rust/onnxruntime/examples/issue22.rs +++ b/rust/onnxruntime/examples/issue22.rs @@ -51,5 +51,5 @@ fn main() { let outputs = session.run(inputs).unwrap(); - print!("outputs: {:#?}", outputs[0].float_array().unwrap()); + print!("outputs: {:#?}", outputs[0].float_array().unwrap().view()); } diff --git a/rust/onnxruntime/examples/sample.rs b/rust/onnxruntime/examples/sample.rs index 9af5cf733ccae..b6f351b2082ed 100644 --- a/rust/onnxruntime/examples/sample.rs +++ b/rust/onnxruntime/examples/sample.rs @@ -73,10 +73,11 @@ fn run() -> Result<(), Error> { let outputs = session.run(input_tensor_values)?; let output = outputs[0].float_array().unwrap(); + let view = output.view(); - assert_eq!(output.shape(), output0_shape.as_slice()); + assert_eq!(view.shape(), output0_shape.as_slice()); for i in 0..5 { - println!("Score for class [{}] = {}", i, output[[0, i, 0, 0]]); + println!("Score for class [{}] = {}", i, view[[0, i, 0, 0]]); } Ok(()) diff --git a/rust/onnxruntime/src/session.rs b/rust/onnxruntime/src/session.rs index 326426e35982c..d475d1b724111 100644 --- a/rust/onnxruntime/src/session.rs +++ b/rust/onnxruntime/src/session.rs @@ -410,10 +410,10 @@ impl Session { /// /// Note that ONNX models can have multiple inputs; a `Vec<_>` is thus /// used for the input data here. - pub fn run<'input, 'output>( - &'output self, + pub fn run<'input>( + &self, mut input_arrays: impl AsMut<[Box]> + 'input, - ) -> Result>> { + ) -> Result> { let mut output_tensor_extractors_ptrs: Vec<*mut sys::OrtValue> = vec![std::ptr::null_mut(); self.outputs.len()]; diff --git a/rust/onnxruntime/src/tensor/ort_output_tensor.rs b/rust/onnxruntime/src/tensor/ort_output_tensor.rs index 83663c0d303f8..727cae1db0ef4 100644 --- a/rust/onnxruntime/src/tensor/ort_output_tensor.rs +++ b/rust/onnxruntime/src/tensor/ort_output_tensor.rs @@ -71,22 +71,41 @@ impl Drop for OrtOutputTensor { } /// An Output tensor with the ptr and the item that will copy from the ptr. -#[derive(Debug)] -pub struct WithOutputTensor<'a, T> { - #[allow(dead_code)] +/// +/// The view is materialized on each access via [`view()`](Self::view) to ensure the +/// borrowed lifetime is tied to `&self`, preventing the view from outliving the +/// underlying buffer owned by the `OrtOutputTensor`. +pub struct WithOutputTensor { pub(crate) tensor: OrtOutputTensor, - item: ArrayView<'a, T, ndarray::IxDyn>, + data_ptr: *const T, + shape: Vec, } -impl<'a, T> std::ops::Deref for WithOutputTensor<'a, T> { - type Target = ArrayView<'a, T, ndarray::IxDyn>; +impl Debug for WithOutputTensor { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WithOutputTensor") + .field("tensor", &self.tensor) + .field("data_ptr", &self.data_ptr) + .field("shape", &self.shape) + .finish() + } +} - fn deref(&self) -> &Self::Target { - &self.item +// SAFETY: The data pointer is derived from OrtOutputTensor which owns the allocation. +// Access is only possible through &self (via view()), so Send/Sync follow from T: Send/Sync. +unsafe impl Send for WithOutputTensor {} +unsafe impl Sync for WithOutputTensor {} + +impl WithOutputTensor { + /// Returns an [`ArrayView`] over the output tensor data. + /// + /// The returned view borrows `self`, so it cannot outlive the tensor owner. + pub fn view(&self) -> ArrayView<'_, T, ndarray::IxDyn> { + unsafe { ArrayView::from_shape_ptr(ndarray::IxDyn(&self.shape), self.data_ptr) } } } -impl<'a, T> TryFrom for WithOutputTensor<'a, T> +impl TryFrom for WithOutputTensor where T: TypeToTensorElementDataType, { @@ -110,45 +129,45 @@ where status_to_result(status).map_err(OrtError::IsTensor)?; assert_ne!(output_array_ptr, std::ptr::null_mut()); - let array_view = - unsafe { ArrayView::from_shape_ptr(ndarray::IxDyn(&value.shape), output_array_ptr) }; + let shape = value.shape.clone(); Ok(WithOutputTensor { tensor: value, - item: array_view, + data_ptr: output_array_ptr, + shape, }) } } /// The onnxruntime Run output type. -pub enum OrtOutput<'a> { +pub enum OrtOutput { /// Tensor of f32s - Float(WithOutputTensor<'a, f32>), + Float(WithOutputTensor), /// Tensor of f64s - Double(WithOutputTensor<'a, f64>), + Double(WithOutputTensor), /// Tensor of u8s - UInt8(WithOutputTensor<'a, u8>), + UInt8(WithOutputTensor), /// Tensor of u16s - UInt16(WithOutputTensor<'a, u16>), + UInt16(WithOutputTensor), /// Tensor of u32s - UInt32(WithOutputTensor<'a, u32>), + UInt32(WithOutputTensor), /// Tensor of u64s - UInt64(WithOutputTensor<'a, u64>), + UInt64(WithOutputTensor), /// Tensor of i8s - Int8(WithOutputTensor<'a, i8>), + Int8(WithOutputTensor), /// Tensor of i16s - Int16(WithOutputTensor<'a, i16>), + Int16(WithOutputTensor), /// Tensor of i32s - Int32(WithOutputTensor<'a, i32>), + Int32(WithOutputTensor), /// Tensor of i64s - Int64(WithOutputTensor<'a, i64>), + Int64(WithOutputTensor), /// Tensor of Strings - String(WithOutputTensor<'a, String>), + String(WithOutputTensor), } -impl<'a> OrtOutput<'a> { - /// Return `WithOutputTensor<'a, f32>` which derefs into an `ArrayView`. - pub fn float_array(&self) -> Option<&WithOutputTensor<'a, f32>> { +impl OrtOutput { + /// Return `WithOutputTensor` which provides a `view()` method for an `ArrayView`. + pub fn float_array(&self) -> Option<&WithOutputTensor> { if let Self::Float(item) = self { Some(item) } else { @@ -156,8 +175,8 @@ impl<'a> OrtOutput<'a> { } } - /// Return `WithOutputTensor<'a, f64>` which derefs into an `ArrayView`. - pub fn double_array(&self) -> Option<&WithOutputTensor<'a, f64>> { + /// Return `WithOutputTensor` which provides a `view()` method for an `ArrayView`. + pub fn double_array(&self) -> Option<&WithOutputTensor> { if let Self::Double(item) = self { Some(item) } else { @@ -165,8 +184,8 @@ impl<'a> OrtOutput<'a> { } } - /// Return `WithOutputTensor<'a, u8>` which derefs into an `ArrayView`. - pub fn uint8_array(&self) -> Option<&WithOutputTensor<'a, u8>> { + /// Return `WithOutputTensor` which provides a `view()` method for an `ArrayView`. + pub fn uint8_array(&self) -> Option<&WithOutputTensor> { if let Self::UInt8(item) = self { Some(item) } else { @@ -174,8 +193,8 @@ impl<'a> OrtOutput<'a> { } } - /// Return `WithOutputTensor<'a, u16>` which derefs into an `ArrayView`. - pub fn uint16_array(&self) -> Option<&WithOutputTensor<'a, u16>> { + /// Return `WithOutputTensor` which provides a `view()` method for an `ArrayView`. + pub fn uint16_array(&self) -> Option<&WithOutputTensor> { if let Self::UInt16(item) = self { Some(item) } else { @@ -183,8 +202,8 @@ impl<'a> OrtOutput<'a> { } } - /// Return `WithOutputTensor<'a, u32>` which derefs into an `ArrayView`. - pub fn uint32_array(&self) -> Option<&WithOutputTensor<'a, u32>> { + /// Return `WithOutputTensor` which provides a `view()` method for an `ArrayView`. + pub fn uint32_array(&self) -> Option<&WithOutputTensor> { if let Self::UInt32(item) = self { Some(item) } else { @@ -192,8 +211,8 @@ impl<'a> OrtOutput<'a> { } } - /// Return `WithOutputTensor<'a, u64>` which derefs into an `ArrayView`. - pub fn uint64_array(&self) -> Option<&WithOutputTensor<'a, u64>> { + /// Return `WithOutputTensor` which provides a `view()` method for an `ArrayView`. + pub fn uint64_array(&self) -> Option<&WithOutputTensor> { if let Self::UInt64(item) = self { Some(item) } else { @@ -201,8 +220,8 @@ impl<'a> OrtOutput<'a> { } } - /// Return `WithOutputTensor<'a, i8>` which derefs into an `ArrayView`. - pub fn int8_array(&self) -> Option<&WithOutputTensor<'a, i8>> { + /// Return `WithOutputTensor` which provides a `view()` method for an `ArrayView`. + pub fn int8_array(&self) -> Option<&WithOutputTensor> { if let Self::Int8(item) = self { Some(item) } else { @@ -210,8 +229,8 @@ impl<'a> OrtOutput<'a> { } } - /// Return `WithOutputTensor<'a, i16>` which derefs into an `ArrayView`. - pub fn int16_array(&self) -> Option<&WithOutputTensor<'a, i16>> { + /// Return `WithOutputTensor` which provides a `view()` method for an `ArrayView`. + pub fn int16_array(&self) -> Option<&WithOutputTensor> { if let Self::Int16(item) = self { Some(item) } else { @@ -219,8 +238,8 @@ impl<'a> OrtOutput<'a> { } } - /// Return `WithOutputTensor<'a, i32>` which derefs into an `ArrayView`. - pub fn int32_array(&self) -> Option<&WithOutputTensor<'a, i32>> { + /// Return `WithOutputTensor` which provides a `view()` method for an `ArrayView`. + pub fn int32_array(&self) -> Option<&WithOutputTensor> { if let Self::Int32(item) = self { Some(item) } else { @@ -228,8 +247,8 @@ impl<'a> OrtOutput<'a> { } } - /// Return `WithOutputTensor<'a, i64>` which derefs into an `ArrayView`. - pub fn int64_array(&self) -> Option<&WithOutputTensor<'a, i64>> { + /// Return `WithOutputTensor` which provides a `view()` method for an `ArrayView`. + pub fn int64_array(&self) -> Option<&WithOutputTensor> { if let Self::Int64(item) = self { Some(item) } else { @@ -237,8 +256,8 @@ impl<'a> OrtOutput<'a> { } } - /// Return `WithOutputTensor<'a, String>` which derefs into an `ArrayView`. - pub fn string_array(&self) -> Option<&WithOutputTensor<'a, String>> { + /// Return `WithOutputTensor` which provides a `view()` method for an `ArrayView`. + pub fn string_array(&self) -> Option<&WithOutputTensor> { if let Self::String(item) = self { Some(item) } else { @@ -247,10 +266,10 @@ impl<'a> OrtOutput<'a> { } } -impl<'a> TryFrom for OrtOutput<'a> { +impl TryFrom for OrtOutput { type Error = OrtError; - fn try_from(value: OrtOutputTensor) -> Result> { + fn try_from(value: OrtOutputTensor) -> Result { unsafe { let mut shape_info = std::ptr::null_mut(); diff --git a/rust/onnxruntime/tests/integration_tests.rs b/rust/onnxruntime/tests/integration_tests.rs index 7843fe269e5e4..1c096400eccf7 100644 --- a/rust/onnxruntime/tests/integration_tests.rs +++ b/rust/onnxruntime/tests/integration_tests.rs @@ -112,6 +112,7 @@ mod download { // and iterate on resulting probabilities, creating an index to later access labels. let output = outputs[0].float_array().unwrap(); let mut probabilities: Vec<(usize, f32)> = output + .view() .softmax(ndarray::Axis(1)) .iter() .copied() @@ -209,6 +210,7 @@ mod download { let output = outputs[0].float_array().unwrap(); let mut probabilities: Vec<(usize, f32)> = output + .view() .softmax(ndarray::Axis(1)) .iter() .copied() @@ -301,6 +303,7 @@ mod download { let output = &outputs[0].float_array().unwrap(); let mut probabilities: Vec<(usize, f32)> = output + .view() .softmax(ndarray::Axis(1)) .iter() .copied() @@ -398,6 +401,7 @@ mod download { let output = &outputs[0].float_array().unwrap(); let mut probabilities: Vec<(usize, f32)> = output + .view() .softmax(ndarray::Axis(1)) .iter() .copied() @@ -515,7 +519,7 @@ mod download { let output = outputs[0].float_array().unwrap(); // The image should have doubled in size - assert_eq!(output.shape(), [1, 448, 448, 3]); + assert_eq!(output.view().shape(), [1, 448, 448, 3]); } }