ort 2.0.0-rc.15

//! ONNX Runtime doesn't (currently) expose an API for inter-device copies, so we instead use a dummy model to copy the
//! tensor & `IoBinding` to configure where the copy ends up.

use alloc::{
	format,
	string::{String, ToString}
};
use core::ops::{Deref, DerefMut};

use crate::{
	Error, OnceLock, Result,
	environment::Environment,
	ep,
	memory::{AllocationDevice, AllocatorType, MemoryInfo, MemoryType},
	session::{IoBinding, NoSelectedOutputs, RunOptions, Session, builder::GraphOptimizationLevel},
	util::{MiniMap, Mutex, MutexGuard},
	value::{AsValueMut, DefiniteTensorValueTypeMarker, DynTensorRef, TensorElementType, Value}
};

#[derive(Debug, Clone, PartialEq, Eq)]
struct IdentitySessionKey {
	src_device: String,
	src_device_id: i32,
	target_device: String,
	target_device_id: i32,
	dtype: ort_sys::ONNXTensorElementDataType
}

struct IdentitySession {
	session: Session,
	binding: IoBinding
}

/// A simple graph in `.ort` format with a single `Identity` node between the input & output.
static IDENTITY_MODEL: &[u8] = include_bytes!("./identity.ort");
static SESSIONS: OnceLock<Mutex<MiniMap<IdentitySessionKey, IdentitySession>>> = OnceLock::new();
/// `RunOptions` with [`RunOptions::disable_device_sync`], shared across `to_async()` calls to reduce allocations.
static ASYNC_RUN_OPTIONS: OnceLock<RunOptions<NoSelectedOutputs>> = OnceLock::new();

fn ep_for_device(device: &AllocationDevice, device_id: i32) -> Result<ep::ExecutionProviderDispatch> {
	if *device == AllocationDevice::CPU {
		return Ok(ep::CPU::default().with_arena_allocator(false).build());
	}
	#[cfg(feature = "cuda")]
	if *device == AllocationDevice::CUDA || *device == AllocationDevice::CUDA_PINNED {
		return Ok(ep::CUDA::default()
			.with_device_id(device_id)
			.with_arena_extend_strategy(ep::ArenaExtendStrategy::SameAsRequested)
			.with_conv_max_workspace(false)
			.with_conv_algorithm_search(ep::cuda::ConvAlgorithmSearch::Default)
			.build());
	};
	#[cfg(feature = "directml")]
	if *device == AllocationDevice::DIRECTML {
		return Ok(ep::DirectML::default().with_device_id(device_id).build());
	}
	#[cfg(feature = "cann")]
	if *device == AllocationDevice::CANN || *device == AllocationDevice::CANN_PINNED {
		return Ok(ep::CANN::default()
			.with_arena_extend_strategy(ep::ArenaExtendStrategy::SameAsRequested)
			.with_cann_graph(false)
			.with_device_id(device_id)
			.build());
	}
	#[cfg(feature = "openvino")]
	if *device == AllocationDevice::OPENVINO_CPU || *device == AllocationDevice::OPENVINO_GPU {
		return Ok(ep::OpenVINO::default()
			.with_num_threads(1)
			.with_device_type(if *device == AllocationDevice::OPENVINO_CPU {
				"CPU".to_string()
			} else {
				format!("GPU.{device_id}")
			})
			.build());
	}
	#[cfg(any(feature = "migraphx", feature = "rocm"))]
	#[allow(unreachable_code)]
	if *device == AllocationDevice::HIP || *device == AllocationDevice::HIP_PINNED {
		#[cfg(feature = "migraphx")]
		return Ok(ep::MIGraphX::default()
			.with_arena_extend_strategy(ep::ArenaExtendStrategy::SameAsRequested)
			.build());
		#[cfg(feature = "rocm")]
		return Ok(ep::ROCm::default()
			.with_arena_extend_strategy(ep::ArenaExtendStrategy::SameAsRequested)
			.with_hip_graph(false)
			.with_exhaustive_conv_search(false)
			.with_device_id(device_id)
			.build());
	}

	Err(crate::Error::new(format!("Unsupported allocation device `{}` for tensor copy target", device.as_str())))
}

