1use alloc::string::ToString;
2
3use super::{ExecutionProvider, ExecutionProviderOptions, RegisterError};
4use crate::{error::Result, session::builder::SessionBuilder};
5
6#[derive(Debug, Default, Clone)]
7pub struct TensorRT {
8 options: ExecutionProviderOptions
9}
10
11super::impl_ep!(arbitrary; TensorRT);
12
13impl TensorRT {
14 #[must_use]
15 pub fn with_device_id(mut self, device_id: i32) -> Self {
16 self.options.set("device_id", device_id.to_string());
17 self
18 }
19
20 #[must_use]
23 pub unsafe fn with_compute_stream(mut self, stream: *mut ()) -> Self {
24 self.options.set("has_user_compute_stream", "1");
25 self.options.set("user_compute_stream", (stream as usize).to_string());
26 self
27 }
28
29 #[must_use]
30 pub fn with_max_workspace_size(mut self, max_size: usize) -> Self {
31 self.options.set("trt_max_workspace_size", max_size.to_string());
32 self
33 }
34
35 #[must_use]
36 pub fn with_min_subgraph_size(mut self, min_size: usize) -> Self {
37 self.options.set("trt_min_subgraph_size", min_size.to_string());
38 self
39 }
40
41 #[must_use]
42 pub fn with_max_partition_iterations(mut self, iterations: u32) -> Self {
43 self.options.set("trt_max_partition_iterations", iterations.to_string());
44 self
45 }
46
47 #[must_use]
48 pub fn with_fp16(mut self, enable: bool) -> Self {
49 self.options.set("trt_fp16_enable", if enable { "1" } else { "0" });
50 self
51 }
52
53 #[must_use]
54 pub fn with_int8(mut self, enable: bool) -> Self {
55 self.options.set("trt_int8_enable", if enable { "1" } else { "0" });
56 self
57 }
58
59 #[must_use]
60 pub fn with_dla(mut self, enable: bool) -> Self {
61 self.options.set("trt_dla_enable", if enable { "1" } else { "0" });
62 self
63 }
64
65 #[must_use]
66 pub fn with_dla_core(mut self, core: u32) -> Self {
67 self.options.set("trt_dla_core", core.to_string());
68 self
69 }
70
71 #[must_use]
72 pub fn with_int8_calibration_table_name(mut self, name: impl ToString) -> Self {
73 self.options.set("trt_int8_calibration_table_name", name.to_string());
74 self
75 }
76
77 #[must_use]
78 pub fn with_int8_use_native_calibration_table(mut self, enable: bool) -> Self {
79 self.options.set("trt_int8_use_native_calibration_table", if enable { "1" } else { "0" });
80 self
81 }
82
83 #[must_use]
84 pub fn with_engine_cache(mut self, enable: bool) -> Self {
85 self.options.set("trt_engine_cache_enable", if enable { "1" } else { "0" });
86 self
87 }
88
89 #[must_use]
90 pub fn with_engine_cache_path(mut self, path: impl ToString) -> Self {
91 self.options.set("trt_engine_cache_path", path.to_string());
92 self
93 }
94
95 #[must_use]
96 pub fn with_dump_subgraphs(mut self, enable: bool) -> Self {
97 self.options.set("trt_dump_subgraphs", if enable { "1" } else { "0" });
98 self
99 }
100
101 #[must_use]
102 pub fn with_engine_cache_prefix(mut self, prefix: impl ToString) -> Self {
103 self.options.set("trt_engine_cache_prefix", prefix.to_string());
104 self
105 }
106
107 #[must_use]
108 pub fn with_weight_stripped_engine(mut self, enable: bool) -> Self {
109 self.options.set("trt_weight_stripped_engine_enable", if enable { "1" } else { "0" });
110 self
111 }
112
113 #[must_use]
114 pub fn with_onnx_model_folder_path(mut self, path: impl ToString) -> Self {
115 self.options.set("trt_onnx_model_folder_path", path.to_string());
116 self
117 }
118
119 #[must_use]
120 pub fn with_engine_decryption(mut self, enable: bool) -> Self {
121 self.options.set("trt_engine_decryption_enable", if enable { "1" } else { "0" });
122 self
123 }
124
125 #[must_use]
126 pub fn with_engine_decryption_lib_path(mut self, lib_path: impl ToString) -> Self {
127 self.options.set("trt_engine_decryption_lib_path", lib_path.to_string());
128 self
129 }
130
131 #[must_use]
132 pub fn with_force_sequential_engine_build(mut self, enable: bool) -> Self {
133 self.options.set("trt_force_sequential_engine_build", if enable { "1" } else { "0" });
134 self
135 }
136
137 #[must_use]
138 pub fn with_context_memory_sharing(mut self, enable: bool) -> Self {
139 self.options.set("trt_context_memory_sharing_enable", if enable { "1" } else { "0" });
140 self
141 }
142
143 #[must_use]
144 pub fn with_layer_norm_fp32_fallback(mut self, enable: bool) -> Self {
145 self.options.set("trt_layer_norm_fp32_fallback", if enable { "1" } else { "0" });
146 self
147 }
148
149 #[must_use]
150 pub fn with_timing_cache(mut self, enable: bool) -> Self {
151 self.options.set("trt_timing_cache_enable", if enable { "1" } else { "0" });
152 self
153 }
154
155 #[must_use]
156 pub fn with_timing_cache_path(mut self, path: impl ToString) -> Self {
157 self.options.set("trt_timing_cache_path", path.to_string());
158 self
159 }
160
161 #[must_use]
162 pub fn with_force_timing_cache(mut self, enable: bool) -> Self {
163 self.options.set("trt_force_timing_cache", if enable { "1" } else { "0" });
164 self
165 }
166
167 #[must_use]
168 pub fn with_detailed_build_log(mut self, enable: bool) -> Self {
169 self.options.set("trt_detailed_build_log", if enable { "1" } else { "0" });
170 self
171 }
172
173 #[must_use]
174 pub fn with_build_heuristics(mut self, enable: bool) -> Self {
175 self.options.set("trt_build_heuristics_enable", if enable { "1" } else { "0" });
176 self
177 }
178
179 #[must_use]
180 pub fn with_sparsity(mut self, enable: bool) -> Self {
181 self.options.set("trt_sparsity_enable", if enable { "1" } else { "0" });
182 self
183 }
184
185 #[must_use]
186 pub fn with_builder_optimization_level(mut self, level: u8) -> Self {
187 self.options.set("trt_builder_optimization_level", level.to_string());
188 self
189 }
190
191 #[must_use]
192 pub fn with_auxiliary_streams(mut self, streams: i8) -> Self {
193 self.options.set("trt_auxiliary_streams", streams.to_string());
194 self
195 }
196
197 #[must_use]
198 pub fn with_tactic_sources(mut self, sources: impl ToString) -> Self {
199 self.options.set("trt_tactic_sources", sources.to_string());
200 self
201 }
202
203 #[must_use]
204 pub fn with_extra_plugin_lib_paths(mut self, paths: impl ToString) -> Self {
205 self.options.set("trt_extra_plugin_lib_paths", paths.to_string());
206 self
207 }
208
209 #[must_use]
210 pub fn with_profile_min_shapes(mut self, shapes: impl ToString) -> Self {
211 self.options.set("trt_profile_min_shapes", shapes.to_string());
212 self
213 }
214
215 #[must_use]
216 pub fn with_profile_max_shapes(mut self, shapes: impl ToString) -> Self {
217 self.options.set("trt_profile_max_shapes", shapes.to_string());
218 self
219 }
220
221 #[must_use]
222 pub fn with_profile_opt_shapes(mut self, shapes: impl ToString) -> Self {
223 self.options.set("trt_profile_opt_shapes", shapes.to_string());
224 self
225 }
226
227 #[must_use]
228 pub fn with_cuda_graph(mut self, enable: bool) -> Self {
229 self.options.set("trt_cuda_graph_enable", if enable { "1" } else { "0" });
230 self
231 }
232
233 #[must_use]
234 pub fn with_dump_ep_context_model(mut self, enable: bool) -> Self {
235 self.options.set("trt_dump_ep_context_model", if enable { "1" } else { "0" });
236 self
237 }
238
239 #[must_use]
240 pub fn with_ep_context_file_path(mut self, path: impl ToString) -> Self {
241 self.options.set("trt_ep_context_file_path", path.to_string());
242 self
243 }
244
245 #[must_use]
246 pub fn with_ep_context_embed_mode(mut self, mode: u8) -> Self {
247 self.options.set("trt_ep_context_embed_mode", mode.to_string());
248 self
249 }
250
251 #[must_use]
252 pub fn with_engine_hw_compatible(mut self, enable: bool) -> Self {
253 self.options.set("trt_engine_hw_compatible", if enable { "1" } else { "0" });
254 self
255 }
256}
257
258impl ExecutionProvider for TensorRT {
259 fn name(&self) -> &'static str {
260 "TensorrtExecutionProvider"
261 }
262
263 fn supported_by_platform(&self) -> bool {
264 cfg!(any(all(target_os = "linux", any(target_arch = "aarch64", target_arch = "x86_64")), all(target_os = "windows", target_arch = "x86_64")))
265 }
266
267 #[allow(unused, unreachable_code)]
268 fn register(&self, session_builder: &mut SessionBuilder) -> Result<(), RegisterError> {
269 #[cfg(any(feature = "load-dynamic", feature = "tensorrt"))]
270 {
271 use core::ptr;
272
273 use crate::{AsPointer, ortsys, util};
274
275 let mut trt_options: *mut ort_sys::OrtTensorRTProviderOptionsV2 = ptr::null_mut();
276 ortsys![unsafe CreateTensorRTProviderOptions(&mut trt_options)?];
277 let _guard = util::run_on_drop(|| {
278 ortsys![unsafe ReleaseTensorRTProviderOptions(trt_options)];
279 });
280
281 let ffi_options = self.options.to_ffi();
282 ortsys![unsafe UpdateTensorRTProviderOptions(
283 trt_options,
284 ffi_options.key_ptrs(),
285 ffi_options.value_ptrs(),
286 ffi_options.len()
287 )?];
288
289 ortsys![unsafe SessionOptionsAppendExecutionProvider_TensorRT_V2(session_builder.ptr_mut(), trt_options)?];
290 return Ok(());
291 }
292
293 Err(RegisterError::MissingFeature)
294 }
295}