Add in-tree isaac_deploy_trt and optimize_graph TensorRT rewrite.

.github/workflows/ci.yml

@@ -42,15 +42,15 @@ jobs:
 
       # Visualization modules ship with leapp; gated deps install on 3.11+.
       - name: Install package and test dependencies
-        run: python -m pip install -e ".[dev,test,warp-cu12]" cmake ninja
+        run: python -m pip install -e ".[dev,test,warp-cu12,graph-surgery]" cmake ninja
 
       - name: Build Warp runtime
         run: leapp-build-warp-runtime
 
       - name: Run pytest
         env:
           CMAKE_GENERATOR: Ninja
-        run: pytest tests/ packages/leapp-visualization/tests/ -v
+        run: pytest tests/ packages/leapp-visualization/tests/ packages/isaac_deploy_trt/tests/ -v
 
   docs:
     runs-on: ubuntu-latest

AGENTS.md

@@ -21,8 +21,9 @@ Python 3.10+ and PyTorch 2.6.0+. Visualization extras install only on
 | Warp tracing (`TracedWpArray`, APIC, `warp_op`) | `warp_support.md` |
 | Public annotate/lifecycle API or docs examples | `docs/source/api/index.rst` plus the matching guide under `docs/source/` |
 | Graph PNG layout | `packages/leapp-visualization/` |
+| TensorRT ONNX optimize / catalog / pipeline merge | `packages/isaac_deploy_trt/` |
 | Export backends | `leapp/backends/` |
-| YAML / `InferenceManager` | `leapp/leapp.py` (`compile_graph`), `leapp/inference_manager.py`, `docs/source/generated_configs.rst` |
+| YAML / `InferenceManager` | `leapp/leapp.py` (`compile_graph`, `optimize_graph`), `leapp/inference_manager.py`, `docs/source/generated_configs.rst` |
 
 Do not mix the skill's current `(context_obj, output_port)` identity with
 legacy `leapp_tag` helpers. The skill states which checkout you have.
@@ -31,7 +32,7 @@ legacy `leapp_tag` helpers. The skill states which checkout you have.
 
 ```
 leapp/                      # installable package
-  leapp.py                  # start / stop / compile_graph / annotate facade
+  leapp.py                  # start / stop / compile_graph / optimize_graph / annotate facade
   export_manager.py         # annotate.* implementations
   inference_manager.py      # runtime loader for exported YAML
   backends/                 # jit, onnx-dynamo, onnx-torchscript, pt2, none
@@ -41,6 +42,7 @@ leapp/                      # installable package
     datatypes/              # TracedData, TracedTensor, TracedNpArray, TracedWpArray
   utils/                    # logging, TensorSemantics, enums, GraphConfigs
 packages/leapp-visualization/   # static PNG layout (ships with leapp)
+packages/isaac_deploy_trt/      # TensorRT ONNX optimize, catalogs, merge-pipeline
 docs/source/                # Sphinx (NVIDIA theme)
 examples/                   # runnable samples; CI runs tests/test_examples/
 tests/
@@ -54,12 +56,21 @@ allowed on `leapp.annotate` are listed in `AnnotateAPI._ALLOWED_APIS` in
 `leapp/leapp.py`. Do not add a consumer API without updating both, the
 docs, and tests.
 
+Do not call `compile_graph()` before `stop()`. Call
+`leapp.optimize_graph(trt_compatible=True)` after `compile_graph()` when
+exported ONNX should be rewritten for TensorRT (`isaac_deploy_trt` /
+`graph-surgery` extra). Use `leapp.start()`, `leapp.stop()`, and
+`leapp.compile_graph()` for graph lifecycle. Use `annotate` only for
+annotation APIs such as `method()`, `input_tensors()`, and
+`output_tensors()`.
+
 ## Commands
 
 ```bash
 python -m pip install -e ".[dev,test,warp-cu12]"  # CUDA 12
 python -m pip install -e ".[dev,test,warp-cu13]"  # CUDA 13
