rfc: support for torch-tensorrt

[hietalajulius]

RFC: Add Torch-TensorRT Support

Summary

This PR adds support for NVIDIA's Torch-TensorRT in tch-rs, enabling GPU inference optimization through TensorRT while maintaining the PyTorch ergonomics.

Motivation

Torch-TensorRT is NVIDIA's official PyTorch-TensorRT integration that can provide:

• Up to 5x faster inference compared to eager execution
• Automatic optimization of PyTorch models
• Seamless integration with existing PyTorch workflows
• Support for both dynamic and static shapes

Implementation Details

Build System Changes

Added support through:

1. New torch-tensorrt feature flag in Cargo.toml
2. Build script modifications to link TensorRT libraries
3. C++ preprocessor flag USE_TORCH_TENSORRT

Important Prerequisites & Caveats

1. Python Environment Requirement

   • Torch-TensorRT must be installed via pip in your Python environment:
     pip install torch-tensorrt
   • The build script will detect TensorRT through this installation

2. Rust Nightly Requirement

   • Requires Rust nightly toolchain due to the use of the -as-needed linker option

   • Add to your project:

     [toolchain]
     channel = "nightly"
     

   • Run with RUSTFLAGS="-Zunstable-options"

3. Runtime Environment

• LD_LIBRARY_PATH must include TensorRT library paths
• E.g.: export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/python3.12/site-packages/torch_tensorrt/lib:/path/to/python3.12/site-packages/tensorrt_libs

Usage

Enable in your project's Cargo.toml e.g.:

[features]
torch-tensorrt = ["tch/torch-tensorrt"]

Questions

• Overall is this within the scope of tch-rs?
• Only having support on nightly is not great, this can be avoided by using e.g. LD_PRELOAD (https://pytorch.org/TensorRT/user_guide/runtime.html) to overcome the -as-needed issue which is not amazing either
• Installation via python is convenient in dev, but an actual production env would probably want to install the libs without the python dependency
• Further tweaking of the linker flags probably needed to avoid setting LD_LIBRARY_PATH

[hietalajulius]

@LaurentMazare let me know if this is something that could be relevant for tch-rs, otherwise I could consider building a new crate that just depends on tch-rs 👍