impl<Type: DefiniteTensorValueTypeMarker + ?Sized> Value<Type> {
	/// Copies the contents of this tensor to another device, returning the newly created tensor value.
	///
	/// ```
	/// # use ort::{memory::{Allocator, AllocatorType, AllocationDevice, MemoryInfo, MemoryType}, session::Session, value::Tensor};
	/// # fn main() -> ort::Result<()> {
	/// # let env = ort::test_util::test_env().clone();
	/// # if false {
	/// # let session = Session::builder(&env)?.commit_from_file("tests/data/upsample.onnx")?;
	/// let cuda_allocator = Allocator::new(
	/// 	&session,
	/// 	MemoryInfo::new(AllocationDevice::CUDA, 0, AllocatorType::Device, MemoryType::Default)?
	/// )?;
	/// let cuda_tensor = Tensor::<f32>::new(&cuda_allocator, [1_usize, 3, 224, 224])?;
	/// # }
	/// # let cuda_tensor = Tensor::<f32>::new(&Allocator::default(), [1_usize, 3, 224, 224])?;
	///
	/// let cpu_tensor = cuda_tensor.to(AllocationDevice::CPU, 0)?;
	/// assert_eq!(cpu_tensor.memory_info().allocation_device(), AllocationDevice::CPU);
	/// assert_eq!(**cpu_tensor.shape(), [1, 3, 224, 224]);
	/// # Ok(())
	/// # }
	/// ```
	#[cfg_attr(docsrs, doc(cfg(not(all(target_family = "wasm", target_os = "unknown")))))]
	pub fn to(&self, device: AllocationDevice, device_id: i32) -> Result<Value<Type>> {
		let mut identity_session = IdentitySessionHandle::new(self.upcast_ref(), &device, device_id, *self.data_type())?;

		let target_memory_info = MemoryInfo::new(device, device_id, AllocatorType::Device, MemoryType::Default)?;
		identity_session.binding.bind_output_to_device("output", &target_memory_info)?;

		let output = {
			let identity_session = &mut *identity_session;
			identity_session
				.session
				.run_binding(&identity_session.binding)?
				.remove("output")
				.expect("identity model should have single output")
		};

		Ok(unsafe { output.transmute_type() })
	}

	/// Asynchronously copies the contents of this tensor to another device.
	///
	/// Unlike [`Tensor::to`][crate::value::Tensor::to], the device's stream will *not* be synchronized (like via
	/// `cudaStreamSynchronize`); thus this function is most useful for host-to-device transfers.
	///
	/// ```
	/// # use ort::{memory::{Allocator, AllocatorType, AllocationDevice, MemoryInfo, MemoryType}, session::Session, value::Tensor};
	/// # fn main() -> ort::Result<()> {
	/// # let env = ort::test_util::test_env().clone();
	/// # let tensor = Tensor::<f32>::new(&Allocator::default(), [1_usize, 3, 224, 224])?;
	/// # if false {
	/// let cuda_tensor = tensor.to_async(AllocationDevice::CUDA, 0)?;
	/// // pass to other CUDA code, or to session input
	/// # }
	/// # Ok(())
	/// # }
	/// ```
	#[cfg_attr(docsrs, doc(cfg(not(all(target_family = "wasm", target_os = "unknown")))))]
	pub fn to_async(&self, device: AllocationDevice, device_id: i32) -> Result<Value<Type>> {
		let mut identity_session = IdentitySessionHandle::new(self.upcast_ref(), &device, device_id, *self.data_type())?;

		let options = ASYNC_RUN_OPTIONS.get_or_try_init(|| -> Result<RunOptions> {
			let mut options = RunOptions::new()?;
			options.disable_device_sync()?;
			Ok(options)
		})?;
		let target_memory_info = MemoryInfo::new(device, device_id, AllocatorType::Device, MemoryType::Default)?;
		identity_session.binding.bind_output_to_device("output", &target_memory_info)?;
		let output = {
			let identity_session = &mut *identity_session;
			identity_session
				.session
				.run_binding_with_options(&identity_session.binding, options)?
				.remove("output")
				.expect("identity model should have single output")
		};

		Ok(unsafe { output.transmute_type() })
	}

