ort 2.0.0-rc.13

A safe Rust wrapper for ONNX Runtime 1.28 - Optimize and accelerate machine learning inference & training
Documentation
use alloc::string::ToString;

use super::{ExecutionProvider, ExecutionProviderOptions};
use crate::{AsPointer, error::Result, ortsys, session::builder::SessionBuilder};

/// [OpenVINO execution provider](https://onnxruntime.ai/docs/execution-providers/OpenVINO-ExecutionProvider.html) for
/// Intel CPUs/GPUs/NPUs.
#[derive(Default, Debug, Clone)]
pub struct OpenVINO {
	options: ExecutionProviderOptions
}

super::impl_ep!(arbitrary; OpenVINO);

impl OpenVINO {
	/// Overrides the accelerator hardware type and precision.
	///
	/// `device_type` should be in the format `CPU`, `NPU`, `GPU`, `GPU.0`, `GPU.1`, etc. Heterogenous combinations are
	/// also supported in the format `HETERO:NPU,GPU`.
	///
	/// ```
	/// # use ort::{ep, session::Session};
	/// # fn main() -> ort::Result<()> {
	/// let ep = ep::OpenVINO::default().with_device_type("GPU.0").build();
	/// # Ok(())
	/// # }
	/// ```
	#[must_use]
	pub fn with_device_type(mut self, device_type: impl AsRef<str>) -> Self {
		self.options.set("device_type", device_type.as_ref());
		self
	}

	/// Overrides the accelerator default value of number of threads with this value at runtime. If this option is not
	/// explicitly set, default value of 8 is used during build time.
	#[must_use]
	pub fn with_num_threads(mut self, num_threads: usize) -> Self {
		self.options.set("num_of_threads", num_threads.to_string());
		self
	}

	/// Explicitly specify the path to save and load the blobs, enabling model caching.
	#[must_use]
	pub fn with_cache_dir(mut self, dir: impl AsRef<str>) -> Self {
		self.options.set("cache_dir", dir.as_ref());
		self
	}

	/// This option enables OpenCL queue throttling for GPU devices (reduces CPU utilization when using GPU).
	#[must_use]
	pub fn with_opencl_throttling(mut self, enable: bool) -> Self {
		self.options.set("enable_opencl_throttling", if enable { "true" } else { "false" });
		self
	}

	#[must_use]
	pub fn with_qdq_optimizer(mut self, enable: bool) -> Self {
		self.options.set("enable_qdq_optimizer", if enable { "true" } else { "false" });
		self
	}

	/// This option if enabled works for dynamic shaped models whose shape will be set dynamically based on the infer
	/// input image/data shape at run time in CPU. This gives best result for running multiple inferences with varied
	/// shaped images/data.
	#[must_use]
	pub fn with_dynamic_shapes(mut self, enable: bool) -> Self {
		self.options.set("disable_dynamic_shapes", if enable { "false" } else { "true" });
		self
	}

	#[must_use]
	pub fn with_num_streams(mut self, num_streams: u8) -> Self {
		self.options.set("num_streams", num_streams.to_string());
		self
	}

	#[must_use]
	pub fn with_precision(mut self, precision: impl AsRef<str>) -> Self {
		self.options.set("precision", precision.as_ref());
		self
	}
}

impl ExecutionProvider for OpenVINO {
	fn name(&self) -> &'static str {
		"OpenVINOExecutionProvider"
	}

	fn register(&self, session_builder: &mut SessionBuilder) -> Result<()> {
		let ffi_options = self.options.to_ffi();
		ortsys![unsafe SessionOptionsAppendExecutionProvider_OpenVINO_V2(
			session_builder.ptr_mut(),
			ffi_options.key_ptrs(),
			ffi_options.value_ptrs(),
			ffi_options.len()
		)?];
		Ok(())
	}
}