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

trtx_sys/
lib.rs
