ort 2.0.0-rc.15

use alloc::{string::String, sync::Arc};
use core::{
	fmt::Debug,
	ptr::{self, NonNull}
};

use crate::{
	AsPointer, Error, ErrorCode,
	error::Result,
	memory::MemoryInfo,
	ortsys,
	session::{Session, SharedSessionInner},
	util::{MiniMap, with_cstr},
	value::{AsValueMut, DynTensor, Value, ValueInner, ValueTypeMarker}
};

/// Enables binding of session inputs and/or outputs to pre-allocated memory.
///
/// [`IoBinding`] minimizes copies between a device (like a GPU) and the host (CPU) by allowing you to bind a
/// certain input/output to a pre-allocated value on a specific device.
///
/// [`IoBinding`] is most suitable for:
/// - An ensemble of models in which the output from one model is the input to another and does not need to pass through
///   the CPU to perform additional processing.
/// - Situations where an output should stay on a device (e.g. to perform additional hardware-accelerated processing).
/// - Models that accept an input that does not change for multiple subsequent runs (like the conditional embedding for
///   a diffusion model).
///
/// [`IoBinding`] will not provide any meaningful benefit for:
/// - Models where every input changes with each invocation, such as a causal language model or object recognition
///   model.
/// - Pipelines that go straight from CPU -> GPU -> CPU.
///
/// # Example
/// A diffusion model which takes a text condition input.
///
/// ```no_run
/// # use ort::{
/// # 	ep,
/// # 	memory::{Allocator, AllocatorType, AllocationDevice, MemoryInfo, MemoryType},
/// # 	session::{Session, IoBinding},
/// # 	value::Tensor
/// # };
/// # fn main() -> ort::Result<()> {
/// # let env = ort::test_util::test_env().clone();
/// let mut text_encoder = Session::builder(&env)?
/// 	.with_execution_providers([
/// 		#[cfg(feature = "cuda")]
/// 		ep::CUDA::default().build()
/// 	])?
/// 	.commit_from_file("text_encoder.onnx")?;
/// let mut unet = Session::builder(&env)?
/// 	.with_execution_providers([
/// 		#[cfg(feature = "cuda")]
/// 		ep::CUDA::default().build()
/// 	])?
/// 	.commit_from_file("unet.onnx")?;
///
/// let text_condition = text_encoder
/// 	.run(ort::inputs![Tensor::<i64>::from_array((
/// 		vec![27],
/// 		vec![
/// 			23763, 15460, 473, 68, 312, 265, 17463, 4098, 304, 1077, 283, 198, 7676, 5976, 272, 285, 3609, 435,
/// 			21680, 321, 265, 300, 1689, 64, 285, 4763, 64
/// 		]
/// 	))?])?
/// 	.remove("output0")
/// 	.unwrap();
///
/// let input_allocator = Allocator::new(
/// 	&unet,
/// 	MemoryInfo::new(AllocationDevice::CUDA_PINNED, 0, AllocatorType::Device, MemoryType::CPUInput)?
/// )?;
/// let mut latents = Tensor::<f32>::new(&input_allocator, [1_usize, 4, 64, 64])?;
///
/// let mut io_binding = unet.create_binding()?;
/// io_binding.bind_input("condition", &text_condition)?;
///
/// let output_allocator = Allocator::new(
/// 	&unet,
/// 	MemoryInfo::new(AllocationDevice::CUDA_PINNED, 0, AllocatorType::Device, MemoryType::CPUOutput)?
/// )?;
/// io_binding.bind_output("noise_pred", Tensor::<f32>::new(&output_allocator, [1_usize, 4, 64, 64])?)?;
///
/// for _ in 0..20 {
/// 	io_binding.bind_input("latents", &latents)?;
/// 	let noise_pred = unet.run_binding(&io_binding)?.remove("noise_pred").unwrap();
///
/// 	let mut latents = latents.extract_array_mut();
/// 	latents += &noise_pred.try_extract_array::<f32>()?;
/// }
/// # Ok(())
/// # }
/// ```
///
/// [`IoBinding`] may provide a decent speedup in this example since the `condition` tensor is unchanging between runs.
/// If we were to use normal session inference, the `condition` tensor would be needlessly copied with each invocation
/// of `unet.run()`, and this copying can come with significant latency & overhead. With [`IoBinding`], the `condition`
/// tensor is only copied to the device once instead of 20 times.
#[derive(Debug)]
pub struct IoBinding {
	ptr: NonNull<ort_sys::OrtIoBinding>,
	held_inputs: MiniMap<String, Arc<ValueInner>>,
	pub(crate) output_values: MiniMap<String, Option<Arc<ValueInner>>>,
	_session: Arc<SharedSessionInner>
}

