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}