	/// Copies the contents of this tensor to another tensor potentially residing on a separate device.
	///
	/// ```
	/// # use ort::{memory::{Allocator, AllocatorType, AllocationDevice, MemoryInfo, MemoryType}, session::Session, value::Tensor};
	/// # fn main() -> ort::Result<()> {
	/// # let env = ort::test_util::test_env().clone();
	/// # if false {
	/// # let session = Session::builder(&env)?.commit_from_file("tests/data/upsample.onnx")?;
	/// let cuda_allocator = Allocator::new(
	/// 	&session,
	/// 	MemoryInfo::new(AllocationDevice::CUDA, 0, AllocatorType::Device, MemoryType::Default)?
	/// )?;
	/// let cuda_tensor = Tensor::<f32>::new(&cuda_allocator, [1_usize, 3, 224, 224])?;
	/// # }
	/// # let cuda_tensor = Tensor::<f32>::new(&Allocator::default(), [1_usize, 3, 224, 224])?;
	/// let mut cpu_tensor = Tensor::<f32>::new(&Allocator::default(), [1_usize, 3, 224, 224])?;;
	///
	/// cuda_tensor.copy_into(&mut cpu_tensor)?;
	/// # Ok(())
	/// # }
	/// ```
	#[cfg_attr(docsrs, doc(cfg(not(all(target_family = "wasm", target_os = "unknown")))))]
	pub fn copy_into(&self, mut target: impl AsValueMut<Type>) -> Result<()> {
		let target_inner = target.value_inner();
		if self.dtype() != &target_inner.dtype {
			return Err(Error::new("target data type or shape does not match source data type or shape"));
		}

		// SAFETY: tensors always have memory info
		let target_memory_info = unsafe { target_inner.memory_info.as_ref().unwrap_unchecked() };
		let mut identity_session =
			IdentitySessionHandle::new(self.upcast_ref(), &target_memory_info.allocation_device(), target_memory_info.device_id(), *self.data_type())?;

		unsafe { identity_session.binding.bind_output_inner("output", target_inner) }?;
		{
			let identity_session = &mut *identity_session;
			identity_session.session.run_binding(&identity_session.binding)?;
		}

		Ok(())
	}

