async_tensorrt/
builder.rs

1use async_cuda::runtime::Future;
2
3use crate::ffi::builder_config::BuilderConfig;
4use crate::ffi::memory::HostBuffer;
5use crate::ffi::network::{NetworkDefinition, NetworkDefinitionCreationFlags};
6use crate::ffi::optimization_profile::OptimizationProfile;
7use crate::ffi::sync::builder::Builder as InnerBuilder;
8
9type Result<T> = std::result::Result<T, crate::error::Error>;
10
11/// Builds an engine from a network definition.
12///
13/// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder.html)
14pub struct Builder {
15    inner: InnerBuilder,
16}
17
18impl Builder {
19    /// Create a new [`Builder`].
20    ///
21    /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/namespacenvinfer1_1_1anonymous__namespace_02_nv_infer_8h_03.html)
22    pub async fn new() -> Result<Self> {
23        let inner = Future::new(InnerBuilder::new).await?;
24        Ok(Builder { inner })
25    }
26
27    /// Create a builder configuration object.
28    ///
29    /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder.html#a8fac4203e688430dff87483fc9db6bf2)
30    ///
31    /// # Return value
32    ///
33    /// A [`BuilderConfig`] that can later be passed to `build_serialized_network`.
34    #[inline(always)]
35    pub async fn config(&mut self) -> BuilderConfig {
36        Future::new(|| self.inner.config()).await
37    }
38
39    /// Create a new optimization profile.
40    ///
41    /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder.html#a68a8b59fbf86e42762b7087e6ffe6fb4)
42    #[inline(always)]
43    pub fn optimization_profile(&mut self) -> Result<OptimizationProfile> {
44        self.inner.optimization_profile()
45    }
46
47    /// Create a new optimization profile. This allocates an empty optimization profile, which
48    /// may or may not actually affect the building process later.
49    ///
50    /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder.html#a68a8b59fbf86e42762b7087e6ffe6fb4)
51    #[inline(always)]
52    pub fn add_default_optimization_profile(&mut self) -> Result<()> {
53        self.inner.add_default_optimization_profile()
54    }
55
56    /// Create a new optimization profile. This allocates an empty optimization profile, which
57    /// may or may not actually affect the building process later.
58    ///
59    /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder.html#a68a8b59fbf86e42762b7087e6ffe6fb4)
60    #[inline(always)]
61    pub fn with_default_optimization_profile(mut self) -> Result<Self> {
62        self.inner.add_default_optimization_profile()?;
63        Ok(self)
64    }
65
66    /// Create a network definition object.
67    ///
68    /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder.html#a853122d044b70383b2c9ebe7fdf11e07)
69    ///
70    /// # Arguments
71    ///
72    /// * `flags` - Flags for specifying network properties.
73    #[inline(always)]
74    pub fn network_definition(
75        &mut self,
76        flags: NetworkDefinitionCreationFlags,
77    ) -> NetworkDefinition {
78        self.inner.network_definition(flags)
79    }
80
81    /// Builds and serializes a network for the provided [`crate::ffi::network::NetworkDefinition`]
82    /// and [`BuilderConfig`].
83    ///
84    /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder.html#ab25ed4aec280df7d64e82930aa6b41c7)
85    ///
86    /// # Arguments
87    ///
88    /// * `network_definition` - Network definition.
89    /// * `config` - Builder configuration.
90    pub async fn build_serialized_network(
91        &mut self,
92        network_definition: &mut NetworkDefinition,
93        config: BuilderConfig,
94    ) -> Result<HostBuffer> {
95        Future::new(move || {
96            self.inner
97                .build_serialized_network(network_definition, config)
98        })
99        .await
100    }
101
102    /// Determine whether the platform has fast native INT8.
103    ///
104    /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder.html#ab09433c57e3ef02f7aad672ec4235ea4)
105    #[inline(always)]
106    pub fn platform_has_fast_int8(&self) -> bool {
107        self.inner.platform_has_fast_int8()
108    }
109
110    /// Determine whether the platform has fast native FP16.
111    ///
112    /// [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder.html#a6e42dd3ecb449ba54ffb823685a7ac47)
113    #[inline(always)]
114    pub fn platform_has_fast_fp16(&self) -> bool {
115        self.inner.platform_has_fast_fp16()
116    }
117}