trtx 0.4.0

Safe Rust bindings to NVIDIA TensorRT-RTX (EXPERIMENTAL - NOT FOR PRODUCTION)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

//! Engine inspector (layer / engine introspection as text or JSON).
//!
//! [`EngineInspector`] wraps [`trtx_sys::nvinfer1::IEngineInspector`] (C++ [`nvinfer1::IEngineInspector`](https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/_static/cpp-api/classnvinfer1_1_1_i_engine_inspector.html)).

use crate::{Error, Result};
use std::{ffi::CStr, marker::PhantomData};

use autocxx::cxx::UniquePtr;
use trtx_sys::{nvinfer1, LayerInformationFormat};

/// [`trtx_sys::nvinfer1::IEngineInspector`] — C++ [`nvinfer1::IEngineInspector`](https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/_static/cpp-api/classnvinfer1_1_1_i_engine_inspector.html).
pub struct EngineInspector<'engine> {
    pub(crate) inner: UniquePtr<nvinfer1::IEngineInspector>,
    pub(crate) _engine: PhantomData<&'engine nvinfer1::ICudaEngine>,
}

impl EngineInspector<'_> {
    /// Returns layer information for the given layer index in the requested format.
    /// See [`trtx_sys::nvinfer1::IEngineInspector::getLayerInformation`].
    pub fn get_layer_information(
        &mut self,
        layer_index: i32,
        format: LayerInformationFormat,
    ) -> Result<String> {
        let ptr = self.inner.getLayerInformation(layer_index, format.into());
        Ok(if ptr.is_null() {
            return Err(Error::Runtime(
                "Could not get layer information".to_string(),
            ));
        } else {
            unsafe { CStr::from_ptr(ptr) }.to_str()?.to_string()
        })
    }

    /// Returns engine information in the requested format.
    /// See [`trtx_sys::nvinfer1::IEngineInspector::getEngineInformation`].
    pub fn get_engine_information(&self, format: LayerInformationFormat) -> Result<String> {
        let ptr = self.inner.getEngineInformation(format.into());
        Ok(if ptr.is_null() {
            return Err(Error::Runtime(
                "Could not get layer information".to_string(),
            ));
        } else {
            unsafe { CStr::from_ptr(ptr) }.to_str()?.to_string()
        })
    }
}