	/// Asynchronously copies the contents of this tensor to another tensor.
	///
	/// Unlike [`Tensor::copy_into`][crate::value::Tensor::copy_into], the device's stream will *not* be synchronized
	/// (like via `cudaStreamSynchronize`); thus this function is most useful for host-to-device transfers.
	///
	/// ```
	/// # use ort::{memory::{Allocator, AllocatorType, AllocationDevice, MemoryInfo, MemoryType}, session::Session, value::Tensor};
	/// # fn main() -> ort::Result<()> {
	/// # let env = ort::test_util::test_env().clone();
	/// let cpu_tensor = Tensor::<f32>::new(&Allocator::default(), [1_usize, 3, 224, 224])?;;
	/// # if false {
	/// # let session = Session::builder(&env)?.commit_from_file("tests/data/upsample.onnx")?;
	/// let cuda_allocator = Allocator::new(
	/// 	&session,
	/// 	MemoryInfo::new(AllocationDevice::CUDA, 0, AllocatorType::Device, MemoryType::Default)?
	/// )?;
	/// let mut cuda_tensor = Tensor::<f32>::new(&cuda_allocator, [1_usize, 3, 224, 224])?;
	/// # }
	/// # let mut cuda_tensor = Tensor::<f32>::new(&Allocator::default(), [1_usize, 3, 224, 224])?;
	///
	/// cpu_tensor.copy_into_async(&mut cuda_tensor)?;
	/// # Ok(())
	/// # }
	/// ```
	#[cfg_attr(docsrs, doc(cfg(not(all(target_family = "wasm", target_os = "unknown")))))]
	pub fn copy_into_async(&self, target: &mut Value<Type>) -> Result<()> {
		let target_inner = target.value_inner();
		if self.dtype() != &target_inner.dtype {
			return Err(Error::new("target data type or shape does not match source data type or shape"));
		}

		// SAFETY: tensors always have memory info
		let target_memory_info = unsafe { target_inner.memory_info.as_ref().unwrap_unchecked() };
		let mut identity_session =
			IdentitySessionHandle::new(self.upcast_ref(), &target_memory_info.allocation_device(), target_memory_info.device_id(), *self.data_type())?;

		let options = ASYNC_RUN_OPTIONS.get_or_try_init(|| -> Result<RunOptions> {
			let mut options = RunOptions::new()?;
			options.disable_device_sync()?;
			Ok(options)
		})?;
		unsafe { identity_session.binding.bind_output_inner("output", target_inner) }?;
		{
			let identity_session = &mut *identity_session;
			identity_session.session.run_binding_with_options(&identity_session.binding, options)?;
		}

		Ok(())
	}
}

#[cfg_attr(docsrs, doc(cfg(not(all(target_family = "wasm", target_os = "unknown")))))]
impl<Type: DefiniteTensorValueTypeMarker + ?Sized> Clone for Value<Type> {
	/// Creates a copy of this tensor and its data on the same device it resides on.
	///
	/// ```
	/// # use ort::{value::Tensor, AsPointer};
	/// # fn main() -> ort::Result<()> {
	/// # let _env = ort::test_util::test_env().clone();
	/// let array = vec![1_i64, 2, 3, 4, 5];
	/// let tensor = Tensor::from_array(([array.len()], array.into_boxed_slice()))?;
	///
	/// let new_tensor = tensor.clone();
	///
	/// // same data
	/// assert_eq!(tensor.extract_tensor(), new_tensor.extract_tensor());
	///
	/// // different allocations
	/// assert_ne!(tensor.ptr(), new_tensor.ptr());
	/// assert_ne!(tensor.data_ptr(), new_tensor.data_ptr());
	/// # 	Ok(())
	/// # }
	/// ```
	#[cfg_attr(docsrs, doc(cfg(not(all(target_family = "wasm", target_os = "unknown")))))]
	fn clone(&self) -> Self {
		let memory_info = self.memory_info();
		self.to(memory_info.allocation_device(), memory_info.device_id())
			.expect("Failed to clone tensor")
	}
}

struct IdentitySessionHandle<'s> {
	inner: &'static mut IdentitySession,
	_src: DynTensorRef<'s>,
	_guard: MutexGuard<'static, MiniMap<IdentitySessionKey, IdentitySession>>
}