impl IoBinding {
	pub(crate) fn new(session: &Session) -> Result<Self> {
		let mut ptr: *mut ort_sys::OrtIoBinding = ptr::null_mut();
		ortsys![unsafe CreateIoBinding(session.ptr().cast_mut(), &mut ptr)?; nonNull(ptr)];
		crate::logging::create!(IoBinding, ptr);
		Ok(Self {
			ptr,
			held_inputs: MiniMap::new(),
			output_values: MiniMap::new(),
			_session: session.inner()
		})
	}

	/// Bind a [`Value`] to a session input.
	///
	/// The value given here will be passed as the input specified by `name` in all future invocations of
	/// [`Session::run_binding`] until this input is overridden (by calling [`IoBinding::bind_input`] again) or until
	/// all inputs are cleared (via [`IoBinding::clear_inputs`] or [`IoBinding::clear`]).
	///
	/// If the session wants the input to be placed on a different device than the one `value` is on (as is often the
	/// case for CPU inputs to CUDA sessions), the data will be queued to be asynchronously copied to said device.
	/// Otherwise, and if `value` was created from borrowed data (i.e. with
	/// [`TensorRef::from_array_view`](crate::value::TensorRef::from_array_view)), a **second clone will be made
	/// first**, so it is recommended to only use `IoBinding` when the inputs' devices differ from the session's
	/// device(s) (or all inputs are owned).
	///
	/// The data is only copied upon invocation of this function. Any changes to the given value afterwards will
	/// not affect the data seen by the session until the value is re-bound.
	pub fn bind_input<T: ValueTypeMarker + ?Sized, S: Into<String>>(&mut self, name: S, value: &Value<T>) -> Result<()> {
		let name: String = name.into();

		#[allow(unused_mut)]
		let mut may_alias = value.is_tensor() && !value.inner.flags.upgradable();
		// if the input wasnt created from borrowed data, then we're keeping the owned data alive via held_inputs, so no
		// need to worry about uaf from aliasing
		#[cfg(feature = "api-23")]
		if let Some(wanted_memory_info) = self._session.inputs_memory_info.get(&name)
			&& let Some(input_memory_info) = value.inner.memory_info.as_ref()
			&& (wanted_memory_info.allocation_device() != input_memory_info.allocation_device()
				|| wanted_memory_info.device_id() != input_memory_info.device_id())
		{
			// ort will never alias if the devices do not match
			may_alias = false;
		}

		let value_inner = if may_alias {
			debug_assert!(value.is_tensor());

			let cloned: DynTensor = unsafe { value.transmute_type_ref() }.clone();
			cloned.inner
		} else {
			Arc::clone(&value.inner)
		};

		let ptr = self.ptr_mut();
		with_cstr(name.as_bytes(), &|name| {
			ortsys![unsafe BindInput(ptr, name.as_ptr(), value_inner.ptr())?];
			Ok(())
		})?;
		self.held_inputs.insert(name, value_inner);
		Ok(())
	}

	/// Bind a session output to a pre-allocated [`Value`].
	///
	/// This allows for the pre-allocation and reuse of memory in the session output (see [`Tensor::new`]). Any
	/// subsequent runs via [`Session::run_binding`] will reuse the same tensor to store the output instead of creating
	/// a new one each time.
	///
	/// The output will be accessible in the value returned by [`Session::run_binding`], under the name specified by
	/// `name`.
	///
	/// [`Tensor::new`]: crate::value::Tensor::new
	pub fn bind_output<T: ValueTypeMarker + ?Sized, S: Into<String>>(&mut self, name: S, mut value: impl AsValueMut<T>) -> Result<()> {
		let inner = value.value_inner();
		if inner.is_backed() {
			// this is purely because i am too lazy to modify `run_binding` to have preallocated outputs share the
			// original value's backing
			return Err(Error::new_with_code(ErrorCode::InvalidArgument, "only tensors created with Tensor::new can be bound to outputs"));
		}

		let name: String = name.into();
		unsafe { self.bind_output_inner(name.as_bytes(), inner) }?;
		self.output_values.insert(name, Some(Arc::clone(inner)));
		Ok(())
	}

