diff --git a/backends/nvidia/tensorrt/__init__.py b/backends/nvidia/tensorrt/__init__.py index 8b75df599b1..73144d28a32 100644 --- a/backends/nvidia/tensorrt/__init__.py +++ b/backends/nvidia/tensorrt/__init__.py @@ -9,3 +9,25 @@ This module provides TensorRT delegate support for accelerating PyTorch models on NVIDIA GPUs. """ + +import os +import platform +import sys + +# On Jetson (aarch64 with JetPack), make system and user-installed packages +# available inside venvs. This is needed for TensorRT (JetPack system package) +# and optionally for packages like onnx installed via pip at system level. +if platform.machine() == "aarch64" and os.path.exists("/etc/nv_tegra_release"): + _py_ver = f"python{sys.version_info.major}.{sys.version_info.minor}" + for _pkg_dir in [ + f"/usr/lib/{_py_ver}/dist-packages", + os.path.expanduser(f"~/.local/lib/{_py_ver}/site-packages"), + ]: + if os.path.isdir(_pkg_dir) and _pkg_dir not in sys.path: + sys.path.append(_pkg_dir) + +from executorch.backends.nvidia.tensorrt.backend import TensorRTBackend + +__all__ = [ + "TensorRTBackend", +] diff --git a/backends/nvidia/tensorrt/backend.py b/backends/nvidia/tensorrt/backend.py new file mode 100644 index 00000000000..a5237e3082f --- /dev/null +++ b/backends/nvidia/tensorrt/backend.py @@ -0,0 +1,42 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""TensorRT backend implementation for ExecuTorch.""" + +from typing import final, List + +from executorch.exir.backend.backend_details import ( + BackendDetails, + CompileSpec, + PreprocessResult, +) +from torch.export.exported_program import ExportedProgram + + +@final +class TensorRTBackend(BackendDetails): + """TensorRT backend for accelerating models on NVIDIA GPUs. + + This backend compiles ExecuTorch edge programs to TensorRT engines + for optimized inference on NVIDIA hardware. + """ + + @staticmethod + def preprocess( + edge_program: ExportedProgram, + compile_specs: List[CompileSpec], + ) -> PreprocessResult: + """Compile edge program to TensorRT engine. + + Args: + edge_program: The edge dialect program to compile. + compile_specs: Backend-specific compilation options. + + Returns: + PreprocessResult containing the serialized TensorRT engine. + """ + + return PreprocessResult(processed_bytes=b"") diff --git a/backends/nvidia/tensorrt/targets.bzl b/backends/nvidia/tensorrt/targets.bzl index c8a8755f221..60e4ce0fbaa 100644 --- a/backends/nvidia/tensorrt/targets.bzl +++ b/backends/nvidia/tensorrt/targets.bzl @@ -7,10 +7,25 @@ def define_common_targets(): TARGETS and BUCK files that call this function. """ + runtime.python_library( + name = "backend", + srcs = [ + "backend.py", + ], + visibility = ["PUBLIC"], + deps = [ + "//caffe2:torch", + "//executorch/exir/backend:backend_details", + ], + ) + runtime.python_library( name = "tensorrt", srcs = [ "__init__.py", ], visibility = ["PUBLIC"], + deps = [ + ":backend", + ], )