impl<'s> IdentitySessionHandle<'s> {
	fn new(src: DynTensorRef<'s>, target_device: &AllocationDevice<'_>, target_device_id: i32, dtype: TensorElementType) -> Result<Self> {
		let tensor_type = ort_sys::ONNXTensorElementDataType::from(dtype);
		let src_memory_info = src.memory_info();
		let session_key = IdentitySessionKey {
			src_device: src_memory_info.allocation_device().as_str().to_string(),
			src_device_id: src_memory_info.device_id(),
			target_device: target_device.as_str().to_string(),
			target_device_id,
			dtype: tensor_type
		};
		let mut sessions = SESSIONS.get_or_init(|| Mutex::new(MiniMap::new())).lock();
		let identity_session = match sessions.get_mut(&session_key) {
			Some(entry) => entry,
			None => {
				let mut model_bytes = IDENTITY_MODEL.to_vec();
				// Override the expected element type of the input & output nodes, respectively.
				model_bytes[544] = tensor_type as u8;
				model_bytes[604] = tensor_type as u8;

				let (source_ep, target_ep) = (
					// We enable `.error_on_failure()` here since `IoBinding::bind_output_to_device` will silently fall back to binding to CPU if the target
					// device doesn't have an EP registered.
					ep_for_device(&src_memory_info.allocation_device(), src_memory_info.device_id())?.error_on_failure(),
					ep_for_device(target_device, target_device_id)?.error_on_failure()
				);

				let env = Environment::current().ok_or_else(|| Error::new("an environment must be active to copy tensors"))?;

				let mut builder = Session::builder(&env)?
					.with_optimization_level(GraphOptimizationLevel::Level3)?
					// since these sessions are persistent for the lifetime of the program, keep them as lean as possible
					// by disabling threading & memory optimizations (there's only 1 operation in the graph anyway)
					.with_intra_threads(1)?
					.with_inter_threads(1)?
					.with_inter_op_spinning(false)?
					.with_intra_op_spinning(false)?
					.with_memory_pattern(false)?
					.with_no_environment_execution_providers()?;
				// Avoid registering the same EP twice, since that's an error.
				if source_ep.name() != target_ep.name() {
					builder = builder.with_execution_providers([source_ep, target_ep])?;
				} else {
					builder = builder.with_execution_providers([source_ep])?;
				}

				let session = builder.commit_from_memory(&model_bytes)?;
				let binding = session.create_binding()?;
				sessions.insert(session_key.clone(), IdentitySession { session, binding });
				sessions.get_mut(&session_key).expect("insert should have worked")
			}
		};

		identity_session.binding.bind_input("input", &src)?;

		Ok(Self {
			inner: unsafe { core::mem::transmute::<&mut IdentitySession, &'static mut IdentitySession>(identity_session) },
			_src: src,
			_guard: sessions
		})
	}
}

impl Deref for IdentitySessionHandle<'_> {
	type Target = IdentitySession;
	fn deref(&self) -> &Self::Target {
		self.inner
	}
}

impl DerefMut for IdentitySessionHandle<'_> {
	fn deref_mut(&mut self) -> &mut Self::Target {
		self.inner
	}
}

impl Drop for IdentitySessionHandle<'_> {
	fn drop(&mut self) {
		// dont let the binding continue to hold onto the source or target
		self.binding.clear();
	}
}

#[cfg(test)]
mod tests {
	use crate::{
		memory::{AllocationDevice, Allocator},
		test_util::test_env,
		value::Tensor
	};

	#[test]
	fn test_clone_tensor() -> crate::Result<()> {
		let _env = test_env();
		let tensor = Tensor::<f32>::from_array(([1, 5], vec![2.167892, 333., 1.0, -0.0, f32::EPSILON]))?;
		let clone = tensor.clone();
		assert_eq!(tensor.extract_tensor(), clone.extract_tensor());
		Ok(())
	}

	#[test]
	fn test_tensor_to_cpu() -> crate::Result<()> {
		let _env = test_env();
		let tensor = Tensor::<f32>::from_array(([1, 5], vec![2.167892, 333., 1.0, -0.0, f32::EPSILON]))?;

		let tensor2 = tensor.to_async(AllocationDevice::CPU, 0)?;
		let tensor3 = tensor2.to(AllocationDevice::CPU, 0)?;
		assert_eq!(tensor.extract_tensor(), tensor3.extract_tensor());

		Ok(())
	}

	#[test]
	fn test_bad_copy_into() -> crate::Result<()> {
		let _env = test_env();
		let tensor = Tensor::<f32>::from_array(([1, 5], vec![2.167892, 333., 1.0, -0.0, f32::EPSILON]))?;

		let mut target_bad_shape = Tensor::<f32>::new(&Allocator::default(), [3i64])?;
		assert!(tensor.copy_into(&mut target_bad_shape).is_err());

		let mut target_f64 = Tensor::<f64>::new(&Allocator::default(), [1i64, 5])?.upcast();
		assert!(tensor.upcast().copy_into(&mut target_f64).is_err());

		Ok(())
	}