	pub(crate) unsafe fn bind_output_inner<S: AsRef<[u8]>>(&mut self, name: S, value: &ValueInner) -> Result<()> {
		let ptr = self.ptr_mut();
		with_cstr(name.as_ref(), &|name| {
			ortsys![unsafe BindOutput(ptr, name.as_ptr(), value.ptr())?];
			Ok(())
		})?;
		Ok(())
	}

	/// Bind a session output to a device which is specified by `mem_info`. Used when the shape of the output is
	/// unknown.
	///
	/// Note that the same output buffer will be reused across runs, so the output must be
	/// [`.clone()`](crate::value::Tensor::clone)d to avoid race conditions.
	///
	/// If the shape of the output is known, prefer [`IoBinding::bind_output`] instead.
	pub fn bind_output_to_device<S: Into<String>>(&mut self, name: S, mem_info: &MemoryInfo) -> Result<()> {
		let name: String = name.into();
		let ptr = self.ptr_mut();
		with_cstr(name.as_bytes(), &|name| {
			ortsys![unsafe BindOutputToDevice(ptr, name.as_ptr(), mem_info.ptr())?];
			Ok(())
		})?;
		self.output_values.insert(name, None);
		Ok(())
	}

	/// Clears all bound inputs specified by [`IoBinding::bind_input`].
	pub fn clear_inputs(&mut self) {
		ortsys![unsafe ClearBoundInputs(self.ptr_mut())];
		drop(self.held_inputs.drain());
	}
	/// Clears all bound outputs specified by [`IoBinding::bind_output`] or [`IoBinding::bind_output_to_device`].
	pub fn clear_outputs(&mut self) {
		ortsys![unsafe ClearBoundOutputs(self.ptr_mut())];
		drop(self.output_values.drain());
	}
	/// Clears both the bound inputs & outputs; equivalent to [`IoBinding::clear_inputs`] followed by
	/// [`IoBinding::clear_outputs`].
	pub fn clear(&mut self) {
		self.clear_inputs();
		self.clear_outputs();
	}

	// technically the synchronize methods do not have to be &mut like the other methods, as they do not mutate any
	// state on the C side - all they do is iterate through all nodes' EPs and call their device synchronize
	// function (cudaDeviceSynchronize/hipDeviceSynchronize), which I assume would be thread-safe

	/// Synchronize all bound inputs, ensuring any pending asynchronous transfers are completed.
	pub fn synchronize_inputs(&self) -> Result<()> {
		ortsys![unsafe SynchronizeBoundInputs(self.ptr().cast_mut())?];
		Ok(())
	}
	/// Synchronize all bound outputs, ensuring any pending asynchronous transfers are completed.
	pub fn synchronize_outputs(&self) -> Result<()> {
		ortsys![unsafe SynchronizeBoundOutputs(self.ptr().cast_mut())?];
		Ok(())
	}
	/// Synchronizes both inputs & outputs; equivalent to [`IoBinding::synchronize_inputs`] followed by
	/// [`IoBinding::synchronize_outputs`].
	pub fn synchronize(&self) -> Result<()> {
		self.synchronize_inputs()?;
		self.synchronize_outputs()?;
		Ok(())
	}
}

unsafe impl Send for IoBinding {}

impl AsPointer for IoBinding {
	type Sys = ort_sys::OrtIoBinding;

	fn ptr(&self) -> *const Self::Sys {
		self.ptr.as_ptr()
	}
}

impl Drop for IoBinding {
	fn drop(&mut self) {
		ortsys![unsafe ReleaseIoBinding(self.ptr_mut())];
		crate::logging::drop!(IoBinding, self.ptr());
	}
}

