fix: manage RuntimeConfig, RuntimeCache with Arc/Rc (#102)

* fix: typo in argument of `create_execution_context_with_config`

* refactor!: put RuntimeConfig for ExecutionContext into a Rc

* refactor: put RuntimeCache into an Arc<Mutex>

Author: theHamsta

trtx/src/cuda_engine.rs

@@ -3,6 +3,7 @@
 //! [`CudaEngine`] wraps [`nvinfer1::ICudaEngine`] (C++ [`nvinfer1::ICudaEngine`](https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/_static/cpp-api/classnvinfer1_1_1_i_cuda_engine.html)).
 //! [`SerializationConfig`] wraps [`nvinfer1::ISerializationConfig`] (C++ [`nvinfer1::ISerializationConfig`](https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/_static/cpp-api/classnvinfer1_1_1_i_serialization_config.html)).
 
+use std::rc::Rc;
 use std::{ffi::CStr, marker::PhantomData};
 
 use crate::engine_inspector::EngineInspector;
@@ -349,16 +350,16 @@ impl<'engine> CudaEngine<'engine> {
                 .inner
                 .pin_mut()
                 .createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kSTATIC);
-            Ok(unsafe { ExecutionContext::from_ptr(context_ptr)? })
+            Ok(unsafe { ExecutionContext::from_ptr(context_ptr, None)? })
         }
         #[cfg(feature = "mock_runtime")]
-        Ok(unsafe { ExecutionContext::from_ptr(std::ptr::null_mut())? })
+        Ok(unsafe { ExecutionContext::from_ptr(std::ptr::null_mut(), None)? })
     }
 
     /// See [nvinfer1::ICudaEngine::createExecutionContext1]
     pub fn create_execution_context_with_config(
         &'_ mut self,
-        runtime_conifg: &'engine RuntimeConfig<'engine>,
+        runtime_config: Rc<RuntimeConfig<'engine>>,
     ) -> Result<ExecutionContext<'engine>> {
         #[cfg(not(feature = "mock_runtime"))]
         {
@@ -367,12 +368,12 @@ impl<'engine> CudaEngine<'engine> {
             let context_ptr = unsafe {
                 self.inner
                     .pin_mut()
-                    .createExecutionContext1(runtime_conifg.inner.as_mut_ptr())
+                    .createExecutionContext1(runtime_config.inner.as_mut_ptr())
             };
-            Ok(unsafe { ExecutionContext::from_ptr(context_ptr)? })
+            Ok(unsafe { ExecutionContext::from_ptr(context_ptr, Some(runtime_config))? })
         }
         #[cfg(feature = "mock_runtime")]
-        Ok(unsafe { ExecutionContext::from_ptr(std::ptr::null_mut())? })
+        Ok(unsafe { ExecutionContext::from_ptr(std::ptr::null_mut(), None)? })
     }
 
     /// See [nvinfer1::ICudaEngine::createSerializationConfig]

trtx/src/execution_context.rs

@@ -1,5 +1,6 @@
 use std::ffi::{CStr, CString};
 use std::pin::Pin;
+use std::rc::Rc;
 
 use cxx::UniquePtr;
 use trtx_sys::nvinfer1;
@@ -10,14 +11,16 @@ use crate::error::{Error, PropertySetAttempt, Result};
 use crate::interfaces::{
     DebugListener, ErrorRecorder, ProcessDebugTensor, Profiler, RecordError, ReportLayerTime,
 };
+use crate::RuntimeConfig;
 
 /// [`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] / [Profiler]
 /// pointers before their Rust wrappers run destructors.
-pub struct ExecutionContext<'a> {
-    _engine: std::marker::PhantomData<&'a CudaEngine<'a>>,
+pub struct ExecutionContext<'engine> {
+    _engine: std::marker::PhantomData<&'engine CudaEngine<'engine>>,
+    _config: Option<Rc<RuntimeConfig<'engine>>>,
     debug_listener: Option<Pin<Box<DebugListener>>>,
     profiler: Option<Pin<Box<Profiler>>>,
     error_recorder: Option<Pin<Box<ErrorRecorder>>>,
@@ -32,9 +35,10 @@ impl std::fmt::Debug for ExecutionContext<'_> {
     }
 }
 
-impl<'a> ExecutionContext<'a> {
+impl<'engine> ExecutionContext<'engine> {
     pub(crate) unsafe fn from_ptr(
         execution_context: *mut nvinfer1::IExecutionContext,
+        config: Option<Rc<RuntimeConfig<'engine>>>,
     ) -> Result<Self> {
         #[cfg(not(feature = "mock_runtime"))]
         if execution_context.is_null() {
@@ -44,6 +48,7 @@ impl<'a> ExecutionContext<'a> {
         }
         Ok(ExecutionContext {
             _engine: Default::default(),
+            _config: config,
             debug_listener: None,
             error_recorder: None,
             profiler: None,

trtx/src/runtime_config.rs

@@ -3,6 +3,8 @@
 //! [`RuntimeConfig`] wraps [`trtx_sys::nvinfer1::IRuntimeConfig`] (C++ [`nvinfer1::IRuntimeConfig`](https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/_static/cpp-api/classnvinfer1_1_1_i_runtime_config.html)).
 
 use std::marker::PhantomData;
+#[cfg(not(feature = "enterprise"))]
+use std::sync::{Arc, Mutex};
 
 #[cfg(not(feature = "enterprise"))]
 use crate::error::PropertySetAttempt;
@@ -21,6 +23,11 @@ use trtx_sys::{CudaGraphStrategy, DynamicShapesKernelSpecializationStrategy};
 pub struct RuntimeConfig<'engine> {
     pub(crate) inner: UniquePtr<IRuntimeConfig>,
     _engine: PhantomData<&'engine nvinfer1::ICudaEngine>,
+    // actually IRuntimeCache has its mutex, so we could omit this if we made mut methods of RuntimeCache (e.g. deserialize &self)
+    // this also makes it safe when we modify through our mutex, while cpp calls are made through
+    // IExecution calls
+    #[cfg(not(feature = "enterprise"))]
+    _cache: Option<Arc<Mutex<RuntimeCache<'engine>>>>,
 }
 
 impl std::fmt::Debug for RuntimeConfig<'_> {
@@ -40,6 +47,8 @@ impl<'engine> RuntimeConfig<'engine> {
         Ok(Self {
             inner: unsafe { UniquePtr::from_raw(runtime_config) },
             _engine: Default::default(),
+            #[cfg(not(feature = "enterprise"))]
+            _cache: None,
         })
     }
 
@@ -75,14 +84,17 @@ impl<'engine> RuntimeConfig<'engine> {
 
     #[cfg(not(feature = "enterprise"))]
     /// See [IRuntimeConfig::setRuntimeCache].
-    pub fn set_runtime_cache(&mut self, cache: &RuntimeCache<'engine>) -> Result<()> {
+    pub fn set_runtime_cache(&mut self, cache: Arc<Mutex<RuntimeCache<'engine>>>) -> Result<()> {
         if cfg!(not(feature = "mock")) {
             if self.inner.pin_mut().setRuntimeCache(
                 cache
+                    .lock()
+                    .unwrap()
                     .inner
                     .as_ref()
                     .expect("RuntimeCache inner must be non-null"),
             ) {
+                self._cache = Some(cache);
                 Ok(())
             } else {
                 Err(Error::FailedToSetProperty(