Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 6 additions & 10 deletions trtx/src/runtime_cache.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,49 +2,45 @@
//!
//! [`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 struct RuntimeCache {
pub(crate) inner: UniquePtr<IRuntimeCache>,
_engine: PhantomData<&'engine nvinfer1::ICudaEngine>,
}

/// # Safety
///
/// IRuntimeCache is internally protected by a shared mutex and
/// UniquePtr holds after initialization a valid IRuntimeCache (or nullptr in mock mode)
unsafe impl Send for RuntimeCache<'_> {}
unsafe impl Sync for RuntimeCache<'_> {}
unsafe impl Send for RuntimeCache {}
unsafe impl Sync for RuntimeCache {}

impl std::fmt::Debug for RuntimeCache<'_> {
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> {
impl RuntimeCache {
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>> {
pub fn serialize(&self) -> Result<HostMemory<'_>> {
#[cfg(not(feature = "mock"))]
{
let host_mem = unsafe { self.inner.serialize().as_mut() }
Expand Down
8 changes: 4 additions & 4 deletions trtx/src/runtime_config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,8 @@ pub struct RuntimeConfig<'engine> {
// 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>>>>, // Mutex, could now be removed with a
// breaking change to set_runtime_cache
_cache: Option<Arc<Mutex<RuntimeCache>>>, // Mutex, could now be removed with a
// breaking change to set_runtime_cache
}

impl std::fmt::Debug for RuntimeConfig<'_> {
Expand Down Expand Up @@ -83,7 +83,7 @@ impl<'engine> RuntimeConfig<'engine> {

#[cfg(not(feature = "enterprise"))]
/// See [IRuntimeConfig::createRuntimeCache].
pub fn create_runtime_cache(&self) -> Result<RuntimeCache<'engine>> {
pub fn create_runtime_cache(&self) -> Result<RuntimeCache> {
#[cfg(not(feature = "mock"))]
let cache_ptr = self.inner.createRuntimeCache();
#[cfg(feature = "mock")]
Expand All @@ -93,7 +93,7 @@ impl<'engine> RuntimeConfig<'engine> {

#[cfg(not(feature = "enterprise"))]
/// See [IRuntimeConfig::setRuntimeCache].
pub fn set_runtime_cache(&mut self, cache: Arc<Mutex<RuntimeCache<'engine>>>) -> Result<()> {
pub fn set_runtime_cache(&mut self, cache: Arc<Mutex<RuntimeCache>>) -> Result<()> {
if cfg!(not(feature = "mock")) {
if self.inner.pin_mut().setRuntimeCache(
cache
Expand Down
Loading