async_tensorrt/ffi/optimization_profile.rs
1use std::marker::PhantomData;
2
3use cpp::cpp;
4
5use crate::ffi::result;
6use crate::ffi::sync::builder::Builder;
7
8/// Defined in `NvInferRuntimeBase.h`
9const MAX_DIMS: usize = 8;
10
11type Result<T> = std::result::Result<T, crate::error::Error>;
12
13/// Synchronous implementation of [`crate::OptimizationProfile`].
14///
15/// Refer to [`crate::OptimizationProfile`] for documentation.
16pub struct OptimizationProfile<'builder>(*mut std::ffi::c_void, PhantomData<&'builder ()>);
17
18/// Implements [`Send`] for [`OptimizationProfile`].
19///
20/// # Safety
21///
22/// The TensorRT API is thread-safe with regards to all operations on [`OptimizationProfile`].
23unsafe impl<'builder> Send for OptimizationProfile<'builder> {}
24
25/// Implements [`Sync`] for [`OptimizationProfile`].
26///
27/// # Safety
28///
29/// The TensorRT API is thread-safe with regards to all operations on [`OptimizationProfile`].
30unsafe impl<'builder> Sync for OptimizationProfile<'builder> {}
31
32/// Optimization profile selector.
33///
34/// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/namespacenvinfer1.html#afd20e1d227abd394fdd3af0cb1525104)
35#[derive(Copy, Clone, Debug)]
36#[repr(i32)]
37enum OptimizationProfileSelector {
38 /// This is used to set or get the minimum permitted value for dynamic dimensions etc.
39 Min = 0,
40 /// This is used to set or get the value that is used in the optimization (kernel selection).
41 Opt = 1,
42 /// This is used to set or get the maximum permitted value for dynamic dimensions etc.
43 Max = 2,
44}
45
46impl<'builder> OptimizationProfile<'builder> {
47 /// Wrap internal pointer as [`OptimizationProfile`].
48 ///
49 /// # Arguments
50 ///
51 /// * `internal` - Pointer to wrap.
52 /// * `_builder` - Reference to builder to tie lifetime of optimization profile to.
53 ///
54 /// # Safety
55 ///
56 /// The pointer must point to a valid `IOptimizationProfile` object.
57 #[inline]
58 pub(crate) fn wrap(internal: *mut std::ffi::c_void, _builder: &'builder Builder) -> Self {
59 OptimizationProfile(internal, PhantomData)
60 }
61
62 /// Set the minimum values for an input shape tensor.
63 ///
64 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#ad89508bb5e59d46d106cb74d701934850
65 ///
66 /// # Arguments
67 ///
68 /// * `input_name` - Name of input tensor.
69 /// * `values` - Shape values.
70 ///
71 /// # Return value
72 ///
73 /// `false` if an inconsistency was detected.
74 #[inline]
75 pub fn set_min_shape_values(&mut self, input_name: &str, values: &[i32]) -> bool {
76 self.set_shape_values(input_name, OptimizationProfileSelector::Min as i32, values)
77 }
78
79 /// Set the optimium values for an input shape tensor.
80 ///
81 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#ad89508bb5e59d46d106cb74d701934850
82 ///
83 /// # Arguments
84 ///
85 /// * `input_name` - Name of input tensor.
86 /// * `values` - Shape values.
87 ///
88 /// # Return value
89 ///
90 /// `false` if an inconsistency was detected.
91 #[inline]
92 pub fn set_opt_shape_values(&mut self, input_name: &str, values: &[i32]) -> bool {
93 self.set_shape_values(input_name, OptimizationProfileSelector::Opt as i32, values)
94 }
95
96 /// Set the maximum values for an input shape tensor.
97 ///
98 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#ad89508bb5e59d46d106cb74d701934850
99 ///
100 /// # Arguments
101 ///
102 /// * `input_name` - Name of input tensor.
103 /// * `values` - Shape values.
104 ///
105 /// # Return value
106 ///
107 /// `false` if an inconsistency was detected.
108 #[inline]
109 pub fn set_max_shape_values(&mut self, input_name: &str, values: &[i32]) -> bool {
110 self.set_shape_values(input_name, OptimizationProfileSelector::Max as i32, values)
111 }
112
113 /// Set the minimum / optimum / maximum values for an input shape tensor.
114 ///
115 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#ad89508bb5e59d46d106cb74d701934850
116 ///
117 /// # Arguments
118 ///
119 /// * `input_name` - Name of input tensor.
120 /// * `select` - Optimization profile selector as integer.
121 /// * `values` - Shape values.
122 ///
123 /// # Return value
124 ///
125 /// `false` if an inconsistency was detected.
126 fn set_shape_values(&mut self, input_name: &str, select: i32, values: &[i32]) -> bool {
127 let internal = self.as_mut_ptr();
128 let input_name_cstr = std::ffi::CString::new(input_name).unwrap();
129 let input_name_ptr = input_name_cstr.as_ptr();
130 let nb_values = values.len() as i32;
131 let values_ptr = values.as_ptr();
132 let res = cpp!(unsafe [
133 internal as "void*",
134 input_name_ptr as "const char*",
135 select as "OptProfileSelector",
136 values_ptr as "const int32_t*",
137 nb_values as "int32_t"
138 ] -> bool as "bool" {
139 return ((IOptimizationProfile*) internal)->setShapeValues(input_name_ptr, select, values_ptr, nb_values);
140 });
141 res
142 }
143
144 /// Get the minimum values for an input shape tensor.
145 ///
146 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#a0654f6beafd1e4004950d5cd45ecab2b)
147 ///
148 /// # Arguments
149 ///
150 /// * `input_name` - Name of input tensor.
151 ///
152 /// # Return value
153 ///
154 /// Input shape if previously set.
155 pub fn get_min_shape_values(&self, input_name: &str) -> Result<Option<Vec<i32>>> {
156 self.get_shape_values(input_name, OptimizationProfileSelector::Min as i32)
157 }
158
159 /// Get the optimum values for an input shape tensor.
160 ///
161 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#a0654f6beafd1e4004950d5cd45ecab2b)
162 ///
163 /// # Arguments
164 ///
165 /// * `input_name` - Name of input tensor.
166 ///
167 /// # Return value
168 ///
169 /// Input shape if previously set.
170 pub fn get_opt_shape_values(&self, input_name: &str) -> Result<Option<Vec<i32>>> {
171 self.get_shape_values(input_name, OptimizationProfileSelector::Opt as i32)
172 }
173
174 /// Get the maximum values for an input shape tensor.
175 ///
176 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#a0654f6beafd1e4004950d5cd45ecab2b)
177 ///
178 /// # Arguments
179 ///
180 /// * `input_name` - Name of input tensor.
181 ///
182 /// # Return value
183 ///
184 /// Input shape if previously set.
185 pub fn get_max_shape_values(&self, input_name: &str) -> Result<Option<Vec<i32>>> {
186 self.get_shape_values(input_name, OptimizationProfileSelector::Max as i32)
187 }
188
189 /// Get the minimum / optimum / maximum values for an input shape tensor.
190 ///
191 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#a0654f6beafd1e4004950d5cd45ecab2b)
192 ///
193 /// # Arguments
194 ///
195 /// * `input_name` - Name of input tensor.
196 /// * `select` - Optimization profile selector as integer.
197 ///
198 /// # Return value
199 ///
200 /// Input shape if previously set.
201 fn get_shape_values(&self, input_name: &str, select: i32) -> Result<Option<Vec<i32>>> {
202 let internal = self.as_ptr();
203 let input_name_cstr = std::ffi::CString::new(input_name).unwrap();
204 let input_name_ptr = input_name_cstr.as_ptr();
205
206 let nb_shape_values = cpp!(unsafe [
207 internal as "void*",
208 input_name_ptr as "const char*"
209 ] -> i32 as "int32_t" {
210 return ((const IOptimizationProfile*) internal)->getNbShapeValues(input_name_ptr);
211 });
212
213 if nb_shape_values < 0 {
214 return Ok(None);
215 }
216 let nb_shape_values = nb_shape_values as usize;
217 let mut values = Vec::with_capacity(nb_shape_values);
218
219 let shape_values = cpp!(unsafe [
220 internal as "void*",
221 input_name_ptr as "const char*",
222 select as "OptProfileSelector"
223 ] -> *const i32 as "const int32_t*" {
224 return ((const IOptimizationProfile*) internal)->getShapeValues(input_name_ptr, select);
225 });
226 let shape_values = result!(shape_values, shape_values)?;
227 for i in 0..nb_shape_values {
228 let dim = unsafe { *shape_values.add(i) };
229 values.push(dim)
230 }
231 Ok(Some(values))
232 }
233
234 /// Set the minimum dimensions for a dynamic input tensor.
235 ///
236 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#ab723695382d6b03d4a0463b8cbe2b19f)
237 ///
238 /// # Arguments
239 ///
240 /// * `input_name` - Name of input tensor.
241 /// * `dims` - Dimensions.
242 ///
243 /// # Return value
244 ///
245 /// `false` if an inconsistency was detected.
246 pub fn set_min_dimensions(&mut self, input_name: &str, dims: &[i32]) -> bool {
247 self.set_dimensions(input_name, OptimizationProfileSelector::Min as i32, dims)
248 }
249
250 /// Set the optimum dimensions for a dynamic input tensor.
251 ///
252 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#ab723695382d6b03d4a0463b8cbe2b19f)
253 ///
254 /// # Arguments
255 ///
256 /// * `input_name` - Name of input tensor.
257 /// * `dims` - Dimensions.
258 ///
259 /// # Return value
260 ///
261 /// `false` if an inconsistency was detected.
262 pub fn set_opt_dimensions(&mut self, input_name: &str, dims: &[i32]) -> bool {
263 self.set_dimensions(input_name, OptimizationProfileSelector::Opt as i32, dims)
264 }
265
266 /// Set the maximum dimensions for a dynamic input tensor.
267 ///
268 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#ab723695382d6b03d4a0463b8cbe2b19f)
269 ///
270 /// # Arguments
271 ///
272 /// * `input_name` - Name of input tensor.
273 /// * `dims` - Dimensions.
274 ///
275 /// # Return value
276 ///
277 /// `false` if an inconsistency was detected.
278 pub fn set_max_dimensions(&mut self, input_name: &str, dims: &[i32]) -> bool {
279 self.set_dimensions(input_name, OptimizationProfileSelector::Max as i32, dims)
280 }
281
282 /// Set the minimum / optimum / maximum dimensions for a dynamic input tensor.
283 ///
284 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#ab723695382d6b03d4a0463b8cbe2b19f)
285 ///
286 /// # Arguments
287 ///
288 /// * `input_name` - Name of input tensor.
289 /// * `select` - Optimization profile selector as integer.
290 /// * `dims` - Dimensions.
291 ///
292 /// # Return value
293 ///
294 /// `false` if an inconsistency was detected.
295 fn set_dimensions(&mut self, input_name: &str, select: i32, dims: &[i32]) -> bool {
296 let internal = self.as_mut_ptr();
297 let input_name_cstr = std::ffi::CString::new(input_name).unwrap();
298 let input_name_ptr = input_name_cstr.as_ptr();
299 let nb_dims = dims.len() as i32;
300 let dims_ptr = dims.as_ptr();
301
302 let res = cpp!(unsafe [
303 internal as "void*",
304 input_name_ptr as "const char*",
305 select as "OptProfileSelector",
306 dims_ptr as "const int32_t*",
307 nb_dims as "int32_t"
308 ] -> bool as "bool" {
309 nvinfer1::Dims xdims;
310 xdims.nbDims = nb_dims;
311 for (int i = 0; i < xdims.nbDims; ++i) {
312 xdims.d[i] = dims_ptr[i];
313 }
314
315 return ((IOptimizationProfile*) internal)->setDimensions(input_name_ptr, select, xdims);
316 });
317 res
318 }
319
320 /// Get the minimum dimensions for a dynamic input tensor.
321 ///
322 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#a495725c79864f3e4059055307a8cc59d)
323 ///
324 /// # Arguments
325 ///
326 /// * `input_name` - Name of input tensor.
327 /// * `dims` - Dimensions.
328 ///
329 /// # Return value
330 ///
331 /// Dimensions if they have been previously set.
332 pub fn get_min_dimensions(&self, input_name: &str) -> Option<Vec<i32>> {
333 self.get_dimensions(input_name, OptimizationProfileSelector::Min as i32)
334 }
335
336 /// Get the optimum dimensions for a dynamic input tensor.
337 ///
338 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#a495725c79864f3e4059055307a8cc59d)
339 ///
340 /// # Arguments
341 ///
342 /// * `input_name` - Name of input tensor.
343 /// * `dims` - Dimensions.
344 ///
345 /// # Return value
346 ///
347 /// Dimensions if they have been previously set.
348 pub fn get_opt_dimensions(&self, input_name: &str) -> Option<Vec<i32>> {
349 self.get_dimensions(input_name, OptimizationProfileSelector::Opt as i32)
350 }
351
352 /// Get the maximum dimensions for a dynamic input tensor.
353 ///
354 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#a495725c79864f3e4059055307a8cc59d)
355 ///
356 /// # Arguments
357 ///
358 /// * `input_name` - Name of input tensor.
359 /// * `dims` - Dimensions.
360 ///
361 /// # Return value
362 ///
363 /// Dimensions if they have been previously set.
364 pub fn get_max_dimensions(&self, input_name: &str) -> Option<Vec<i32>> {
365 self.get_dimensions(input_name, OptimizationProfileSelector::Max as i32)
366 }
367
368 /// Get the minimum / optimum / maximum dimensions for a dynamic input tensor.
369 ///
370 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#a495725c79864f3e4059055307a8cc59d)
371 ///
372 /// # Arguments
373 ///
374 /// * `input_name` - Name of input tensor.
375 /// * `select` - Optimization profile selector as integer.
376 /// * `dims` - Dimensions.
377 ///
378 /// # Return value
379 ///
380 /// Dimensions if they have been previously set.
381 fn get_dimensions(&self, input_name: &str, select: i32) -> Option<Vec<i32>> {
382 let internal = self.as_ptr();
383 let input_name_cstr = std::ffi::CString::new(input_name).unwrap();
384 let input_name_ptr = input_name_cstr.as_ptr();
385 let mut dims = Vec::with_capacity(MAX_DIMS);
386 let dims_ptr = dims.as_mut_ptr();
387
388 let num_dimensions = cpp!(unsafe [
389 internal as "void*",
390 input_name_ptr as "const char*",
391 select as "OptProfileSelector",
392 dims_ptr as "int32_t*"
393 ] -> i32 as "int32_t" {
394 auto dims = ((const IOptimizationProfile*) internal)->getDimensions(input_name_ptr, select);
395 if (dims.nbDims > 0) {
396 for (int i = 0; i < dims.nbDims; ++i) {
397 dims_ptr[i] = dims.d[i];
398 }
399 }
400 return dims.nbDims;
401 });
402 if num_dimensions >= 0 {
403 // Safety: The vec has been initialized up until num_dimensions elements
404 unsafe {
405 dims.set_len(num_dimensions as usize);
406 }
407 Some(dims)
408 } else {
409 None
410 }
411 }
412
413 /// Set a target for extra GPU memory that may be used by this profile.
414 ///
415 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#abc9215e02ad6b5d911b35d45d59236e7)
416 ///
417 /// # Arguments
418 ///
419 /// * `target` - Additional memory that the builder should aim to maximally allocate for this profile, as a fraction of the memory it would use if the user did not impose any constraints on memory.
420 ///
421 /// # Return value
422 ///
423 /// `true` if the input is in the valid range (between 0 and 1 inclusive), else `false`.
424 pub fn set_extra_memory_target(&mut self, target: f32) -> bool {
425 let internal = self.as_ptr();
426 cpp!(unsafe [
427 internal as "const void*",
428 target as "float"
429 ] -> bool as "bool" {
430 return ((IOptimizationProfile*) internal)->setExtraMemoryTarget(target);
431 })
432 }
433
434 /// Get the extra memory target that has been defined for this profile.
435 ///
436 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#aa5339baa4f134993667bc2df94cb0c2e)
437 pub fn get_extra_memory_target(&self) -> f32 {
438 let internal = self.as_ptr();
439 cpp!(unsafe [
440 internal as "const void*"
441 ] -> f32 as "float" {
442 return ((const IOptimizationProfile*) internal)->getExtraMemoryTarget();
443 })
444 }
445
446 /// Check whether the optimization profile is valid.
447 ///
448 /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_optimization_profile.html#ae817a3cfb3f528a7b00173336521a187)
449 pub fn is_valid(&self) -> bool {
450 let internal = self.as_ptr();
451 cpp!(unsafe [
452 internal as "const void*"
453 ] -> bool as "bool" {
454 return ((const IOptimizationProfile*) internal)->isValid();
455 })
456 }
457
458 /// Get internal readonly pointer.
459 #[inline(always)]
460 pub fn as_ptr(&self) -> *const std::ffi::c_void {
461 let OptimizationProfile(internal, _) = *self;
462 internal
463 }
464
465 /// Get internal mutable pointer.
466 #[inline(always)]
467 pub fn as_mut_ptr(&mut self) -> *mut std::ffi::c_void {
468 let OptimizationProfile(internal, _) = *self;
469 internal
470 }
471}