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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

use alloc::string::ToString;
use core::{ffi::c_void, ptr};

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

#[derive(Debug, Default, Clone)]
pub struct ROCm {
	options: ExecutionProviderOptions
}

super::impl_ep!(arbitrary; ROCm);

impl ROCm {
	#[must_use]
	pub fn with_device_id(mut self, device_id: i32) -> Self {
		self.options.set("device_id", device_id.to_string());
		self
	}

	#[must_use]
	pub fn with_exhaustive_conv_search(mut self, enable: bool) -> Self {
		self.options.set("miopen_conv_exhaustive_search", if enable { "1" } else { "0" });
		self
	}

	#[must_use]
	pub fn with_conv_use_max_workspace(mut self, enable: bool) -> Self {
		self.options.set("miopen_conv_use_max_workspace", if enable { "1" } else { "0" });
		self
	}

	#[must_use]
	pub fn with_mem_limit(mut self, limit: usize) -> Self {
		self.options.set("gpu_mem_limit", limit.to_string());
		self
	}

	#[must_use]
	pub fn with_arena_extend_strategy(mut self, strategy: ArenaExtendStrategy) -> Self {
		self.options.set(
			"arena_extend_strategy",
			match strategy {
				ArenaExtendStrategy::NextPowerOfTwo => "kNextPowerOfTwo",
				ArenaExtendStrategy::SameAsRequested => "kSameAsRequested"
			}
		);
		self
	}

	#[must_use]
	pub fn with_copy_in_default_stream(mut self, enable: bool) -> Self {
		self.options.set("do_copy_in_default_stream", if enable { "1" } else { "0" });
		self
	}

	#[must_use]
	pub fn with_compute_stream(mut self, ptr: *mut c_void) -> Self {
		self.options.set("has_user_compute_stream", "1");
		self.options.set("user_compute_stream", (ptr as usize).to_string());
		self
	}

	#[must_use]
	pub fn with_hip_graph(mut self, enable: bool) -> Self {
		self.options.set("enable_hip_graph", if enable { "1" } else { "0" });
		self
	}

	#[must_use]
	pub fn with_tunable_op(mut self, enable: bool) -> Self {
		self.options.set("tunable_op_enable", if enable { "1" } else { "0" });
		self
	}

	#[must_use]
	pub fn with_tuning(mut self, enable: bool) -> Self {
		self.options.set("tunable_op_tuning_enable", if enable { "1" } else { "0" });
		self
	}

	#[must_use]
	pub fn with_max_tuning_duration(mut self, ms: i32) -> Self {
		self.options.set("tunable_op_max_tuning_duration_ms", ms.to_string());
		self
	}
}

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

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

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

		ortsys![unsafe SessionOptionsAppendExecutionProvider_ROCM(session_builder.ptr_mut(), rocm_options)?];

		Ok(())
	}
}