ort/ep/
rocm.rs

1use alloc::string::ToString;
2use core::{ffi::c_void, ptr};
3
4use super::{ArenaExtendStrategy, ExecutionProvider, ExecutionProviderOptions};
5use crate::{AsPointer, error::Result, ortsys, session::builder::SessionBuilder, util};
6
7#[derive(Debug, Default, Clone)]
8pub struct ROCm {
9	options: ExecutionProviderOptions
10}
11
12super::impl_ep!(arbitrary; ROCm);
13
14impl ROCm {
15	#[must_use]
16	pub fn with_device_id(mut self, device_id: i32) -> Self {
17		self.options.set("device_id", device_id.to_string());
18		self
19	}
20
21	#[must_use]
22	pub fn with_exhaustive_conv_search(mut self, enable: bool) -> Self {
23		self.options.set("miopen_conv_exhaustive_search", if enable { "1" } else { "0" });
24		self
25	}
26
27	#[must_use]
28	pub fn with_conv_use_max_workspace(mut self, enable: bool) -> Self {
29		self.options.set("miopen_conv_use_max_workspace", if enable { "1" } else { "0" });
30		self
31	}
32
33	#[must_use]
34	pub fn with_mem_limit(mut self, limit: usize) -> Self {
35		self.options.set("gpu_mem_limit", limit.to_string());
36		self
37	}
38
39	#[must_use]
40	pub fn with_arena_extend_strategy(mut self, strategy: ArenaExtendStrategy) -> Self {
41		self.options.set(
42			"arena_extend_strategy",
43			match strategy {
44				ArenaExtendStrategy::NextPowerOfTwo => "kNextPowerOfTwo",
45				ArenaExtendStrategy::SameAsRequested => "kSameAsRequested"
46			}
47		);
48		self
49	}
50
51	#[must_use]
52	pub fn with_copy_in_default_stream(mut self, enable: bool) -> Self {
53		self.options.set("do_copy_in_default_stream", if enable { "1" } else { "0" });
54		self
55	}
56
57	#[must_use]
58	pub fn with_compute_stream(mut self, ptr: *mut c_void) -> Self {
59		self.options.set("has_user_compute_stream", "1");
60		self.options.set("user_compute_stream", (ptr as usize).to_string());
61		self
62	}
63
64	#[must_use]
65	pub fn with_hip_graph(mut self, enable: bool) -> Self {
66		self.options.set("enable_hip_graph", if enable { "1" } else { "0" });
67		self
68	}
69
70	#[must_use]
71	pub fn with_tunable_op(mut self, enable: bool) -> Self {
72		self.options.set("tunable_op_enable", if enable { "1" } else { "0" });
73		self
74	}
75
76	#[must_use]
77	pub fn with_tuning(mut self, enable: bool) -> Self {
78		self.options.set("tunable_op_tuning_enable", if enable { "1" } else { "0" });
79		self
80	}
81
82	#[must_use]
83	pub fn with_max_tuning_duration(mut self, ms: i32) -> Self {
84		self.options.set("tunable_op_max_tuning_duration_ms", ms.to_string());
85		self
86	}
87}
88
89impl ExecutionProvider for ROCm {
90	fn name(&self) -> &'static str {
91		"ROCMExecutionProvider"
92	}
93
94	fn register(&self, session_builder: &mut SessionBuilder) -> Result<()> {
95		let mut rocm_options: *mut ort_sys::OrtROCMProviderOptions = ptr::null_mut();
96		ortsys![unsafe CreateROCMProviderOptions(&mut rocm_options)?];
97		let _guard = util::run_on_drop(|| {
98			ortsys![unsafe ReleaseROCMProviderOptions(rocm_options)];
99		});
100
101		let ffi_options = self.options.to_ffi();
102		ortsys![unsafe UpdateROCMProviderOptions(
103			rocm_options,
104			ffi_options.key_ptrs(),
105			ffi_options.value_ptrs(),
106			ffi_options.len()
107		)?];
108
109		ortsys![unsafe SessionOptionsAppendExecutionProvider_ROCM(session_builder.ptr_mut(), rocm_options)?];
110
111		Ok(())
112	}
113}
ort/ep/rocm.rs

ort/ep/
rocm.rs
