diff --git a/crates/openvino/src/core.rs b/crates/openvino/src/core.rs index 577cac7..7a8acb2 100644 --- a/crates/openvino/src/core.rs +++ b/crates/openvino/src/core.rs @@ -4,16 +4,21 @@ use crate::error::LoadingError; use crate::{cstr, drop_using_function, try_unsafe, util::Result}; use crate::{model::CompiledModel, Model}; -use crate::{DeviceType, SetupError, Tensor}; +use crate::{DeviceType, PropertyKey, RwPropertyKey, SetupError, Tensor, Version}; use openvino_sys::{ - self, ov_core_compile_model, ov_core_create, ov_core_create_with_config, ov_core_free, - ov_core_read_model, ov_core_read_model_from_memory_buffer, ov_core_t, + self, ov_available_devices_free, ov_core_compile_model, ov_core_create, + ov_core_create_with_config, ov_core_free, ov_core_get_available_devices, ov_core_get_property, + ov_core_get_versions_by_device_name, ov_core_read_model, ov_core_read_model_from_memory_buffer, + ov_core_set_property, ov_core_t, ov_core_versions_free, }; -use std::ffi::CString; - +use std::ffi::{CStr, CString}; use std::os::raw::c_char; +use std::slice; +use std::str::FromStr; + +const EMPTY_C_STR: &CStr = unsafe { CStr::from_bytes_with_nul_unchecked(b"\0") }; -/// See [`Core`](https://docs.openvino.ai/2023.3/api/c_cpp_api/group__ov__core__c__api.html). +/// See [`Core`](https://docs.openvino.ai/2024/api/c_cpp_api/group__ov__core__c__api.html). pub struct Core { ptr: *mut ov_core_t, } @@ -40,6 +45,155 @@ impl Core { Ok(Core { ptr }) } + /// Gets device plugins version information. + /// Device name can be complex and identify multiple devices at once like `HETERO:CPU,GPU`; + /// in this case, the returned map contains multiple entries, each per device. + pub fn versions(&self, device_name: &str) -> Result> { + let device_name = cstr!(device_name); + let mut ov_version_list = openvino_sys::ov_core_version_list_t { + versions: std::ptr::null_mut(), + size: 0, + }; + try_unsafe!(ov_core_get_versions_by_device_name( + self.ptr, + device_name, + std::ptr::addr_of_mut!(ov_version_list) + ))?; + + let ov_versions = + unsafe { slice::from_raw_parts(ov_version_list.versions, ov_version_list.size) }; + + let mut versions: Vec<(DeviceType, Version)> = Vec::with_capacity(ov_version_list.size); + for ov_version in ov_versions { + let c_str_device_name = unsafe { std::ffi::CStr::from_ptr(ov_version.device_name) }; + let device_name = c_str_device_name.to_string_lossy(); + let device_type = DeviceType::from_str(device_name.as_ref()).unwrap(); + versions.push((device_type, Version::from(&ov_version.version))); + } + + unsafe { ov_core_versions_free(std::ptr::addr_of_mut!(ov_version_list)) }; + Ok(versions) + } + + /// Gets devices available for inference. + pub fn available_devices(&self) -> Result> { + let mut ov_available_devices = openvino_sys::ov_available_devices_t { + devices: std::ptr::null_mut(), + size: 0, + }; + try_unsafe!(ov_core_get_available_devices( + self.ptr, + std::ptr::addr_of_mut!(ov_available_devices) + ))?; + + let ov_devices = unsafe { + slice::from_raw_parts(ov_available_devices.devices, ov_available_devices.size) + }; + + let mut devices = Vec::with_capacity(ov_available_devices.size); + for ov_device in ov_devices { + let c_str_device_name = unsafe { std::ffi::CStr::from_ptr(*ov_device) }; + let device_name = c_str_device_name.to_string_lossy(); + let device_type = DeviceType::from_str(device_name.as_ref()).unwrap(); + devices.push(device_type); + } + + unsafe { ov_available_devices_free(std::ptr::addr_of_mut!(ov_available_devices)) }; + Ok(devices) + } + + /// Gets properties related to this Core. + /// The method extracts information that can be set via the [set_property] method. + pub fn get_property(&self, key: PropertyKey) -> Result { + let ov_prop_key = cstr!(key.as_ref()); + let mut ov_prop_value = std::ptr::null_mut(); + try_unsafe!(ov_core_get_property( + self.ptr, + EMPTY_C_STR.as_ptr(), + ov_prop_key, + std::ptr::addr_of_mut!(ov_prop_value) + ))?; + let rust_prop = unsafe { CStr::from_ptr(ov_prop_value) } + .to_str() + .unwrap() + .to_owned(); + Ok(rust_prop) + } + + /// Sets a property for this Core instance. + pub fn set_property(&mut self, key: RwPropertyKey, value: &str) -> Result<()> { + let ov_prop_key = cstr!(key.as_ref()); + let ov_prop_value = cstr!(value); + try_unsafe!(ov_core_set_property( + self.ptr, + EMPTY_C_STR.as_ptr(), + ov_prop_key, + ov_prop_value, + ))?; + Ok(()) + } + + /// Sets properties for this Core instance. + pub fn set_properties<'a>( + &mut self, + properties: impl IntoIterator, + ) -> Result<()> { + for (prop_key, prop_value) in properties { + self.set_property(prop_key, prop_value)?; + } + Ok(()) + } + + /// Gets properties related to device behaviour. + /// The method extracts information that can be set via the [set_device_property] method. + pub fn get_device_property(&self, device_name: &str, key: PropertyKey) -> Result { + let ov_device_name = cstr!(device_name); + let ov_prop_key = cstr!(key.as_ref()); + let mut ov_prop_value = std::ptr::null_mut(); + try_unsafe!(ov_core_get_property( + self.ptr, + ov_device_name, + ov_prop_key, + std::ptr::addr_of_mut!(ov_prop_value) + ))?; + let rust_prop = unsafe { CStr::from_ptr(ov_prop_value) } + .to_str() + .unwrap() + .to_owned(); + Ok(rust_prop) + } + + /// Sets a property for a device. + pub fn set_device_property( + &mut self, + device_name: &str, + key: RwPropertyKey, + value: &str, + ) -> Result<()> { + let ov_device_name = cstr!(device_name); + let ov_prop_key = cstr!(key.as_ref()); + let ov_prop_value = cstr!(value); + try_unsafe!(ov_core_set_property( + self.ptr, + ov_device_name, + ov_prop_key, + ov_prop_value, + ))?; + Ok(()) + } + + /// Sets properties for a device. + pub fn set_device_properties<'a>( + &mut self, + device_name: &str, + properties: impl IntoIterator, + ) -> Result<()> { + for (prop_key, prop_value) in properties { + self.set_device_property(device_name, prop_key, prop_value)?; + } + Ok(()) + } + /// Read a Model from a pair of files: `model_path` points to an XML file containing the /// OpenVINO model IR and `weights_path` points to the binary weights file. pub fn read_model_from_file(&mut self, model_path: &str, weights_path: &str) -> Result {