-pytest tests/ packages/leapp-visualization/tests/ -v
+python -m pip install -e ".[dev,graph-surgery]"   # TensorRT ONNX optimize extras
+pytest tests/ packages/leapp-visualization/tests/ packages/isaac_deploy_trt/tests/ -v
 ```
 
 Narrower (use these while iterating):
@@ -71,6 +82,7 @@ Narrower (use these while iterating):
 | Annotate / I/O / state | `pytest tests/functional_tests/ -v` |
 | Examples | `pytest tests/test_examples/ -v` |
 | Visualization | `pytest packages/leapp-visualization/tests/ -v` (Python 3.11+) |
+| TensorRT optimize / catalogs | `pytest packages/isaac_deploy_trt/tests/ tests/unit_tests/test_optimize_graph.py tests/unit_tests/test_tensorrt_leapp_bundle.py -v` |
 | Warp tracing | see `warp_support.md` |
 
 ### Warp runtime build
@@ -140,14 +152,26 @@ diff (`git diff origin/main...HEAD`) is the true delta.
 | `guides/export.rst` | Backends and `export_with` |
 | `guides/graph.rst` | State, `annotate.module`, feedback |
 | `guides/buffers.rst` | `static_outputs`, `mirror_leapp_tags` |
-| `guides/debugging.rst` | Validation, logs, dry_run, `non_traced` |
+| `guides/debugging.rst` | Validation, logs, dry_run, `non_traced`, TensorRT optimize |
 | `leapp_runtime.rst` | `InferenceManager` |
 | `semantics/` | `TensorSemantics`, kinds, `TemporalAxis` |
 | `generated_configs.rst` | YAML shape |
 | `api/index.rst` | Public signatures |
 
 User-facing integration recipes belong here, not in this file.
 
+TensorRT GraphSurgeon optimizations, op catalogs, and LEAPP DAG merge live in
+`packages/isaac_deploy_trt`. After `compile_graph()`,
+`optimize_graph(trt_compatible=True)` optimizes ONNX in the graph directory
+and updates YAML checksums / `parameters.optimization` /
+`parameters.tensorrt_version`. When
+`--tensorrt` is omitted, `check`, `optimize`, `merge-pipeline --optimize`,
+and that optimize pass use the minimal catalog
+(`tensorrt_catalogs/minimal/tensorrt_layers.json`): operators and native
+layers supported by every checked-in release and platform. Merge a DAG
+for one Isaac `TensorRTNode` with
+`isaac_deploy_trt merge-pipeline exported.yaml -o pipeline.onnx --optimize`.
+
 ## Guardrails when changing behavior
 
 - Prefer extending `TracedData` once (`_map_structure`, index lowering,

docs/source/api/index.rst

@@ -16,6 +16,7 @@ annotations.
    #     or annotate.input_tensors() / annotate.output_tensors()
    leapp.stop()
    leapp.compile_graph(visualize=True)
+   leapp.optimize_graph(trt_compatible=True)
 
 Graph lifecycle
 ===============
@@ -210,6 +211,43 @@ The method performs the complete pipeline:
    running.
 #. Log graph statistics.
 
+``leapp.optimize_graph()``
+--------------------------
+
+Post-compile graph optimizations. Call after ``compile_graph()``.
+
+Signature
+~~~~~~~~~
+
+.. code-block:: python
+
+   leapp.optimize_graph(
+       trt_compatible: bool = False,
+   )
+
+Parameters
+~~~~~~~~~~
+
+* ``trt_compatible`` (bool, optional): Optimize exported ONNX files in the
+  graph output directory so TensorRT can parse them. Always edits UINT8
+  intermediate Casts, Resize ``antialias``, and rank-9 Reshape. Also
+  decomposes operators the minimal TensorRT catalog marks unsupported.
+  That catalog is the intersection of every checked-in release and platform
+  (``tensorrt_catalogs/minimal/tensorrt_layers.json``). Updates YAML
+  checksums, ``parameters.optimization``, and ``parameters.tensorrt_version``.
+  Graph I/O names, dtypes,
+  and shapes are unchanged. Requires ``isaac_deploy_trt``. Defaults to
+  ``False``.
+
+Behavior
+~~~~~~~~
+
+* Requires ``stop()`` and ``compile_graph()`` first (the graph YAML must
+  exist in the output directory).
+* With no options enabled, logs a warning and returns ``{}``.
+* Returns a dict of applied passes (for example
+  ``{"trt_compatible": <bundle result>}``).
+
 ``GraphConfigs``
 ----------------
 

docs/source/generated_configs.rst

@@ -145,6 +145,28 @@ The ``models`` mapping is deterministically ordered. LEAPP guarantees:
       use these fields to verify that the model file has not changed since
       export.
 
+   ``opset_version``
+      Present on ONNX backends. The ONNX opset used at export.
+
+   ``optimization``
+      Present on ONNX backends unless
+      ``backend_params={"annotate_tensorrt_compatible": False}``.
+      ``true`` means every ONNX operator in the exported graph is supported
+      by the minimal TensorRT catalog: the intersection of every checked-in
+      release and platform. If ``validate_tensorrt_build`` was enabled, it
+      also means the installed TensorRT package parsed and built the model.
+      This does **not** switch the runtime ``backend``; Isaac Deploy / Triton
+      still load ``backend: onnx`` unless the deployer maps this flag itself.
+      Operators that fail the catalog are omitted from the YAML. They are
+      written to ``log.txt`` as ``tensorrt_blockers <model.onnx>: ...``.
+      Catalog ``true`` is not a guarantee that TensorRT can build engines
+      with dynamic shapes or shape-tensor inputs without optimization profiles.
+
+   ``tensorrt_version``
+      Written next to ``optimization``. ``minimal`` means validation and
+      optimize used the intersection of every checked-in TensorRT release.
+      A release such as ``10.16`` means that catalog was selected.
+
 ``pipeline``
 ------------
 

docs/source/guides/debugging.rst

@@ -176,3 +176,46 @@ Disable session patches
 LEAPP patches for the tracing session (for example NumPy and torch
 interop). Use this only to test whether those patches are involved in
 a failure.
+
+TensorRT catalog and optimize
+==============================
+
+TensorRT op-catalog checks, GraphSurgeon optimizations, and DAG merge live in
+the in-tree ``isaac_deploy_trt`` package (also the ``graph-surgery`` extra),
+not in the LEAPP tracing CLI. After ``compile_graph()``, optimize exported
+ONNX in place with ``leapp.optimize_graph(trt_compatible=True)``. That
+updates YAML checksums and ``parameters.optimization``.
+
+With no ``--tensorrt`` flag, ``check``, ``optimize``, ``optimize-dir``, and
+``merge-pipeline --optimize`` use the minimal catalog,
+``tensorrt_catalogs/minimal/tensorrt_layers.json``. An ONNX operator or
+native layer counts as supported only when every checked-in TensorRT release
+and platform marks it supported. ``leapp.optimize_graph(trt_compatible=True)``
+uses that catalog for unsupported-operator optimizations. UINT8 intermediate
+Casts, Resize ``antialias``, and rank-9 Reshape edits always run.
+
+``--tensorrt`` selects one release, for example ``10.16``. ``--platform``
+(``x86``, ``orin``, ``thor``, or ``dgx``) applies only together with
+``--tensorrt``. ``--tensorrt 11.2`` with no platform uses the built-in 11.2
+parser table. ``--catalog`` replaces those flags with one JSON file.
+
+.. code-block:: bash
+
+   pip install -e ".[graph-surgery]"
+   isaac_deploy_trt check --onnx path/to/model.onnx --strict
+   isaac_deploy_trt check --onnx path/to/model.onnx --tensorrt 10.16 --platform orin --strict
+   isaac_deploy_trt check --onnx-dir path/to/export_dir --build --strict
+   isaac_deploy_trt optimize-dir path/to/export_dir -o path/to/export_dir_trt
+   isaac_deploy_trt optimize-dir path/to/export_dir -o path/to/export_dir_trt --tensorrt 10.16 --platform orin
+   isaac_deploy_trt merge-pipeline path/to/export.yaml -o pipeline.onnx --optimize
+
+``--build`` also runs TensorRT parse/build (requires ``tensorrt``).
+``--strict`` exits non-zero if catalog or TensorRT checks fail.
+
+The YAML fields ``parameters.optimization`` and ``parameters.tensorrt_version``
+record that check. ``tensorrt_version`` is ``minimal`` when no release was
+selected. Operators that fail the catalog are written to ``log.txt`` as
+``tensorrt_blockers <model.onnx>: ...`` and are not stored in the YAML.
+Set ``backend_params={"annotate_tensorrt_compatible": False}`` to omit the
+flag, or ``{"validate_tensorrt": True, "validate_tensorrt_build": True}``
+to require a TensorRT parse/build before marking ``true``.

leapp/__init__.py

@@ -30,7 +30,7 @@
 # side effect, so it is available to tracing/export without explicit setup.
 from .leapp_graph.custom_operator_registry import warp_operator  # noqa: F401
 from .leapp_graph.datatypes.patching import FunctionPatch
-from .leapp import annotate, start, stop, compile_graph
+from .leapp import annotate, start, stop, compile_graph, optimize_graph
 from .utils.enums import InputKindEnum, OutputKindEnum
 from .utils.tensor_description import GraphConfigs, TensorSemantics, TemporalAxis
 
@@ -49,6 +49,7 @@
     "start",
     "stop",
     "compile_graph",
+    "optimize_graph",
     "__version__",
     "TensorSemantics",
     "TemporalAxis",

leapp/backends/onnx_export_backend.py

@@ -3,11 +3,30 @@
 # SPDX-License-Identifier: Apache-2.0
 #
 
+import os
 import torch
 from typing import Tuple
 from leapp.backends.export_backend import ExportBackend, prepare_tensors_for_export, SimplifiedONNXProgram
 from leapp.utils.logging import _get_logger
-import os
+
+
+def log_tensorrt_blockers(onnx_path: str, result) -> None:
+    """Write TensorRT catalog failures to the LEAPP log, not the YAML."""
+    blockers = list(getattr(result, "unsupported_ops", []) or [])
+    blockers.extend(
+        f"unknown:{name}" for name in (getattr(result, "unknown_ops", []) or [])
+    )
+    for error in getattr(result, "errors", []) or []:
+        if str(error).startswith("ONNX ops not "):
+            continue
+        blockers.append(str(error))
+    if not blockers:
+        return
+    _get_logger().warning(
+        f"tensorrt_blockers {os.path.basename(onnx_path)}: {', '.join(blockers)}"
+    )
+
+
 import onnx
 from onnx import AttributeProto, TensorProto, helper
 from onnx import numpy_helper
@@ -20,6 +39,10 @@ class ONNXExportBackend(ExportBackend):
     def get_backend_metadata(self):
         metadata = {}
         metadata['opset_version'] = getattr(self, 'opset_version', _constants.ONNX_DEFAULT_OPSET)
+        tensorrt_compatible = getattr(self, 'tensorrt_compatible', None)
+        if tensorrt_compatible is not None:
+            metadata['optimization'] = bool(tensorrt_compatible)
+            metadata['tensorrt_version'] = getattr(self, 'tensorrt_version', None) or "minimal"
         return metadata
 
     def get_backend_model_type(self):
@@ -253,10 +276,48 @@ def save(self, save_path: str) -> Tuple[str, str, str]:
                 _get_logger().error(
                     f"ONNX model validation warning (model still saved): {e}")
 
+        # Record the YAML ``optimization`` flag (catalog check by default).
+        # Set backend_params={"annotate_tensorrt_compatible": False} to omit it.
+        # Enable parse/build with {"validate_tensorrt": True} and optional
+        # {"validate_tensorrt_build": True, "validate_tensorrt_strict": True}.
+        self._annotate_tensorrt_compatible(onnx_path)
+
         md5sum, sha256sum = self._verify_model_location_and_get_hash(onnx_path)
 
         return onnx_path, md5sum, sha256sum
 
+    def _annotate_tensorrt_compatible(self, onnx_path: str) -> None:
+        annotate = self.backend_params.get('annotate_tensorrt_compatible', True)
+        validate = self.backend_params.get('validate_tensorrt', False)
+        if not annotate and not validate:
+            return
+
+        try:
+            from isaac_deploy_trt.validate import (
+                check_onnx_ops_against_tensorrt,
+                validate_onnx_for_tensorrt,
+            )
+        except ImportError:
+            _get_logger().warning(
+                "isaac_deploy_trt is not installed; skipping TensorRT "
+                "compatibility annotation. pip install -e ./packages/isaac_deploy_trt"
+            )
+            return
+
+        if validate:
+            result = validate_onnx_for_tensorrt(
+                onnx_path,
+                build=self.backend_params.get('validate_tensorrt_build', False),
+                strict=self.backend_params.get('validate_tensorrt_strict', False),
+            )
+        else:
+            result = check_onnx_ops_against_tensorrt(onnx_path)
+
+        log_tensorrt_blockers(onnx_path, result)
+        if annotate:
+            self.tensorrt_compatible = bool(result.ok)
+            self.tensorrt_version = result.tensorrt_version or "minimal"
+
     def compile(self,  m: torch.nn.Module = None) -> SimplifiedONNXProgram:
         raise NotImplementedError(
             "TorchExportBackend does not support compilation, please use torch-script or torch-trace instead")