TensorRT integration

NVIDIA TensorRT is an inference-optimised compiler/runtime for NVIDIA GPUs. TensorFlow integrates with it through tensorflow/compiler/tf2tensorrt/, which fuses compatible TF subgraphs into TRTEngineOps that wrap a TensorRT engine.

Purpose

• Take a SavedModel / frozen graph and identify subgraphs TensorRT can run.
• Compile each such subgraph into a TensorRT engine (ICudaEngine).
• Replace the original subgraph with a single TRTEngineOp that runs the engine inline at inference time.
• Support FP32, FP16, INT8 (with calibration), and the various TensorRT precision modes.

Directory layout

tensorflow/compiler/tf2tensorrt/
├── convert/
│   ├── convert_graph.{h,cc}        # Top-level: graph → TRT-converted graph
│   ├── convert_nodes.{h,cc}        # Per-op TF→TRT lowering
│   ├── trt_optimization_pass.{h,cc}  # Grappler pass entry
│   └── ...
├── kernels/
│   └── trt_engine_op.cc             # Op kernel that runs a TRT engine
├── plugin/                          # TRT plugin support for TF-only ops
├── ops/
├── utils/
├── segment/                          # Subgraph segmentation: which ops fuse together
└── ...

How it works

graph LR
    Graph[TF graph]
    Grappler[Grappler with TRTOptimizationPass]
    Segment[Segment graph]
    Convert[Per-segment TF -> TRT lowering]
    Engine[TensorRT engine]
    TRTOp[TRTEngineOp in TF graph]
    Runtime[TF runtime]
    GPU[NVIDIA GPU]

    Graph --> Grappler
    Grappler --> Segment
    Segment --> Convert
    Convert --> Engine
    Engine --> TRTOp
    Graph -- replaces fused subgraph --> TRTOp
    TRTOp --> Runtime
    Runtime --> GPU

1. Grappler pass. TRTOptimizationPass (tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc) plugs into Grappler. It runs after the standard optimizers.
2. Segmentation. The pass finds maximal contiguous subgraphs of TRT-supported ops (segment/).
3. Conversion. For each segment, Converter (convert_nodes.cc) translates each TF op into the corresponding TensorRT layer using the TensorRT C++ builder API.
4. Engine creation. TensorRT builds an optimised ICudaEngine (this can take seconds-to-minutes); the engine is serialised into an attribute of the new op.
5. TRTEngineOp. A single op replaces the segment. At runtime it deserialises the engine (or builds dynamic shapes JIT) and runs it on the GPU.

Modes

• Static mode — input shapes known at conversion time; engine is fully built at conversion.
• Dynamic mode — input shapes can vary; the op builds engines lazily per shape profile.
• INT8 — requires a calibration dataset; the converter inserts calibration ranges.

Python entrypoint

tf.experimental.tensorrt.Converter (tensorflow/python/compiler/tensorrt/) is the user-facing API:

from tensorflow.python.compiler.tensorrt import trt_convert as trt
converter = trt.TrtGraphConverterV2(input_saved_model_dir="…")
converter.convert()
converter.save("…_trt")

Build flags

TensorRT is optional. The build is gated by --config=tensorrt and TF_NEED_TENSORRT=1 at ./configure time. The TensorRT headers/libraries must be present; the build links against them dynamically.

Limitations

• Not every TF op has a TRT equivalent — unsupported ops fall back to the regular TF runtime, which means TRT engines are usually multiple disjoint segments.
• INT8 calibration requires representative input data.
• Engine build is non-deterministic; engines are typically cached (tensorflow/compiler/tf2tensorrt/utils/) keyed by input signatures.

Integration points

• Grappler — the TRTOptimizationPass runs as part of the optimizer pipeline.
• TF runtime — TRTEngineOp is a normal OpKernel that calls into TensorRT at compute time.
• Plugins — for TF-only ops that have no TRT layer, plugins (tensorflow/compiler/tf2tensorrt/plugin/) implement the layer in CUDA so TRT can call them.

Entry points for modification

• New op support — convert_nodes.cc is the one big switch on op type. Adding a new op means adding a converter function and registering it.
• New plugin — implement TensorRT's plugin interface; register in tensorflow/compiler/tf2tensorrt/plugin/.
• Segmentation policy — tensorflow/compiler/tf2tensorrt/segment/.

Related

• systems/grappler — the TRT pass plugs into Grappler.
• compilers/xla — XLA is an alternative for GPU acceleration; the two don't usually coexist on the same subgraph.