	#[test]
	#[cfg(feature = "cuda")]
	fn test_copy_cuda() -> crate::Result<()> {
		let _env = test_env();
		let tensor = Tensor::<f32>::from_array(([1, 5], vec![2.167892, 333., 1.0, -0.0, f32::EPSILON]))?;

		let cuda_tensor = tensor.to(AllocationDevice::CUDA, 0)?;
		let memory = cuda_tensor.memory_info();
		assert_eq!(memory.allocation_device(), AllocationDevice::CUDA);
		assert!(!memory.is_cpu_accessible());

		let cpu_tensor = cuda_tensor.to(AllocationDevice::CPU, 0)?;
		assert!(cpu_tensor.memory_info().is_cpu_accessible());

		assert_eq!(tensor.extract_tensor(), cpu_tensor.extract_tensor());

		Ok(())
	}

	#[test]
	#[cfg(feature = "cuda")]
	fn test_copy_cuda_async() -> crate::Result<()> {
		let _env = test_env();
		let tensor = Tensor::<f32>::from_array(([1, 5], vec![2.167892, 333., 1.0, -0.0, f32::EPSILON]))?;

		let cuda_tensor = tensor.to_async(AllocationDevice::CUDA, 0)?;
		let cpu_tensor = cuda_tensor.to(AllocationDevice::CPU, 0)?;
		assert_eq!(tensor.extract_tensor(), cpu_tensor.extract_tensor());

		Ok(())
	}

	#[test]
	#[cfg(feature = "cuda")]
	fn test_copy_into_cuda() -> crate::Result<()> {
		use crate::{
			ep,
			memory::{AllocatorType, MemoryInfo, MemoryType},
			session::Session
		};

		let dummy_session = Session::builder(test_env())?
			.with_execution_providers([ep::CUDA::default().build()])?
			.commit_from_file("tests/data/upsample.ort")?;

		let allocator = Allocator::new(&dummy_session, MemoryInfo::new(AllocationDevice::CUDA, 0, AllocatorType::Device, MemoryType::Default)?)?;
		let tensor = Tensor::<f32>::from_array(([1, 5], vec![2.167892, 333., 1.0, -0.0, f32::EPSILON]))?;
		let mut cuda_tensor = Tensor::<f32>::new(&allocator, [1_i64, 5])?;

		tensor.copy_into(&mut cuda_tensor)?;
		let cpu_tensor = cuda_tensor.to(AllocationDevice::CPU, 0)?;
		assert_eq!(tensor.extract_tensor(), cpu_tensor.extract_tensor());

		Ok(())
	}

	#[test]
	#[cfg(feature = "cuda")]
	fn test_copy_into_async_cuda() -> crate::Result<()> {
		use crate::{
			ep,
			memory::{AllocatorType, MemoryInfo, MemoryType},
			session::Session
		};

		let dummy_session = Session::builder(test_env())?
			.with_execution_providers([ep::CUDA::default().build()])?
			.commit_from_file("tests/data/upsample.ort")?;

		let allocator = Allocator::new(&dummy_session, MemoryInfo::new(AllocationDevice::CUDA, 0, AllocatorType::Device, MemoryType::Default)?)?;
		let tensor = Tensor::<f32>::from_array(([1, 5], vec![2.167892, 333., 1.0, -0.0, f32::EPSILON]))?;
		let mut cuda_tensor = Tensor::<f32>::new(&allocator, [1_i64, 5])?;

		tensor.copy_into_async(&mut cuda_tensor)?;
		let cpu_tensor = cuda_tensor.to(AllocationDevice::CPU, 0)?;
		assert_eq!(tensor.extract_tensor(), cpu_tensor.extract_tensor());

		Ok(())
	}
}