ort 2.0.0-rc.13

A safe Rust wrapper for ONNX Runtime 1.28 - Optimize and accelerate machine learning inference & training
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

use core::ptr;

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

/// [oneDNN/DNNL execution provider](https://onnxruntime.ai/docs/execution-providers/oneDNN-ExecutionProvider.html) for
/// Intel CPUs & iGPUs.
#[derive(Debug, Default, Clone)]
#[doc(alias = "DNNL")]
pub struct OneDNN {
	options: ExecutionProviderOptions
}

super::impl_ep!(arbitrary; OneDNN);

impl OneDNN {
	/// Enable/disable the usage of the arena allocator.
	///
	/// ```
	/// # use ort::{ep, session::Session};
	/// # fn main() -> ort::Result<()> {
	/// let ep = ep::OneDNN::default().with_arena_allocator(true).build();
	/// # Ok(())
	/// # }
	/// ```
	#[must_use]
	pub fn with_arena_allocator(mut self, enable: bool) -> Self {
		self.options.set("use_arena", if enable { "1" } else { "0" });
		self
	}
}

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

	fn register(&self, session_builder: &mut SessionBuilder) -> Result<()> {
		let mut dnnl_options: *mut ort_sys::OrtDnnlProviderOptions = ptr::null_mut();
		ortsys![unsafe CreateDnnlProviderOptions(&mut dnnl_options)?];
		let _guard = util::run_on_drop(|| {
			ortsys![unsafe ReleaseDnnlProviderOptions(dnnl_options)];
		});

		let ffi_options = self.options.to_ffi();
		ortsys![unsafe UpdateDnnlProviderOptions(
				dnnl_options,
				ffi_options.key_ptrs(),
				ffi_options.value_ptrs(),
				ffi_options.len()
			)?];

		ortsys![unsafe SessionOptionsAppendExecutionProvider_Dnnl(session_builder.ptr_mut(), dnnl_options)?];

		Ok(())
	}
}