atomr_accel_tensorrt/
sys.rs

1//! Hand-written FFI surface for libnvinfer (and libnvonnxparser when
2//! the `tensorrt-onnx` feature is on).
3//!
4//! TensorRT is a C++ API. We expose just the C-ABI shim functions we
5//! need from a thin C++ glue layer. The functions declared `extern "C"`
6//! here are intentionally empty when the `tensorrt-link` feature is
7//! off — the linker is never asked to resolve them, and the safe
8//! wrappers in `builder.rs`/`engine.rs`/`runtime.rs`/`onnx.rs` only
9//! call them through a `#[cfg(feature = "tensorrt-link")]` gate.
10//!
11//! The opaque pointer types (`IBuilder`, `IBuilderConfig`,
12//! `INetworkDefinition`, `ICudaEngine`, `IExecutionContext`,
13//! `IRuntime`, `IPluginCreator`) are zero-sized stand-ins for the
14//! corresponding TensorRT C++ classes. They are only ever held as
15//! `*mut` raw pointers; `Send`/`Sync` for the safe wrappers is granted
16//! via newtypes (see `engine.rs`).
17//!
18//! The C-ABI shim itself (a hand-written `nvinfer_shim.cpp`) is not
19//! shipped in this Phase 8 skeleton — it lives behind the
20//! `tensorrt-link` feature in a follow-up commit. Until then the FFI
21//! signatures here document the surface area and let downstream code
22//! type-check against a stable shape.
23
24#![allow(non_camel_case_types, dead_code, non_snake_case, unused_imports)]
25
26use std::os::raw::{c_char, c_int, c_void};
27
28// -------- Opaque object pointers --------
29
30#[repr(C)]
31pub struct IBuilder {
32    _private: [u8; 0],
33}
34
35#[repr(C)]
36pub struct IBuilderConfig {
37    _private: [u8; 0],
38}
39
40#[repr(C)]
41pub struct INetworkDefinition {
42    _private: [u8; 0],
43}
44
45#[repr(C)]
46pub struct ICudaEngine {
47    _private: [u8; 0],
48}
49
50#[repr(C)]
51pub struct IExecutionContext {
52    _private: [u8; 0],
53}
54
55#[repr(C)]
56pub struct IRuntime {
57    _private: [u8; 0],
58}
59
60#[repr(C)]
61pub struct IHostMemory {
62    _private: [u8; 0],
63}
64
65#[repr(C)]
66pub struct IRefitter {
67    _private: [u8; 0],
68}
69
70#[repr(C)]
71pub struct IInt8Calibrator {
72    _private: [u8; 0],
73}
74
75#[repr(C)]
76pub struct IPluginCreator {
77    _private: [u8; 0],
78}
79
80#[repr(C)]
81pub struct IPluginV3 {
82    _private: [u8; 0],
83}
84
85#[repr(C)]
86pub struct IOnnxParser {
87    _private: [u8; 0],
88}
89
90// -------- Enums (mirrored from NvInferRuntimeCommon.h / NvInferRuntime.h) --------
91
92#[repr(i32)]
93#[derive(Debug, Clone, Copy, PartialEq, Eq)]
94pub enum DataType {
95    kFLOAT = 0,
96    kHALF = 1,
97    kINT8 = 2,
98    kINT32 = 3,
99    kBOOL = 4,
100    kUINT8 = 5,
101    kFP8 = 6,
102    kBF16 = 7,
103    kINT64 = 8,
104    kINT4 = 9,
105}
106
107#[repr(u32)]
108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
109pub enum BuilderFlag {
110    kFP16 = 0,
111    kINT8 = 1,
112    kDEBUG = 2,
113    kGPU_FALLBACK = 3,
114    kREFIT = 4,
115    kDISABLE_TIMING_CACHE = 5,
116    kTF32 = 6,
117    kSPARSE_WEIGHTS = 7,
118    kSAFETY_SCOPE = 8,
119    kOBEY_PRECISION_CONSTRAINTS = 9,
120    kPREFER_PRECISION_CONSTRAINTS = 10,
121    kDIRECT_IO = 11,
122    kREJECT_EMPTY_ALGORITHMS = 12,
123    kBF16 = 13,
124    kFP8 = 14,
125    kSTRIP_PLAN = 15,
126    kVERSION_COMPATIBLE = 16,
127    kEXCLUDE_LEAN_RUNTIME = 17,
128}
129
130#[repr(i32)]
131#[derive(Debug, Clone, Copy, PartialEq, Eq)]
132pub enum DeviceType {
133    kGPU = 0,
134    kDLA = 1,
135}
136
137#[repr(u32)]
138#[derive(Debug, Clone, Copy, PartialEq, Eq)]
139pub enum TacticSource {
140    kCUBLAS = 0,
141    kCUBLAS_LT = 1,
142    kCUDNN = 2,
143    kEDGE_MASK_CONVOLUTIONS = 3,
144    kJIT_CONVOLUTIONS = 4,
145}
146
147#[repr(i32)]
148#[derive(Debug, Clone, Copy, PartialEq, Eq)]
149pub enum CalibrationAlgoType {
150    kLEGACY_CALIBRATION = 0,
151    kENTROPY_CALIBRATION = 1,
152    kENTROPY_CALIBRATION_2 = 2,
153    kMINMAX_CALIBRATION = 3,
154}
155
156#[repr(C)]
157#[derive(Debug, Clone, Copy)]
158pub struct Dims {
159    pub nb_dims: c_int,
160    pub d: [c_int; 8],
161}
162
163// -------- Function declarations (link probe is gated by `tensorrt-link`) --------
164//
165// The signatures below mirror the C-ABI shim that wraps the TensorRT
166// C++ classes. With the `tensorrt-link` feature off these are present
167// in the source as documentation; they're never referenced and so
168// never produce link errors.
169
170#[cfg(feature = "tensorrt-link")]
171extern "C" {
172    // Builder lifecycle
173    pub fn atomr_trt_builder_create(logger_severity: c_int) -> *mut IBuilder;
174    pub fn atomr_trt_builder_destroy(builder: *mut IBuilder);
175    pub fn atomr_trt_builder_create_network(
176        builder: *mut IBuilder,
177        flags: u32,
178    ) -> *mut INetworkDefinition;
179    pub fn atomr_trt_builder_create_config(builder: *mut IBuilder) -> *mut IBuilderConfig;
180    pub fn atomr_trt_builder_build_serialized(
181        builder: *mut IBuilder,
182        network: *mut INetworkDefinition,
183        config: *mut IBuilderConfig,
184    ) -> *mut IHostMemory;
185
186    // BuilderConfig knobs
187    pub fn atomr_trt_config_destroy(config: *mut IBuilderConfig);
188    pub fn atomr_trt_config_set_flag(config: *mut IBuilderConfig, flag: u32, on: c_int);
189    pub fn atomr_trt_config_set_memory_pool_limit(
190        config: *mut IBuilderConfig,
191        pool: c_int,
192        bytes: usize,
193    );
194    pub fn atomr_trt_config_set_default_device_type(config: *mut IBuilderConfig, dt: c_int);
195    pub fn atomr_trt_config_set_dla_core(config: *mut IBuilderConfig, core: c_int);
196    pub fn atomr_trt_config_set_tactic_sources(config: *mut IBuilderConfig, mask: u32);
197    pub fn atomr_trt_config_set_int8_calibrator(
198        config: *mut IBuilderConfig,
199        calibrator: *mut IInt8Calibrator,
200    );
201    pub fn atomr_trt_config_set_timing_cache(
202        config: *mut IBuilderConfig,
203        blob: *const u8,
204        len: usize,
205    );
206
207    // Engine
208    pub fn atomr_trt_engine_destroy(engine: *mut ICudaEngine);
209    pub fn atomr_trt_engine_create_execution_context(
210        engine: *mut ICudaEngine,
211    ) -> *mut IExecutionContext;
212    pub fn atomr_trt_engine_serialize(engine: *mut ICudaEngine) -> *mut IHostMemory;
213    pub fn atomr_trt_engine_num_io_tensors(engine: *mut ICudaEngine) -> c_int;
214    pub fn atomr_trt_engine_io_tensor_name(engine: *mut ICudaEngine, idx: c_int) -> *const c_char;
215    pub fn atomr_trt_engine_create_refitter(engine: *mut ICudaEngine) -> *mut IRefitter;
216
217    // Refitter
218    pub fn atomr_trt_refitter_destroy(refitter: *mut IRefitter);
219    pub fn atomr_trt_refitter_set_named_weights(
220        refitter: *mut IRefitter,
221        name: *const c_char,
222        weights: *const c_void,
223        bytes: usize,
224        dtype: c_int,
225    ) -> c_int;
226    pub fn atomr_trt_refitter_refit_engine(refitter: *mut IRefitter) -> c_int;
227
228    // ExecutionContext
229    pub fn atomr_trt_context_destroy(ctx: *mut IExecutionContext);
230    pub fn atomr_trt_context_set_input_shape(
231        ctx: *mut IExecutionContext,
232        name: *const c_char,
233        dims: *const Dims,
234    ) -> c_int;
235    pub fn atomr_trt_context_set_tensor_address(
236        ctx: *mut IExecutionContext,
237        name: *const c_char,
238        addr: *mut c_void,
239    ) -> c_int;
240    pub fn atomr_trt_context_enqueue_v3(
241        ctx: *mut IExecutionContext,
242        cuda_stream: *mut c_void,
243    ) -> c_int;
244
245    // Runtime + deserialise
246    pub fn atomr_trt_runtime_create(logger_severity: c_int) -> *mut IRuntime;
247    pub fn atomr_trt_runtime_destroy(runtime: *mut IRuntime);
248    pub fn atomr_trt_runtime_deserialize(
249        runtime: *mut IRuntime,
250        blob: *const u8,
251        len: usize,
252    ) -> *mut ICudaEngine;
253
254    // HostMemory
255    pub fn atomr_trt_host_memory_data(mem: *mut IHostMemory) -> *const u8;
256    pub fn atomr_trt_host_memory_size(mem: *mut IHostMemory) -> usize;
257    pub fn atomr_trt_host_memory_destroy(mem: *mut IHostMemory);
258
259    // Plugin registry (IPluginV3)
260    pub fn atomr_trt_register_plugin_creator(creator: *mut IPluginCreator) -> c_int;
261}
262
263#[cfg(all(feature = "tensorrt-link", feature = "tensorrt-onnx"))]
264extern "C" {
265    pub fn atomr_trt_onnx_parser_create(
266        network: *mut INetworkDefinition,
267        logger_severity: c_int,
268    ) -> *mut IOnnxParser;
269    pub fn atomr_trt_onnx_parser_destroy(parser: *mut IOnnxParser);
270    pub fn atomr_trt_onnx_parser_parse(
271        parser: *mut IOnnxParser,
272        data: *const u8,
273        len: usize,
274        path: *const c_char,
275    ) -> c_int;
276    pub fn atomr_trt_onnx_parser_num_errors(parser: *mut IOnnxParser) -> c_int;
277}
atomr_accel_tensorrt/sys.rs

atomr_accel_tensorrt/
sys.rs
