diff --git a/backends/nvidia/tensorrt/__init__.py b/backends/nvidia/tensorrt/__init__.py index 73144d28a32..8d6d111f375 100644 --- a/backends/nvidia/tensorrt/__init__.py +++ b/backends/nvidia/tensorrt/__init__.py @@ -27,7 +27,9 @@ sys.path.append(_pkg_dir) from executorch.backends.nvidia.tensorrt.backend import TensorRTBackend +from executorch.backends.nvidia.tensorrt.partitioner import TensorRTPartitioner __all__ = [ "TensorRTBackend", + "TensorRTPartitioner", ] diff --git a/backends/nvidia/tensorrt/partitioner/TARGETS b/backends/nvidia/tensorrt/partitioner/TARGETS new file mode 100644 index 00000000000..0a42614a385 --- /dev/null +++ b/backends/nvidia/tensorrt/partitioner/TARGETS @@ -0,0 +1,5 @@ +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +define_common_targets() diff --git a/backends/nvidia/tensorrt/partitioner/__init__.py b/backends/nvidia/tensorrt/partitioner/__init__.py new file mode 100644 index 00000000000..b7ac9d4fe4f --- /dev/null +++ b/backends/nvidia/tensorrt/partitioner/__init__.py @@ -0,0 +1,45 @@ +# 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 partitioner for ExecuTorch.""" + +from typing import Dict, List, Optional + +from executorch.backends.nvidia.tensorrt.backend import TensorRTBackend +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.backend.partitioner import ( + DelegationSpec, + Partitioner, + PartitionResult, +) +from torch.export.exported_program import ExportedProgram + + +class TensorRTPartitioner(Partitioner): + """Partitioner for TensorRT backend. + """ + + def __init__( + self, + compile_specs: Optional[List[CompileSpec]] = None, + ) -> None: + super().__init__() + self.compile_specs = compile_specs or [] + self.delegation_spec = DelegationSpec( + backend_id=TensorRTBackend.__name__, + compile_specs=self.compile_specs, + ) + + def partition(self, exported_program: ExportedProgram) -> PartitionResult: + """Partition the graph for TensorRT delegation. + + Identifies subgraphs that can be lowered to the TensorRT backend. + """ + partition_tags: Dict[str, DelegationSpec] = {} + return PartitionResult( + tagged_exported_program=exported_program, + partition_tags=partition_tags, + ) diff --git a/backends/nvidia/tensorrt/partitioner/targets.bzl b/backends/nvidia/tensorrt/partitioner/targets.bzl new file mode 100644 index 00000000000..4c090b4a767 --- /dev/null +++ b/backends/nvidia/tensorrt/partitioner/targets.bzl @@ -0,0 +1,21 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + """Defines targets that should be shared between fbcode and xplat. + + The directory containing this targets.bzl file should also contain both + TARGETS and BUCK files that call this function. + """ + + runtime.python_library( + name = "partitioner", + srcs = [ + "__init__.py", + ], + visibility = ["PUBLIC"], + deps = [ + "//caffe2:torch", + "//executorch/backends/nvidia/tensorrt:backend", + "//executorch/exir/backend:partitioner", + ], + )