trtx_sys/
lib.rs

1//! Raw FFI bindings to NVIDIA TensorRT-RTX using autocxx
2//!
3//! ⚠️ **EXPERIMENTAL - NOT FOR PRODUCTION USE**
4//!
5//! This crate is in early experimental development. The API is unstable and will change.
6//! This is NOT production-ready software. Use at your own risk.
7//!
8//! This crate provides low-level, unsafe bindings to the TensorRT-RTX C++ library.
9//! For safe, ergonomic Rust API, use the `trtx` crate instead.
10//!
11//! # Architecture
12//!
13//! This crate uses a hybrid approach:
14//! - **autocxx** for direct C++ bindings to TensorRT classes
15//! - **Minimal C wrapper** for Logger callbacks (virtual methods)
16//!
17//! # Safety
18//!
19//! All functions in this crate are `unsafe` as they directly call into C++ code
20//! and perform no safety checks. Callers must ensure:
21//!
22//! - Pointers are valid and properly aligned
23//! - Lifetimes are managed correctly
24//! - Thread safety requirements are met
25//! - CUDA context is properly initialized
26
27#![allow(non_upper_case_globals)]
28#![allow(non_camel_case_types)]
29#![allow(non_snake_case)]
30#![allow(clippy::all)]
31
32#[allow(warnings)]
33mod enums {
34    include!(concat!(env!("OUT_DIR"), "/enums.rs"));
35}
36
37macro_rules! better_enum {
38    ($to:ident) => {
39        pub use crate::enums::$to;
40        impl Into<crate::nvinfer1::$to> for $to {
41            fn into(self) -> crate::nvinfer1::$to {
42                unsafe { transmute(self) }
43            }
44        }
45        impl From<crate::nvinfer1::$to> for $to {
46            fn from(value: crate::nvinfer1::$to) -> Self {
47                unsafe { transmute(value) }
48            }
49        }
50    };
51}
52
53use std::mem::transmute;
54use std::pin::Pin;
55better_enum!(LayerType);
56better_enum!(ActivationType);
57better_enum!(DataType);
58better_enum!(ProfilingVerbosity);
59better_enum!(MemoryPoolType);
60better_enum!(DeviceType);
61better_enum!(EngineCapability);
62better_enum!(BuilderFlag);
63better_enum!(PreviewFeature);
64better_enum!(HardwareCompatibilityLevel);
65better_enum!(RuntimePlatform);
66better_enum!(TilingOptimizationLevel);
67#[cfg(not(feature = "enterprise"))]
68better_enum!(ComputeCapability);
69better_enum!(CumulativeOperation);
70better_enum!(ElementWiseOperation);
71better_enum!(GatherMode);
72better_enum!(InterpolationMode);
73better_enum!(SampleMode);
74better_enum!(MatrixOperation);
75better_enum!(PoolingType);
76better_enum!(ReduceOperation);
77better_enum!(ResizeCoordinateTransformation);
78better_enum!(ResizeSelector);
79better_enum!(ResizeRoundMode);
80better_enum!(ScaleMode);
81better_enum!(ScatterMode);
82better_enum!(PaddingMode);
83better_enum!(UnaryOperation);
84better_enum!(TopKOperation);
85better_enum!(LayerInformationFormat);
86better_enum!(TensorLocation);
87better_enum!(TensorIOMode);
88better_enum!(TensorFormat);
89better_enum!(SerializationFlag);
90better_enum!(OptProfileSelector);
91better_enum!(AttentionNormalizationOp);
92better_enum!(SeekPosition);
93better_enum!(WeightsRole);
94better_enum!(TripLimit);
95better_enum!(LoopOutput);
96#[cfg(feature = "v_1_4")]
97better_enum!(MoEActType);
98#[cfg(feature = "v_1_4")]
99better_enum!(CollectiveOperation);
100
101pub use enums::ErrorCode;
102
103use autocxx::prelude::*;
104
105include_cpp! {
106    #include "NvInfer.h"
107    #include "NvInferRuntime.h"
108    #include "NvOnnxParser.h"
109
110    safety!(unsafe_ffi)
111
112    // Core TensorRT types
113    generate!("nvinfer1::IBuilder")
114    generate!("nvinfer1::IBuilderConfig")
115    generate!("nvinfer1::INetworkDefinition")
116    generate!("nvinfer1::ITensor")
117    generate!("nvinfer1::ILayer")
118    generate!("nvinfer1::IVersionedInterface")
119    generate!("nvinfer1::IProgressMonitor")
120    generate!("nvinfer1::IStreamWriter")
121    generate!("nvinfer1::IStreamReaderV2")
122    generate!("nvinfer1::IErrorRecorder")
123    generate!("nvinfer1::IProfiler")
124    generate!("nvinfer1::IGpuAllocator")
125    generate!("nvinfer1::IDebugListener")
126    generate!("nvinfer1::ISerializationConfig")
127    generate!("nvinfer1::IOptimizationProfile")
128    generate!("nvinfer1::IRefitter")
129
130    // Derived layer types - for inheritance support
131    generate!("nvinfer1::IActivationLayer")
132    generate!("nvinfer1::IConvolutionLayer")
133    generate!("nvinfer1::IPoolingLayer")
134    generate!("nvinfer1::IElementWiseLayer")
135    generate!("nvinfer1::IShuffleLayer")
136    generate!("nvinfer1::IConcatenationLayer")
137    generate!("nvinfer1::IMatrixMultiplyLayer")
138    generate!("nvinfer1::IConstantLayer")
139    generate!("nvinfer1::ISoftMaxLayer")
140    generate!("nvinfer1::IScaleLayer")
141    generate!("nvinfer1::IReduceLayer")
142    generate!("nvinfer1::ISliceLayer")
143    generate!("nvinfer1::IResizeLayer")
144    generate!("nvinfer1::ITopKLayer")
145    generate!("nvinfer1::IGatherLayer")
146    generate!("nvinfer1::IScatterLayer")
147    generate!("nvinfer1::ISelectLayer")
148    generate!("nvinfer1::IUnaryLayer")
149    generate!("nvinfer1::IIdentityLayer")
150    generate!("nvinfer1::IPaddingLayer")
151    generate!("nvinfer1::ICastLayer")
152    generate!("nvinfer1::IDeconvolutionLayer")
153    generate!("nvinfer1::IQuantizeLayer")
154    generate!("nvinfer1::IDequantizeLayer")
155    generate!("nvinfer1::IAssertionLayer")
156    generate!("nvinfer1::ICumulativeLayer")
157    generate!("nvinfer1::ILoop")
158    generate!("nvinfer1::IIfConditional")
159    generate!("nvinfer1::INormalizationLayer")
160    generate!("nvinfer1::ISqueezeLayer")
161    generate!("nvinfer1::IUnsqueezeLayer")
162    generate!("nvinfer1::ILRNLayer")
163    generate!("nvinfer1::IShapeLayer")
164    generate!("nvinfer1::IParametricReLULayer")
165    generate!("nvinfer1::IFillLayer")
166    generate!("nvinfer1::IEinsumLayer")
167    generate!("nvinfer1::IOneHotLayer")
168    generate!("nvinfer1::INonZeroLayer")
169    generate!("nvinfer1::IGridSampleLayer")
170    generate!("nvinfer1::INMSLayer")
171    generate!("nvinfer1::IReverseSequenceLayer")
172    generate!("nvinfer1::IDynamicQuantizeLayer")
173    generate!("nvinfer1::IRotaryEmbeddingLayer")
174    generate!("nvinfer1::IKVCacheUpdateLayer")
175    generate!("nvinfer1::IRaggedSoftMaxLayer")
176    generate!("nvinfer1::ILoopBoundaryLayer")
177    generate!("nvinfer1::IRecurrenceLayer")
178    generate!("nvinfer1::ILoopOutputLayer")
179    generate!("nvinfer1::ITripLimitLayer")
180    generate!("nvinfer1::IIteratorLayer")
181    generate!("nvinfer1::IConditionLayer")
182    generate!("nvinfer1::IIfConditionalOutputLayer")
183    generate!("nvinfer1::IIfConditionalInputLayer")
184    generate!("nvinfer1::IAttentionBoundaryLayer")
185    generate!("nvinfer1::IAttentionInputLayer")
186    generate!("nvinfer1::IAttentionOutputLayer")
187    generate!("nvinfer1::IAttention")
188    generate!("nvinfer1::IMoELayer")
189    generate!("nvinfer1::IDistCollectiveLayer")
190    // NOTE: IRNNv2Layer is deprecated (TRT_DEPRECATED) and autocxx cannot generate bindings for it
191    // RNN operations (lstm, lstmCell, gru, gruCell) remain deferred until we can work around this
192    // generate!("nvinfer1::IRNNv2Layer")
193
194    generate!("nvinfer1::IRuntime")
195    generate!("nvinfer1::ICudaEngine")
196    generate!("nvinfer1::IExecutionContext")
197    generate!("nvinfer1::IEngineInspector")
198    generate!("nvinfer1::IHostMemory")
199    generate!("nvinfer1::LayerInformationFormat")
200
201    // Try generating Dims64 directly (base class, not the typedef alias)
202    generate_pod!("nvinfer1::Dims64")
203
204    generate_pod!("nvinfer1::DataType")
205    generate_pod!("nvinfer1::TensorIOMode")
206    generate_pod!("nvinfer1::MemoryPoolType")
207    generate_pod!("nvinfer1::NetworkDefinitionCreationFlag")
208    generate_pod!("nvinfer1::ActivationType")
209    generate_pod!("nvinfer1::PoolingType")
210    generate_pod!("nvinfer1::PaddingMode")
211    generate_pod!("nvinfer1::ElementWiseOperation")
212    generate_pod!("nvinfer1::MatrixOperation")
213    generate_pod!("nvinfer1::UnaryOperation")
214    generate_pod!("nvinfer1::ReduceOperation")
215    generate_pod!("nvinfer1::CumulativeOperation")
216    generate_pod!("nvinfer1::GatherMode")
217    generate_pod!("nvinfer1::ScatterMode")
218    generate_pod!("nvinfer1::InterpolationMode")
219    generate_pod!("nvinfer1::ResizeCoordinateTransformation")
220    generate_pod!("nvinfer1::ResizeSelector")
221    generate_pod!("nvinfer1::ResizeRoundMode")
222    generate_pod!("nvinfer1::ProfilingVerbosity")
223    generate_pod!("nvinfer1::EngineCapability")
224    generate_pod!("nvinfer1::BuilderFlag")
225    generate_pod!("nvinfer1::BuilderFlags")
226    generate_pod!("nvinfer1::DeviceType")
227    generate_pod!("nvinfer1::TacticSource")
228    generate_pod!("nvinfer1::TacticSources")
229    generate_pod!("nvinfer1::PreviewFeature")
230    generate_pod!("nvinfer1::HardwareCompatibilityLevel")
231    generate_pod!("nvinfer1::RuntimePlatform")
232    generate_pod!("nvinfer1::TilingOptimizationLevel")
233    generate_pod!("nvinfer1::ComputeCapability")
234    generate_pod!("nvinfer1::APILanguage")
235    // NOTE: RNN enums commented out because IRNNv2Layer (deprecated) cannot be generated
236    // generate!("nvinfer1::RNNOperation")
237    // generate!("nvinfer1::RNNDirection")
238    // generate!("nvinfer1::RNNInputMode")
239    // generate!("nvinfer1::RNNGateType")
240    generate_pod!("nvinfer1::Weights")
241    generate_pod!("nvinfer1::Permutation")
242    generate_pod!("nvinfer1::TripLimit")
243    generate_pod!("nvinfer1::LoopOutput")
244    generate_pod!("nvinfer1::AttentionNormalizationOp")
245    generate_pod!("nvinfer1::WeightsRole")
246
247    generate!("nvinfer1::SeekPosition")
248    generate!("nvinfer1::ErrorCode")
249    generate!("nvinfer1::LayerType")
250    generate!("nvinfer1::SerializationFlags")
251    generate!("nvinfer1::SerializationFlag")
252    generate!("nvinfer1::OptProfileSelector")
253
254    // NOTE: createInferBuilder/Runtime moved to logger_bridge.cpp (autocxx struggles with these)
255
256    // ONNX Parser
257    generate!("nvonnxparser::IParser")
258    // NOTE: createParser also moved to logger_bridge.cpp
259
260}
261
262pub unsafe trait AsLayer {
263    fn as_layer(&self) -> &nvinfer1::ILayer {
264        // can't use safe `as_ref() -> &nvinfer1::ILayer` because only implemented for direct
265        // subclasses of ILayer
266        unsafe {
267            (self as *const Self as *const nvinfer1::ILayer)
268                .as_ref()
269                .unwrap()
270        }
271    }
272    fn as_layer_pin_mut(&mut self) -> Pin<&mut nvinfer1::ILayer> {
273        unsafe {
274            Pin::new_unchecked(
275                (self as *mut Self as *mut nvinfer1::ILayer)
276                    .as_mut()
277                    .unwrap(),
278            )
279        }
280    }
281}
282pub unsafe trait AsLayerTyped: AsLayer {
283    const TYPE: LayerType;
284}
285
286unsafe impl AsLayer for nvinfer1::IActivationLayer {}
287unsafe impl AsLayer for nvinfer1::IConvolutionLayer {}
288unsafe impl AsLayer for nvinfer1::ICastLayer {}
289unsafe impl AsLayer for nvinfer1::IPoolingLayer {}
290unsafe impl AsLayer for nvinfer1::ILRNLayer {}
291unsafe impl AsLayer for nvinfer1::IScaleLayer {}
292unsafe impl AsLayer for nvinfer1::ISoftMaxLayer {}
293unsafe impl AsLayer for nvinfer1::IDeconvolutionLayer {}
294unsafe impl AsLayer for nvinfer1::IConcatenationLayer {}
295unsafe impl AsLayer for nvinfer1::IElementWiseLayer {}
296unsafe impl AsLayer for nvinfer1::IUnaryLayer {}
297unsafe impl AsLayer for nvinfer1::IPaddingLayer {}
298unsafe impl AsLayer for nvinfer1::IShuffleLayer {}
299unsafe impl AsLayer for nvinfer1::IReduceLayer {}
300unsafe impl AsLayer for nvinfer1::ITopKLayer {}
301unsafe impl AsLayer for nvinfer1::IGatherLayer {}
302unsafe impl AsLayer for nvinfer1::IMatrixMultiplyLayer {}
303unsafe impl AsLayer for nvinfer1::IRaggedSoftMaxLayer {}
304unsafe impl AsLayer for nvinfer1::IConstantLayer {}
305unsafe impl AsLayer for nvinfer1::IIdentityLayer {}
306unsafe impl AsLayer for nvinfer1::ISliceLayer {}
307unsafe impl AsLayer for nvinfer1::IShapeLayer {}
308unsafe impl AsLayer for nvinfer1::IParametricReLULayer {}
309unsafe impl AsLayer for nvinfer1::IResizeLayer {}
310unsafe impl AsLayer for nvinfer1::ISelectLayer {}
311unsafe impl AsLayer for nvinfer1::IFillLayer {}
312unsafe impl AsLayer for nvinfer1::IQuantizeLayer {}
313unsafe impl AsLayer for nvinfer1::IDequantizeLayer {}
314unsafe impl AsLayer for nvinfer1::IScatterLayer {}
315unsafe impl AsLayer for nvinfer1::IEinsumLayer {}
316unsafe impl AsLayer for nvinfer1::IAssertionLayer {}
317unsafe impl AsLayer for nvinfer1::IOneHotLayer {}
318unsafe impl AsLayer for nvinfer1::INonZeroLayer {}
319unsafe impl AsLayer for nvinfer1::IGridSampleLayer {}
320unsafe impl AsLayer for nvinfer1::INMSLayer {}
321unsafe impl AsLayer for nvinfer1::IReverseSequenceLayer {}
322unsafe impl AsLayer for nvinfer1::INormalizationLayer {}
323unsafe impl AsLayer for nvinfer1::ISqueezeLayer {}
324unsafe impl AsLayer for nvinfer1::IUnsqueezeLayer {}
325unsafe impl AsLayer for nvinfer1::ICumulativeLayer {}
326unsafe impl AsLayer for nvinfer1::IDynamicQuantizeLayer {}
327unsafe impl AsLayer for nvinfer1::IRotaryEmbeddingLayer {}
328unsafe impl AsLayer for nvinfer1::IKVCacheUpdateLayer {}
329
330unsafe impl AsLayer for nvinfer1::IAttentionInputLayer {}
331unsafe impl AsLayer for nvinfer1::IAttentionOutputLayer {}
332unsafe impl AsLayer for nvinfer1::ILoopBoundaryLayer {}
333unsafe impl AsLayer for nvinfer1::ILoopOutputLayer {}
334unsafe impl AsLayer for nvinfer1::IRecurrenceLayer {}
335unsafe impl AsLayer for nvinfer1::ITripLimitLayer {}
336unsafe impl AsLayer for nvinfer1::IIteratorLayer {}
337unsafe impl AsLayer for nvinfer1::IConditionLayer {}
338unsafe impl AsLayer for nvinfer1::IIfConditionalOutputLayer {}
339unsafe impl AsLayer for nvinfer1::IIfConditionalInputLayer {}
340unsafe impl AsLayer for nvinfer1::IAttentionBoundaryLayer {}
341#[cfg(feature = "v_1_4")]
342unsafe impl AsLayer for nvinfer1::IMoELayer {}
343#[cfg(feature = "v_1_4")]
344unsafe impl AsLayer for nvinfer1::IDistCollectiveLayer {}
345
346// this one is not concrete
347unsafe impl AsLayer for nvinfer1::ILayer {}
348
349// indirect subclasses of ILayer e.g. via ILoopBoundaryLayer, IAttentionBoundaryLayer, IIfConditionalBoundaryLayer
350
351unsafe impl AsLayerTyped for nvinfer1::IActivationLayer {
352    const TYPE: LayerType = LayerType::kACTIVATION;
353}
354unsafe impl AsLayerTyped for nvinfer1::IConvolutionLayer {
355    const TYPE: LayerType = LayerType::kCONVOLUTION;
356}
357unsafe impl AsLayerTyped for nvinfer1::ICastLayer {
358    const TYPE: LayerType = LayerType::kCAST;
359}
360unsafe impl AsLayerTyped for nvinfer1::IPoolingLayer {
361    const TYPE: LayerType = LayerType::kPOOLING;
362}
363unsafe impl AsLayerTyped for nvinfer1::ILRNLayer {
364    const TYPE: LayerType = LayerType::kLRN;
365}
366unsafe impl AsLayerTyped for nvinfer1::IScaleLayer {
367    const TYPE: LayerType = LayerType::kSCALE;
368}
369unsafe impl AsLayerTyped for nvinfer1::ISoftMaxLayer {
370    const TYPE: LayerType = LayerType::kSOFTMAX;
371}
372unsafe impl AsLayerTyped for nvinfer1::IDeconvolutionLayer {
373    const TYPE: LayerType = LayerType::kDECONVOLUTION;
374}
375unsafe impl AsLayerTyped for nvinfer1::IConcatenationLayer {
376    const TYPE: LayerType = LayerType::kCONCATENATION;
377}
378unsafe impl AsLayerTyped for nvinfer1::IElementWiseLayer {
379    const TYPE: LayerType = LayerType::kELEMENTWISE;
380}
381unsafe impl AsLayerTyped for nvinfer1::IUnaryLayer {
382    const TYPE: LayerType = LayerType::kUNARY;
383}
384unsafe impl AsLayerTyped for nvinfer1::IPaddingLayer {
385    const TYPE: LayerType = LayerType::kPADDING;
386}
387unsafe impl AsLayerTyped for nvinfer1::IShuffleLayer {
388    const TYPE: LayerType = LayerType::kSHUFFLE;
389}
390unsafe impl AsLayerTyped for nvinfer1::IReduceLayer {
391    const TYPE: LayerType = LayerType::kREDUCE;
392}
393unsafe impl AsLayerTyped for nvinfer1::ITopKLayer {
394    const TYPE: LayerType = LayerType::kTOPK;
395}
396unsafe impl AsLayerTyped for nvinfer1::IGatherLayer {
397    const TYPE: LayerType = LayerType::kGATHER;
398}
399unsafe impl AsLayerTyped for nvinfer1::IMatrixMultiplyLayer {
400    const TYPE: LayerType = LayerType::kMATRIX_MULTIPLY;
401}
402unsafe impl AsLayerTyped for nvinfer1::IRaggedSoftMaxLayer {
403    const TYPE: LayerType = LayerType::kRAGGED_SOFTMAX;
404}
405unsafe impl AsLayerTyped for nvinfer1::IConstantLayer {
406    const TYPE: LayerType = LayerType::kCONSTANT;
407}
408unsafe impl AsLayerTyped for nvinfer1::IIdentityLayer {
409    const TYPE: LayerType = LayerType::kIDENTITY;
410}
411unsafe impl AsLayerTyped for nvinfer1::ISliceLayer {
412    const TYPE: LayerType = LayerType::kSLICE;
413}
414unsafe impl AsLayerTyped for nvinfer1::IShapeLayer {
415    const TYPE: LayerType = LayerType::kSHAPE;
416}
417unsafe impl AsLayerTyped for nvinfer1::IParametricReLULayer {
418    const TYPE: LayerType = LayerType::kPARAMETRIC_RELU;
419}
420unsafe impl AsLayerTyped for nvinfer1::IResizeLayer {
421    const TYPE: LayerType = LayerType::kRESIZE;
422}
423unsafe impl AsLayerTyped for nvinfer1::ISelectLayer {
424    const TYPE: LayerType = LayerType::kSELECT;
425}
426unsafe impl AsLayerTyped for nvinfer1::IFillLayer {
427    const TYPE: LayerType = LayerType::kFILL;
428}
429unsafe impl AsLayerTyped for nvinfer1::IQuantizeLayer {
430    const TYPE: LayerType = LayerType::kQUANTIZE;
431}
432unsafe impl AsLayerTyped for nvinfer1::IDequantizeLayer {
433    const TYPE: LayerType = LayerType::kDEQUANTIZE;
434}
435unsafe impl AsLayerTyped for nvinfer1::IScatterLayer {
436    const TYPE: LayerType = LayerType::kSCATTER;
437}
438unsafe impl AsLayerTyped for nvinfer1::IEinsumLayer {
439    const TYPE: LayerType = LayerType::kEINSUM;
440}
441unsafe impl AsLayerTyped for nvinfer1::IAssertionLayer {
442    const TYPE: LayerType = LayerType::kASSERTION;
443}
444unsafe impl AsLayerTyped for nvinfer1::IOneHotLayer {
445    const TYPE: LayerType = LayerType::kONE_HOT;
446}
447unsafe impl AsLayerTyped for nvinfer1::INonZeroLayer {
448    const TYPE: LayerType = LayerType::kNON_ZERO;
449}
450unsafe impl AsLayerTyped for nvinfer1::IGridSampleLayer {
451    const TYPE: LayerType = LayerType::kGRID_SAMPLE;
452}
453unsafe impl AsLayerTyped for nvinfer1::INMSLayer {
454    const TYPE: LayerType = LayerType::kNMS;
455}
456unsafe impl AsLayerTyped for nvinfer1::IReverseSequenceLayer {
457    const TYPE: LayerType = LayerType::kREVERSE_SEQUENCE;
458}
459unsafe impl AsLayerTyped for nvinfer1::INormalizationLayer {
460    const TYPE: LayerType = LayerType::kNORMALIZATION;
461}
462unsafe impl AsLayerTyped for nvinfer1::ISqueezeLayer {
463    const TYPE: LayerType = LayerType::kSQUEEZE;
464}
465unsafe impl AsLayerTyped for nvinfer1::IUnsqueezeLayer {
466    const TYPE: LayerType = LayerType::kUNSQUEEZE;
467}
468unsafe impl AsLayerTyped for nvinfer1::ICumulativeLayer {
469    const TYPE: LayerType = LayerType::kCUMULATIVE;
470}
471unsafe impl AsLayerTyped for nvinfer1::IDynamicQuantizeLayer {
472    const TYPE: LayerType = LayerType::kDYNAMIC_QUANTIZE;
473}
474unsafe impl AsLayerTyped for nvinfer1::IRotaryEmbeddingLayer {
475    const TYPE: LayerType = LayerType::kROTARY_EMBEDDING;
476}
477unsafe impl AsLayerTyped for nvinfer1::IKVCacheUpdateLayer {
478    const TYPE: LayerType = LayerType::kKVCACHE_UPDATE;
479}
480
481// indirect subclasses of ILayer e.g. via ILoopBoundaryLayer, IAttentionBoundaryLayer, IIfConditionalBoundaryLayer
482
483unsafe impl AsLayerTyped for nvinfer1::IAttentionInputLayer {
484    const TYPE: LayerType = LayerType::kATTENTION_INPUT;
485}
486unsafe impl AsLayerTyped for nvinfer1::IAttentionOutputLayer {
487    const TYPE: LayerType = LayerType::kATTENTION_OUTPUT;
488}
489unsafe impl AsLayerTyped for nvinfer1::ILoopBoundaryLayer {
490    const TYPE: LayerType = LayerType::kTRIP_LIMIT;
491}
492unsafe impl AsLayerTyped for nvinfer1::ILoopOutputLayer {
493    const TYPE: LayerType = LayerType::kLOOP_OUTPUT;
494}
495unsafe impl AsLayerTyped for nvinfer1::IRecurrenceLayer {
496    const TYPE: LayerType = LayerType::kRECURRENCE;
497}
498unsafe impl AsLayerTyped for nvinfer1::ITripLimitLayer {
499    const TYPE: LayerType = LayerType::kTRIP_LIMIT;
500}
501unsafe impl AsLayerTyped for nvinfer1::IIteratorLayer {
502    const TYPE: LayerType = LayerType::kITERATOR;
503}
504unsafe impl AsLayerTyped for nvinfer1::IConditionLayer {
505    const TYPE: LayerType = LayerType::kCONDITION;
506}
507unsafe impl AsLayerTyped for nvinfer1::IIfConditionalOutputLayer {
508    const TYPE: LayerType = LayerType::kCONDITIONAL_OUTPUT;
509}
510unsafe impl AsLayerTyped for nvinfer1::IIfConditionalInputLayer {
511    const TYPE: LayerType = LayerType::kCONDITIONAL_INPUT;
512}
513unsafe impl AsLayerTyped for nvinfer1::IAttentionBoundaryLayer {
514    const TYPE: LayerType = LayerType::kATTENTION_INPUT;
515}
516#[cfg(feature = "v_1_4")]
517unsafe impl AsLayerTyped for nvinfer1::IMoELayer {
518    const TYPE: LayerType = LayerType::kMOE;
519}
520#[cfg(feature = "v_1_4")]
521unsafe impl AsLayerTyped for nvinfer1::IDistCollectiveLayer {
522    const TYPE: LayerType = LayerType::kDIST_COLLECTIVE;
523}
524
525// Logger bridge C functions
526unsafe extern "C" {
527    pub unsafe fn get_tensorrt_version() -> u32;
528    pub unsafe fn create_rust_logger_bridge(
529        callback: RustLogCallback,
530        user_data: *mut std::ffi::c_void,
531    ) -> *mut RustLoggerBridge;
532
533    pub unsafe fn destroy_rust_logger_bridge(logger: *mut RustLoggerBridge);
534
535    pub unsafe fn get_logger_interface(logger: *mut RustLoggerBridge) -> *mut std::ffi::c_void; // Returns ILogger*
536                                                                                                //
537    pub unsafe fn trtx_create_progress_monitor(
538        user_data: *mut std::ffi::c_void,
539        phaseStart: unsafe extern "system" fn(
540            user_data: *mut std::ffi::c_void,
541            phaseName: *const ::std::os::raw::c_char,
542            parentPhase: *const ::std::os::raw::c_char,
543            nbSteps: i32,
544        ),
545        stepComplete: unsafe extern "system" fn(
546            user_data: *mut std::ffi::c_void,
547            phaseName: *const ::std::os::raw::c_char,
548            step: i32,
549        ) -> bool,
550        phaseFinish: unsafe extern "system" fn(
551            user_data: *mut std::ffi::c_void,
552            phaseName: *const ::std::os::raw::c_char,
553        ),
554    ) -> *mut nvinfer1::IProgressMonitor;
555    pub unsafe fn trtx_destroy_progress_monitor(cpp_obj: *mut nvinfer1::IProgressMonitor);
556    pub unsafe fn trtx_create_gpu_allocator(
557        rust_impl: *mut std::ffi::c_void,
558        allocateAsync: unsafe extern "system" fn(
559            this: *const std::ffi::c_void,
560            size: u64,
561            alignment: u64,
562            flags: u32,
563            cuda_stream: *mut std::ffi::c_void,
564        ) -> *mut std::ffi::c_void,
565        reallocate: unsafe extern "system" fn(
566            this: *const std::ffi::c_void,
567            memory: *mut std::ffi::c_void,
568            alignment: u64,
569            new_size: u64,
570        ) -> *mut std::ffi::c_void,
571        deallocateAsync: unsafe extern "system" fn(
572            this: *const std::ffi::c_void,
573            memory: *mut std::ffi::c_void,
574            cuda_stream: *mut std::ffi::c_void,
575        ) -> bool,
576    ) -> *mut nvinfer1::IGpuAllocator;
577    pub unsafe fn trtx_destroy_gpu_allocator(cpp_obj: *mut nvinfer1::IGpuAllocator);
578    pub unsafe fn trtx_create_error_recorder(
579        rust_impl: *mut std::ffi::c_void,
580        getNbErrors: *mut std::ffi::c_void,
581        getErrorCode: *mut std::ffi::c_void,
582        getErrorDesc: *mut std::ffi::c_void,
583        hasOverflowed: *mut std::ffi::c_void,
584        clear: *mut std::ffi::c_void,
585        reportError: *mut std::ffi::c_void,
586        incRefCount: *mut std::ffi::c_void,
587        decRefCount: *mut std::ffi::c_void,
588    ) -> *mut nvinfer1::IErrorRecorder;
589    pub unsafe fn trtx_destroy_error_recorder(cpp_obj: *mut nvinfer1::IErrorRecorder);
590
591    pub unsafe fn trtx_create_debug_listener(
592        rust_impl: *mut std::ffi::c_void,
593        processDebugTensor: unsafe extern "system" fn(
594            this: *const std::ffi::c_void,
595            addr: *const std::ffi::c_void,
596            location: nvinfer1::TensorLocation,
597            type_: nvinfer1::DataType,
598            shape: *const Dims64,
599            name: *const std::ffi::c_char,
600            stream: *mut std::ffi::c_void,
601        ) -> bool,
602    ) -> *mut nvinfer1::IDebugListener;
603
604    // TensorRT factory functions (wrapped as simple C functions)
605    #[cfg(feature = "link_tensorrt_rtx")]
606    pub unsafe fn create_infer_builder(logger: *mut std::ffi::c_void) -> *mut nvinfer1::IBuilder;
607
608    #[cfg(feature = "link_tensorrt_rtx")]
609    pub unsafe fn create_infer_runtime(logger: *mut std::ffi::c_void) -> *mut nvinfer1::IRuntime;
610
611    #[cfg(feature = "link_tensorrt_rtx")]
612    pub fn create_infer_refitter(
613        cuda_engine: *mut std::ffi::c_void,
614        logger: *mut std::ffi::c_void,
615    ) -> *mut nvinfer1::IRefitter; // Returns IRefitter*
616
617    pub unsafe fn trtx_refitter_get_missing(
618        refitter: *mut std::ffi::c_void,
619        size: i32,
620        layer_names: *mut *const std::os::raw::c_char,
621        roles: *mut i32,
622    ) -> i32;
623
624    pub unsafe fn trtx_refitter_get_all(
625        refitter: *mut std::ffi::c_void,
626        size: i32,
627        layer_names: *mut *const std::os::raw::c_char,
628        roles: *mut i32,
629    ) -> i32;
630
631    pub unsafe fn trtx_refitter_get_missing_weights(
632        refitter: *mut std::ffi::c_void,
633        size: i32,
634        weights_names: *mut *const std::os::raw::c_char,
635    ) -> i32;
636
637    pub unsafe fn trtx_refitter_get_all_weights(
638        refitter: *mut std::ffi::c_void,
639        size: i32,
640        weights_names: *mut *const std::os::raw::c_char,
641    ) -> i32;
642
643    // ONNX Parser factory function
644    #[cfg(feature = "link_tensorrt_onnxparser")]
645    pub unsafe fn create_onnx_parser(
646        network: *mut nvinfer1::INetworkDefinition,
647        logger: *mut std::ffi::c_void,
648    ) -> *mut nvonnxparser::IParser;
649
650    pub unsafe fn network_add_concatenation(
651        network: *mut std::ffi::c_void,
652        inputs: *mut *mut std::ffi::c_void,
653        nb_inputs: i32,
654    ) -> *mut std::ffi::c_void;
655
656    // Parser methods
657    pub unsafe fn parser_parse(
658        parser: *mut std::ffi::c_void,
659        data: *const std::ffi::c_void,
660        size: usize,
661    ) -> bool;
662    pub unsafe fn parser_get_nb_errors(parser: *mut std::ffi::c_void) -> i32;
663    pub unsafe fn parser_get_error(
664        parser: *mut std::ffi::c_void,
665        index: i32,
666    ) -> *mut std::ffi::c_void;
667    pub unsafe fn parser_error_desc(error: *mut std::ffi::c_void) -> *const std::os::raw::c_char;
668
669}
670
671// Opaque type for logger bridge
672#[repr(C)]
673pub struct RustLoggerBridge {
674    _unused: [u8; 0],
675}
676
677// Rust callback type for logger
678pub type RustLogCallback = unsafe extern "C" fn(
679    user_data: *mut std::ffi::c_void,
680    severity: i32,
681    msg: *const std::os::raw::c_char,
682);
683
684// Re-export TensorRT types from the private ffi module
685pub mod nvinfer1 {
686    pub use super::ffi::nvinfer1::*;
687}
688
689#[cfg(feature = "onnxparser")]
690pub mod nvonnxparser {
691    pub use super::ffi::nvonnxparser::*;
692}
693
694// Re-export Dims64 as Dims to match TensorRT's typedef
695pub use nvinfer1::Dims64;
696pub type Dims = Dims64;
697
698// Re-export InterpolationMode as ResizeMode to match TensorRT's typedef
699pub type ResizeMode = InterpolationMode;
700
701/// Helper methods for Dims construction (avoiding name collision with generated constructor)
702impl Dims64 {
703    /// Create a Dims from a slice of dimensions
704    pub fn from_slice(dims: &[i64]) -> Self {
705        let mut d = [0i64; 8];
706        let nb_dims = dims.len().min(8) as i32;
707        d[..nb_dims as usize].copy_from_slice(&dims[..nb_dims as usize]);
708        Self { nbDims: nb_dims, d }
709    }
710
711    /// Create a 2D Dims
712    pub fn new_2d(d0: i64, d1: i64) -> Self {
713        Self {
714            nbDims: 2,
715            d: [d0, d1, 0, 0, 0, 0, 0, 0],
716        }
717    }
718
719    /// Create a 3D Dims
720    pub fn new_3d(d0: i64, d1: i64, d2: i64) -> Self {
721        Self {
722            nbDims: 3,
723            d: [d0, d1, d2, 0, 0, 0, 0, 0],
724        }
725    }
726
727    /// Create a 4D Dims
728    pub fn new_4d(d0: i64, d1: i64, d2: i64, d3: i64) -> Self {
729        Self {
730            nbDims: 4,
731            d: [d0, d1, d2, d3, 0, 0, 0, 0],
732        }
733    }
734}
735
736// Re-export Weights
737pub use nvinfer1::Weights;
738
739/// Helper methods for Weights construction
740impl nvinfer1::Weights {
741    /// Create a Weights with FLOAT data type
742    pub fn new_float(values_ptr: *const std::ffi::c_void, count_val: i64) -> Self {
743        Self {
744            type_: nvinfer1::DataType::kFLOAT,
745            values: values_ptr,
746            count: count_val,
747        }
748    }
749
750    /// Create a Weights with specified data type
751    pub fn new_with_type(
752        data_type: nvinfer1::DataType,
753        values_ptr: *const std::ffi::c_void,
754        count_val: i64,
755    ) -> Self {
756        Self {
757            type_: data_type,
758            values: values_ptr,
759            count: count_val,
760        }
761    }
762}
763
764impl DataType {
765    pub const fn size_bits(self) -> usize {
766        match self {
767            DataType::kFLOAT => 32,
768            DataType::kHALF => 16,
769            DataType::kINT8 => 8,
770            DataType::kINT32 => 32,
771            DataType::kBOOL => 8,
772            DataType::kUINT8 => 8,
773            DataType::kFP8 => 8,
774            DataType::kBF16 => 16,
775            DataType::kINT64 => 64,
776            DataType::kINT4 => 4,
777            DataType::kFP4 => 4,
778            DataType::kE8M0 => 8,
779        }
780    }
781}
trtx_sys/lib.rs

trtx_sys/
lib.rs
