async_tensorrt/ffi/
builder_config.rs

1use cpp::cpp;
2
3use crate::OptimizationProfile;
4
5type Result<T> = std::result::Result<T, crate::error::Error>;
6
7/// Holds properties for configuring a builder to produce an engine.
8///
9/// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder_config.html)
10pub struct BuilderConfig(*mut std::ffi::c_void);
11
12/// Implements [`Send`] for [`BuilderConfig`].
13///
14/// # Safety
15///
16/// The TensorRT API is thread-safe with regards to all operations on [`BuilderConfig`].
17unsafe impl Send for BuilderConfig {}
18
19/// Implements [`Sync`] for [`BuilderConfig`].
20///
21/// # Safety
22///
23/// The TensorRT API is thread-safe with regards to all operations on [`BuilderConfig`].
24unsafe impl Sync for BuilderConfig {}
25
26impl BuilderConfig {
27    /// Wrap internal pointer as [`BuilderConfig`].
28    ///
29    /// # Safety
30    ///
31    /// The pointer must point to a valid `IBuilderConfig` object.
32    pub(crate) fn wrap(internal: *mut std::ffi::c_void) -> Self {
33        Self(internal)
34    }
35
36    /// Set the maximum workspace size.
37    ///
38    /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder_config.html#a8209999988ab480c60c8a905dfd2654d)
39    ///
40    /// # Arguments
41    ///
42    /// * `size` - The maximum GPU temporary memory which the engine can use at execution time in
43    ///   bytes.
44    pub fn with_max_workspace_size(mut self, size: usize) -> Self {
45        let internal = self.as_mut_ptr();
46        cpp!(unsafe [
47            internal as "void*",
48            size as "std::size_t"
49        ] {
50            ((IBuilderConfig*) internal)->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, size);
51        });
52        self
53    }
54
55    /// Set the `kSTRICT_TYPES` flag.
56    ///
57    /// [TensorRT documentation for `setFlag`](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder_config.html#ac9821504ae7a11769e48b0e62761837e)
58    /// [TensorRT documentation for `kSTRICT_TYPES`](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/namespacenvinfer1.html#abdc74c40fe7a0c3d05d2caeccfbc29c1ad3ff8ff39475957d8676c2cda337add7)
59    pub fn with_strict_types(mut self) -> Self {
60        let internal = self.as_mut_ptr();
61        cpp!(unsafe [
62            internal as "void*"
63        ] {
64            #if NV_TENSORRT_MAJOR >= 10
65            ((IBuilderConfig*) internal)->setFlag(BuilderFlag::kPREFER_PRECISION_CONSTRAINTS);
66            ((IBuilderConfig*) internal)->setFlag(BuilderFlag::kDIRECT_IO);
67            ((IBuilderConfig*) internal)->setFlag(BuilderFlag::kREJECT_EMPTY_ALGORITHMS);
68            #else
69            ((IBuilderConfig*) internal)->setFlag(BuilderFlag::kSTRICT_TYPES);
70            #endif
71        });
72        self
73    }
74
75    /// Set the `kFP16` flag.
76    ///
77    /// [TensorRT documentation for `setFlag`](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder_config.html#ac9821504ae7a11769e48b0e62761837e)
78    /// [TensorRT documentation for `kFP16`](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/namespacenvinfer1.html#abdc74c40fe7a0c3d05d2caeccfbc29c1a56e4ef5e47a48568bd24c4e0aaabcead)
79    pub fn with_fp16(mut self) -> Self {
80        let internal = self.as_mut_ptr();
81        cpp!(unsafe [
82            internal as "void*"
83        ] {
84            ((IBuilderConfig*) internal)->setFlag(BuilderFlag::kFP16);
85        });
86        self
87    }
88
89    /// Add an optimization profile.
90    ///
91    /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder_config.html#ab97fa40c85fa8afab65fc2659e38da82)
92    pub fn with_optimization_profile(
93        mut self,
94        optimization_profile: OptimizationProfile,
95    ) -> Result<Self> {
96        self.add_optimization_profile(optimization_profile)?;
97        Ok(self)
98    }
99
100    /// Add an optimization profile.
101    ///
102    /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder_config.html#ab97fa40c85fa8afab65fc2659e38da82)
103    pub fn add_optimization_profile(
104        &mut self,
105        optimization_profile: OptimizationProfile,
106    ) -> Result<()> {
107        let internal = self.as_mut_ptr();
108        let optimization_profile = optimization_profile.as_ptr();
109        let index = cpp!(unsafe [
110            internal as "void*",
111            optimization_profile as "const IOptimizationProfile*"
112        ] -> i32 as "std::int32_t" {
113           return ((IBuilderConfig*) internal)->addOptimizationProfile(optimization_profile);
114        });
115        if index >= 0 {
116            Ok(())
117        } else {
118            Err(crate::error::last_error())
119        }
120    }
121
122    /// Get internal readonly pointer.
123    #[inline(always)]
124    pub fn as_ptr(&self) -> *const std::ffi::c_void {
125        let BuilderConfig(internal) = *self;
126        internal
127    }
128
129    /// Get internal mutable pointer.
130    #[inline(always)]
131    pub fn as_mut_ptr(&mut self) -> *mut std::ffi::c_void {
132        let BuilderConfig(internal) = *self;
133        internal
134    }
135}
136
137impl Drop for BuilderConfig {
138    fn drop(&mut self) {
139        let internal = self.as_mut_ptr();
140        cpp!(unsafe [
141            internal as "void*"
142        ] {
143            destroy((IBuilderConfig*) internal);
144        });
145    }
146}
async_tensorrt/ffi/builder_config.rs

async_tensorrt/ffi/
builder_config.rs