#[cfg(test)]
mod tests {
	#[cfg(feature = "ndarray")]
	use crate::test_util::{mnist, test_env};
	use crate::{
		Result,
		memory::{AllocationDevice, Allocator, AllocatorType, MemoryInfo, MemoryType},
		session::Session,
		value::Tensor
	};

	// not terribly useful since CI is CPU-only, but it at least ensures the API won't segfault or something silly
	#[test]
	#[cfg(all(feature = "ndarray", feature = "fetch-models"))]
	fn test_mnist_input_bound() -> Result<()> {
		let mut session = Session::builder(test_env())?.commit_from_url(mnist::MODEL_URL)?;

		let array = mnist::get_image();

		let mut binding = session.create_binding()?;
		binding.bind_input(session.inputs()[0].name(), &Tensor::from_array(array)?)?;
		binding
			.bind_output_to_device(session.outputs()[0].name(), &MemoryInfo::new(AllocationDevice::CPU, 0, AllocatorType::Device, MemoryType::CPUOutput)?)?;

		let outputs = session.run_binding(&binding)?;
		let probabilities = mnist::extract_probabilities(&outputs[0])?;
		assert_eq!(probabilities[0].0, 5);

		Ok(())
	}

	#[test]
	#[cfg(all(feature = "ndarray", feature = "fetch-models"))]
	fn test_mnist_input_output_bound() -> Result<()> {
		let mut session = Session::builder(test_env())?.commit_from_url(mnist::MODEL_URL)?;

		let array = mnist::get_image();

		let mut binding = session.create_binding()?;
		binding.bind_input(session.inputs()[0].name(), &Tensor::from_array(array)?)?;

		binding.bind_output(session.outputs()[0].name(), Tensor::<f32>::new(&Allocator::default(), [1i64, 10])?)?;

		let outputs = session.run_binding(&binding)?;
		let probabilities = mnist::extract_probabilities(&outputs[0])?;
		assert_eq!(probabilities[0].0, 5);

		Ok(())
	}

	#[test]
	#[cfg(all(feature = "ndarray", feature = "fetch-models"))]
	fn test_send_iobinding() -> Result<()> {
		let mut session = Session::builder(test_env())?.commit_from_url(mnist::MODEL_URL)?;

		let array = mnist::get_image();

		let mut binding = session.create_binding()?;
		binding.bind_output(session.outputs()[0].name(), Tensor::<f32>::new(&Allocator::default(), [1i64, 10])?)?;

		let probabilities = std::thread::spawn(move || {
			binding.bind_input(session.inputs()[0].name(), &Tensor::from_array(array)?)?;
			let outputs = session.run_binding(&binding)?;
			let probabilities = mnist::extract_probabilities(&outputs[0])?;
			Ok::<Vec<(usize, f32)>, crate::Error>(probabilities)
		})
		.join()
		.expect("")?;

		assert_eq!(probabilities[0].0, 5);

		Ok(())
	}

	#[test]
	#[cfg(all(feature = "ndarray", feature = "fetch-models"))]
	fn test_mnist_clear_binds() -> Result<()> {
		let mut session = Session::builder(test_env())?.commit_from_url(mnist::MODEL_URL)?;

		let array = mnist::get_image();

		let mut binding = session.create_binding()?;
		binding.bind_input(session.inputs()[0].name(), &Tensor::from_array(array)?)?;

		binding.bind_output(session.outputs()[0].name(), Tensor::<f32>::new(&Allocator::default(), [1i64, 10])?)?;

		{
			let outputs = session.run_binding(&binding)?;
			let probabilities = mnist::extract_probabilities(&outputs[0])?;
			assert_eq!(probabilities[0].0, 5);
		}

		binding.clear_outputs();
		binding
			.bind_output_to_device(session.outputs()[0].name(), &MemoryInfo::new(AllocationDevice::CPU, 0, AllocatorType::Device, MemoryType::CPUOutput)?)?;

		{
			let outputs = session.run_binding(&binding)?;
			let probabilities = mnist::extract_probabilities(&outputs[0])?;
			assert_eq!(probabilities[0].0, 5);
		}

		binding.clear_inputs();
		assert!(session.run_binding(&binding).is_err());

		Ok(())
	}
}