feat: add RuntimeConfig, RuntimeCache (#101)

* feat: add RuntimeConfig, RuntimeCache

* docs: fix doc comment issues

Author: theHamsta

trtexec-rs/src/cli.rs

@@ -53,7 +53,7 @@ pub struct Args {
     #[arg(short, long)]
     /// Capture TensorRT API usage to JSON for possible later replay
     ///
-    /// See https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/inference-library/capture-replay.html#capture-replay
+    /// See <https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/inference-library/capture-replay.html#capture-replay>
     pub api_capture: Option<PathBuf>,
 
     /// Report layer time

trtx-sys/build.rs

@@ -100,8 +100,11 @@ fn prepare_transformed_headers(header_dir: &Path, out_dir: &Path) -> PathBuf {
                 + r#"
 namespace nvinfer1 {
     class IMoELayer;
+    class IRuntimeCache;
     class IDistCollectiveLayer;
     enum class ComputeCapability : int32_t;
+    enum class CudaGraphStrategy : int32_t;
+    enum class DynamicShapesKernelSpecializationStrategy : int32_t;
 }"#;
 
             // trimming of indentation is necessary, so that rustdoc doesn't interpret

trtx-sys/src/lib.rs

@@ -65,6 +65,11 @@ better_enum!(HardwareCompatibilityLevel);
 better_enum!(RuntimePlatform);
 better_enum!(TilingOptimizationLevel);
 #[cfg(not(feature = "enterprise"))]
+better_enum!(CudaGraphStrategy);
+#[cfg(not(feature = "enterprise"))]
+better_enum!(DynamicShapesKernelSpecializationStrategy);
+better_enum!(ExecutionContextAllocationStrategy);
+#[cfg(not(feature = "enterprise"))]
 better_enum!(ComputeCapability);
 better_enum!(CumulativeOperation);
 better_enum!(ElementWiseOperation);
@@ -187,11 +192,10 @@ include_cpp! {
     generate!("nvinfer1::IAttention")
     generate!("nvinfer1::IMoELayer")
     generate!("nvinfer1::IDistCollectiveLayer")
-    // NOTE: IRNNv2Layer is deprecated (TRT_DEPRECATED) and autocxx cannot generate bindings for it
-    // RNN operations (lstm, lstmCell, gru, gruCell) remain deferred until we can work around this
-    // generate!("nvinfer1::IRNNv2Layer")
 
     generate!("nvinfer1::IRuntime")
+    generate!("nvinfer1::IRuntimeConfig")
+    generate!("nvinfer1::IRuntimeCache")
     generate!("nvinfer1::ICudaEngine")
     generate!("nvinfer1::IExecutionContext")
     generate!("nvinfer1::IEngineInspector")
@@ -232,11 +236,10 @@ include_cpp! {
     generate_pod!("nvinfer1::TilingOptimizationLevel")
     generate_pod!("nvinfer1::ComputeCapability")
     generate_pod!("nvinfer1::APILanguage")
-    // NOTE: RNN enums commented out because IRNNv2Layer (deprecated) cannot be generated
-    // generate!("nvinfer1::RNNOperation")
-    // generate!("nvinfer1::RNNDirection")
-    // generate!("nvinfer1::RNNInputMode")
-    // generate!("nvinfer1::RNNGateType")
+    generate_pod!("nvinfer1::CudaGraphStrategy")
+    generate_pod!("nvinfer1::DynamicShapesKernelSpecializationStrategy")
+    generate_pod!("nvinfer1::ExecutionContextAllocationStrategy")
+
     generate_pod!("nvinfer1::Weights")
     generate_pod!("nvinfer1::Permutation")
     generate_pod!("nvinfer1::TripLimit")

trtx/src/builder.rs

@@ -114,6 +114,12 @@ impl<'builder> Builder<'builder> {
         }
     }
 
+    /// See [IBuilder::createBuilderConfig]
+    pub fn create_builder_config(&'_ mut self) -> Result<BuilderConfig<'builder>> {
+        self.create_config()
+    }
+
+    /// See [IBuilder::createBuilderConfig]
     pub fn create_config(&'_ mut self) -> Result<BuilderConfig<'builder>> {
         #[cfg(not(feature = "mock"))]
         let config_ptr = self.inner.pin_mut().createBuilderConfig();

trtx/src/cuda_engine.rs

@@ -8,6 +8,7 @@ use std::{ffi::CStr, marker::PhantomData};
 use crate::engine_inspector::EngineInspector;
 use crate::error::PropertySetAttempt;
 use crate::host_memory::HostMemory;
+use crate::runtime_config::RuntimeConfig;
 use crate::{DataType, Error, ExecutionContext, Result};
 use autocxx::cxx::UniquePtr;
 use trtx_sys::{
@@ -34,7 +35,7 @@ impl SerializationConfig<'_> {
     pub fn flag(&self, flag: SerializationFlag) -> bool {
         self.inner.getFlag(flag.into())
     }
-    /// See [nvinfer1::SerializationConfig::getFlags]
+    /// See [nvinfer1::ISerializationConfig::getFlags]
     pub fn flags(&self) -> u32 {
         self.inner.getFlags()
     }
@@ -354,6 +355,26 @@ impl<'engine> CudaEngine<'engine> {
         Ok(unsafe { ExecutionContext::from_ptr(std::ptr::null_mut())? })
     }
 
+    /// See [nvinfer1::ICudaEngine::createExecutionContext1]
+    pub fn create_execution_context_with_config(
+        &'_ mut self,
+        runtime_conifg: &'engine RuntimeConfig<'engine>,
+    ) -> Result<ExecutionContext<'engine>> {
+        #[cfg(not(feature = "mock_runtime"))]
+        {
+            use crate::ExecutionContext;
+
+            let context_ptr = unsafe {
+                self.inner
+                    .pin_mut()
+                    .createExecutionContext1(runtime_conifg.inner.as_mut_ptr())
+            };
+            Ok(unsafe { ExecutionContext::from_ptr(context_ptr)? })
+        }
+        #[cfg(feature = "mock_runtime")]
+        Ok(unsafe { ExecutionContext::from_ptr(std::ptr::null_mut())? })
+    }
+
     /// See [nvinfer1::ICudaEngine::createSerializationConfig]
     pub fn create_serialization_config(&mut self) -> Result<SerializationConfig<'engine>> {
         let config = unsafe {
@@ -395,6 +416,15 @@ impl<'engine> CudaEngine<'engine> {
         Ok(unsafe { self.inner.isShapeInferenceIO(name_cstr.as_ptr()) })
     }
 
+    /// See [nvinfer1::ICudaEngine::createRuntimeConfig]
+    pub fn create_runtime_config(&'_ mut self) -> Result<RuntimeConfig<'engine>> {
+        #[cfg(not(feature = "mock"))]
+        let config_ptr = self.inner.pin_mut().createRuntimeConfig();
+        #[cfg(feature = "mock")]
+        let config_ptr = std::ptr::null_mut();
+        RuntimeConfig::new(config_ptr)
+    }
+
     /// Returns an iterator over all IO tensor names.
     pub fn io_tensor_names(&self) -> Result<CudaEngineIoTensorNamesIter<'_>> {
         Ok(CudaEngineIoTensorNamesIter {

trtx/src/error.rs

@@ -22,6 +22,9 @@ pub enum PropertySetAttempt {
     ExecutionContextTensorDebugState,
     ExecutionContextNcclCommunicator,
     ExecutionContextInputShape,
+    RuntimeConfigCudaGraphStrategy,
+    RuntimeConfigRuntimeCache,
+    RuntimeCacheDeserialize,
     DequantizeLayerBlockShape,
     QuantizeLayerBlockShape,
     AttentionLayerInput,
@@ -112,6 +115,12 @@ pub enum Error {
     #[error("Failed to create BuilderConfig")]
     BuilderConfigCreationFailed,
 
+    #[error("Failed to create RuntimeConfig")]
+    RuntimeConfigCreationFailed,
+
+    #[error("Failed to create RuntimeCache")]
+    RuntimeCacheCreationFailed,
+
     #[error("Failed to set property: {0:?}")]
     FailedToSetProperty(PropertySetAttempt),
 
@@ -132,6 +141,9 @@ pub enum Error {
 
     #[error("Failed to set IO tensor  address for tensor {tensor_name:?}")]
     FailedToSetTensorAddress { tensor_name: String },
+
+    #[error("Failed to reset to RuntimeCache")]
+    FailedToResetRuntimeCache,
 }
 
 impl<T> From<std::sync::PoisonError<T>> for Error {

trtx/src/execution_context.rs

@@ -14,7 +14,7 @@ use crate::interfaces::{
 /// [`trtx_sys::nvinfer1::IExecutionContext`] — C++ [`nvinfer1::IExecutionContext`](https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/_static/cpp-api/classnvinfer1_1_1_i_execution_context.html).
 ///
 /// `inner` is declared last so it is dropped first (see [`Drop`]): TensorRT must release
-/// [`DebugListener`](crate::interfaces::DebugListener) / [`Profiler`](crate::interfaces::Profiler)
+/// [DebugListener] / [Profiler]
 /// pointers before their Rust wrappers run destructors.
 pub struct ExecutionContext<'a> {
     _engine: std::marker::PhantomData<&'a CudaEngine<'a>>,

trtx/src/lib.rs

@@ -151,6 +151,9 @@ pub mod onnx_parser;
 pub mod optimization_profile;
 pub mod refitter;
 pub mod runtime;
+#[cfg(not(feature = "enterprise"))]
+pub mod runtime_cache;
+pub mod runtime_config;
 pub mod tensor;
 
 // Re-export commonly used types
@@ -169,7 +172,9 @@ pub use network::{ConvWeights, NetworkDefinition, OwnedConvWeights, OwnedWeights
 #[cfg(feature = "onnxparser")]
 pub use onnx_parser::OnnxParser;
 pub use refitter::Refitter;
-pub use runtime::{CudaEngine, EngineInspector, ExecutionContext, Runtime};
+#[cfg(not(feature = "enterprise"))]
+pub use runtime::RuntimeCache;
+pub use runtime::{CudaEngine, EngineInspector, ExecutionContext, Runtime, RuntimeConfig};
 
 #[cfg(feature = "dlopen_tensorrt_rtx")]
 #[cfg(not(any(feature = "link_tensorrt_rtx", feature = "mock")))]
@@ -303,12 +308,14 @@ pub fn dynamically_load_tensorrt_onnxparser(_filename: Option<impl AsFilename>)
 
 // Re-export TensorRT enums
 pub use trtx_sys::{
-    ActivationType, CumulativeOperation, DataType, ElementWiseOperation, GatherMode,
-    InterpolationMode, LayerInformationFormat, LayerType, MatrixOperation, PaddingMode,
-    PoolingType, ReduceOperation, ResizeCoordinateTransformation, ResizeMode, ResizeRoundMode,
-    ResizeSelector, SampleMode, ScaleMode, ScatterMode, TensorFormat, TensorIOMode, TopKOperation,
-    UnaryOperation,
+    ActivationType, CumulativeOperation, DataType, ElementWiseOperation,
+    ExecutionContextAllocationStrategy, GatherMode, InterpolationMode, LayerInformationFormat,
+    LayerType, MatrixOperation, PaddingMode, PoolingType, ReduceOperation,
+    ResizeCoordinateTransformation, ResizeMode, ResizeRoundMode, ResizeSelector, SampleMode,
+    ScaleMode, ScatterMode, TensorFormat, TensorIOMode, TopKOperation, UnaryOperation,
 };
 
 #[cfg(not(feature = "enterprise"))]
-pub use trtx_sys::ComputeCapability;
+pub use trtx_sys::{
+    ComputeCapability, CudaGraphStrategy, DynamicShapesKernelSpecializationStrategy,
+};

trtx/src/runtime.rs

@@ -14,6 +14,9 @@ pub use crate::engine_inspector::EngineInspector;
 use crate::error::{Error, Result};
 pub use crate::execution_context::ExecutionContext;
 use crate::logger::Logger;
+#[cfg(not(feature = "enterprise"))]
+pub use crate::runtime_cache::RuntimeCache;
+pub use crate::runtime_config::RuntimeConfig;
 
 /// [`trtx_sys::nvinfer1::IRuntime`] — C++ [`nvinfer1::IRuntime`](https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/_static/cpp-api/classnvinfer1_1_1_i_runtime.html).
 pub struct Runtime<'logger> {

trtx/src/runtime_cache.rs

@@ -0,0 +1,82 @@
+//! Runtime cache for TensorRT JIT compilation (serialize / deserialize).
+//!
+//! [`RuntimeCache`] wraps [`trtx_sys::nvinfer1::IRuntimeCache`] (C++ [`nvinfer1::IRuntimeCache`](https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/_static/cpp-api/classnvinfer1_1_1_i_runtime_cache.html).
+
+use std::marker::PhantomData;
+
+use crate::error::{PropertySetAttempt, Result};
+use crate::host_memory::HostMemory;
+use crate::Error;
+use cxx::UniquePtr;
+use trtx_sys::nvinfer1::{self, IRuntimeCache};
+
+/// [`trtx_sys::nvinfer1::IRuntimeCache`] — C++ [`nvinfer1::IRuntimeCache`](https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/_static/cpp-api/classnvinfer1_1_1_i_runtime_cache.html).
+pub struct RuntimeCache<'engine> {
+    pub(crate) inner: UniquePtr<IRuntimeCache>,
+    _engine: PhantomData<&'engine nvinfer1::ICudaEngine>,
+}
+
+impl std::fmt::Debug for RuntimeCache<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("RuntimeCache")
+            .field("inner", &format!("{:x}", self.inner.as_ptr() as usize))
+            .finish_non_exhaustive()
+    }
+}
+
+impl<'engine> RuntimeCache<'engine> {
+    pub(crate) fn new(cache: *mut nvinfer1::IRuntimeCache) -> Result<Self> {
+        #[cfg(not(feature = "mock"))]
+        if cache.is_null() {
+            return Err(Error::RuntimeCacheCreationFailed);
+        }
+        Ok(Self {
+            inner: unsafe { UniquePtr::from_raw(cache) },
+            _engine: Default::default(),
+        })
+    }
+
+    /// See [IRuntimeCache::serialize].
+    pub fn serialize(&self) -> Result<HostMemory<'engine>> {
+        #[cfg(not(feature = "mock"))]
+        {
+            let host_mem = unsafe { self.inner.serialize().as_mut() }
+                .ok_or_else(|| Error::Runtime("Failed to serialize IRuntimeCache".to_string()))?;
+            Ok(unsafe { HostMemory::from_raw(host_mem) })
+        }
+        #[cfg(feature = "mock")]
+        Ok(unsafe { HostMemory::from_raw(std::ptr::null_mut()) })
+    }
+
+    /// See [IRuntimeCache::deserialize].
+    pub fn deserialize(&mut self, blob: &[u8]) -> Result<()> {
+        if cfg!(not(feature = "mock")) {
+            if unsafe {
+                self.inner
+                    .pin_mut()
+                    .deserialize(blob.as_ptr() as *const autocxx::c_void, blob.len())
+            } {
+                Ok(())
+            } else {
+                Err(Error::FailedToSetProperty(
+                    PropertySetAttempt::RuntimeCacheDeserialize,
+                ))
+            }
+        } else {
+            Ok(())
+        }
+    }
+
+    /// See [IRuntimeCache::reset].
+    pub fn reset(&mut self) -> Result<()> {
+        if cfg!(not(feature = "mock")) {
+            if self.inner.pin_mut().reset() {
+                Ok(())
+            } else {
+                Err(Error::FailedToResetRuntimeCache)
+            }
+        } else {
+            Ok(())
+        }
+    }
+}