feat: implement EngineValidity, EngineInvalidityDiagnostics

trtx-sys/Cargo.toml

@@ -33,4 +33,5 @@ enterprise = []
 
 [dependencies]
 autocxx = "0.30"
+bitflags = "2.13.1"
 cxx = "1.0"

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",
     ] {

trtx-sys/src/lib.rs

@@ -52,13 +52,47 @@ macro_rules! better_enum {
 
 use std::mem::transmute;
 use std::pin::Pin;
+
 better_enum!(LayerType);
 better_enum!(ActivationType);
 better_enum!(DataType);
 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);

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")]

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,