diff --git a/trtx-sys/Cargo.toml b/trtx-sys/Cargo.toml index 97f6887..820a495 100644 --- a/trtx-sys/Cargo.toml +++ b/trtx-sys/Cargo.toml @@ -33,4 +33,5 @@ enterprise = [] [dependencies] autocxx = "0.30" +bitflags = "2.13.1" cxx = "1.0" diff --git a/trtx-sys/build.rs b/trtx-sys/build.rs index 5dcb25c..7b4ff3a 100644 --- a/trtx-sys/build.rs +++ b/trtx-sys/build.rs @@ -200,6 +200,8 @@ fn generate_enum_bindings(crate_root: &str, out_path: &Path, include_dir: &Path) ".*Role", ".*Limit", ".*AttentionNormalizationOp", + ".*EngineValidity", + ".*EngineInvalidityDiagnostics", ".*SeekPosition", ".*LoopOutput", ] { diff --git a/trtx-sys/src/lib.rs b/trtx-sys/src/lib.rs index 7940580..0f22584 100644 --- a/trtx-sys/src/lib.rs +++ b/trtx-sys/src/lib.rs @@ -52,6 +52,7 @@ macro_rules! better_enum { use std::mem::transmute; use std::pin::Pin; + better_enum!(LayerType); better_enum!(ActivationType); better_enum!(DataType); @@ -59,6 +60,39 @@ better_enum!(ProfilingVerbosity); better_enum!(MemoryPoolType); better_enum!(DeviceType); better_enum!(EngineCapability); +#[cfg(not(feature = "enterprise"))] +better_enum!(EngineValidity); +#[cfg(not(feature = "enterprise"))] +use bitflags::bitflags; +#[cfg(not(feature = "enterprise"))] +bitflags! { + /// Bitmask indicating the reason(s) why an engine is invalid. + /// + /// See https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/_static/c-api/namespacenvinfer1.html#a1ad7701ca8f1b97b3909323096d8d6f8 + #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] + pub struct EngineInvalidityDiagnostics: u64 { + /// TensorRT-RTX version mismatch to when engine was built. + const VERSION_MISMATCH = 1 << 0; + + /// Unsupported compute capability on current system. + const UNSUPPORTED_CC = 1 << 1; + + /// CUDA driver too old (driver downgrade compared to when engine was built). + const OLD_CUDA_DRIVER = 1 << 2; + + /// CUDA runtime too old (runtime downgrade compared to when engine was built). + const OLD_CUDA_RUNTIME = 1 << 3; + + /// Insufficient GPU memory to hold all engine weights. + const INSUFFICIENT_GPU_MEMORY = 1 << 4; + + /// Serialized engine does not conform to the expected format. + const MALFORMED_ENGINE = 1 << 5; + + /// Incorrect installation of the CUDA driver or runtime. + const CUDA_ERROR = 1 << 6; + } +} better_enum!(BuilderFlag); better_enum!(PreviewFeature); better_enum!(HardwareCompatibilityLevel); diff --git a/trtx/src/lib.rs b/trtx/src/lib.rs index 4d4dc56..7bb8c1a 100644 --- a/trtx/src/lib.rs +++ b/trtx/src/lib.rs @@ -322,6 +322,7 @@ pub use trtx_sys::{ #[cfg(not(feature = "enterprise"))] pub use trtx_sys::{ ComputeCapability, CudaGraphStrategy, DynamicShapesKernelSpecializationStrategy, + EngineInvalidityDiagnostics, EngineValidity, }; #[cfg(feature = "v_1_4")] diff --git a/trtx/src/runtime.rs b/trtx/src/runtime.rs index 29414f6..c6b6a4e 100644 --- a/trtx/src/runtime.rs +++ b/trtx/src/runtime.rs @@ -8,6 +8,8 @@ use std::marker::PhantomData; use cxx::UniquePtr; use log::trace; use trtx_sys::nvinfer1; +#[cfg(not(feature = "enterprise"))] +use trtx_sys::EngineInvalidityDiagnostics; pub use crate::cuda_engine::CudaEngine; pub use crate::engine_inspector::EngineInspector; @@ -137,6 +139,61 @@ impl<'runtime> Runtime<'runtime> { self.inner.getDeferredWeightsLoading() } } + + /// Returns the number of bytes required to inspect a serialized engine's header. + /// + /// See [`nvinfer1::IRuntime::getEngineHeaderSize`]. + #[cfg(not(feature = "enterprise"))] + pub fn engine_header_size(&self) -> usize { + if cfg!(feature = "mock_runtime") { + 0 + } else { + self.inner.getEngineHeaderSize() as usize + } + } + + /// Checks whether a serialized engine is likely to be valid on the current system. + /// + /// Returns the header-based validity classification and a bitmask of + /// [`crate::EngineInvalidityDiagnostics`] values. The diagnostics bitmask is zero for valid + /// and suboptimal engines. + /// + /// This only inspects the engine header and cannot detect corruption in the engine body. + /// + /// See [`nvinfer1::IRuntime::getEngineValidity`]. + #[cfg(not(feature = "enterprise"))] + pub fn engine_validity( + &self, + data: &[u8], + ) -> Result<(crate::EngineValidity, EngineInvalidityDiagnostics)> { + let header_size = self.engine_header_size(); + if data.len() < header_size { + return Err(Error::InvalidArgument(format!( + "engine buffer must contain at least {header_size} bytes, got {}", + data.len() + ))); + } + + if cfg!(feature = "mock_runtime") { + return Ok(( + crate::EngineValidity::kVALID, + EngineInvalidityDiagnostics::empty(), + )); + } + + let mut diagnostics = 0; + let validity = unsafe { + self.inner.getEngineValidity( + data.as_ptr() as *const autocxx::c_void, + data.len() as i64, + &mut diagnostics, + ) + }; + Ok(( + validity.into(), + EngineInvalidityDiagnostics::from_bits_retain(diagnostics), + )) + } //pub fn deserialize_cuda_engine_v2( //&'_ mut self, //stream_reader: &'runtime mut StreamReaderV2,