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}