feature: Dynamo implementation of torch_executed_modules

py/torch_tensorrt/dynamo/_compiler.py

@@ -92,7 +92,7 @@ def cross_compile_for_windows(
     require_full_compilation: bool = _defaults.REQUIRE_FULL_COMPILATION,
     min_block_size: int = _defaults.MIN_BLOCK_SIZE,
     torch_executed_ops: Optional[Collection[Target]] = None,
-    torch_executed_modules: Optional[List[str]] = None,
+    torch_executed_modules: Optional[Collection[str]] = None,
     pass_through_build_failures: bool = _defaults.PASS_THROUGH_BUILD_FAILURES,
     max_aux_streams: Optional[int] = _defaults.MAX_AUX_STREAMS,
     version_compatible: bool = _defaults.VERSION_COMPATIBLE,
@@ -175,7 +175,7 @@ def cross_compile_for_windows(
         require_full_compilation (bool): Require modules to be compiled end to end or return an error as opposed to returning a hybrid graph where operations that cannot be run in TensorRT are run in PyTorch
         min_block_size (int): The minimum number of contiguous TensorRT convertible operations in order to run a set of operations in TensorRT
         torch_executed_ops (Collection[Target]): Set of aten operators that must be run in PyTorch. An error will be thrown if this set is not empty but ``require_full_compilation`` is True
-        torch_executed_modules (List[str]): List of modules that must be run in PyTorch. An error will be thrown if this list is not empty but ``require_full_compilation`` is True
+        torch_executed_modules (Collection[str]): Collection of modules that must be run in PyTorch. An error will be thrown if this collection is not empty but ``require_full_compilation`` is True
         pass_through_build_failures (bool): Error out if there are issues during compilation (only applicable to torch.compile workflows)
         max_aux_stream (Optional[int]): Maximum streams in the engine
         version_compatible (bool): Build the TensorRT engines compatible with future versions of TensorRT (Restrict to lean runtime operators to provide version forward compatibility for the engines)
@@ -291,12 +291,6 @@ def cross_compile_for_windows(
 
     engine_capability = EngineCapability._from(engine_capability)
 
-    if torch_executed_modules is not None and torch_executed_modules:
-        logger.warning(
-            f"Detected torch_executed_modules was non-empty: {torch_executed_modules}"
-            "\nThis feature is unimplemented in Torch-TRT Dynamo currently."
-        )
-
     if use_fp32_acc:
         logger.debug(
             "FP32 accumulation for FP16 matmul layers is enabled. If "
@@ -338,6 +332,9 @@ def cross_compile_for_windows(
         "torch_executed_ops": (
             torch_executed_ops if torch_executed_ops is not None else set()
         ),
+        "torch_executed_modules": (
+            torch_executed_modules if torch_executed_modules is not None else set()
+        ),
         "pass_through_build_failures": pass_through_build_failures,
         "max_aux_streams": max_aux_streams,
         "version_compatible": version_compatible,
@@ -456,7 +453,7 @@ def compile(
     require_full_compilation: bool = _defaults.REQUIRE_FULL_COMPILATION,
     min_block_size: int = _defaults.MIN_BLOCK_SIZE,
     torch_executed_ops: Optional[Collection[Target]] = None,
-    torch_executed_modules: Optional[List[str]] = None,
+    torch_executed_modules: Optional[Collection[str]] = None,
     pass_through_build_failures: bool = _defaults.PASS_THROUGH_BUILD_FAILURES,
     max_aux_streams: Optional[int] = _defaults.MAX_AUX_STREAMS,
     version_compatible: bool = _defaults.VERSION_COMPATIBLE,
@@ -554,7 +551,7 @@ def compile(
         require_full_compilation (bool): Require modules to be compiled end to end or return an error as opposed to returning a hybrid graph where operations that cannot be run in TensorRT are run in PyTorch
         min_block_size (int): The minimum number of contiguous TensorRT convertible operations in order to run a set of operations in TensorRT
         torch_executed_ops (Optional[Collection[Target]]): Set of aten operators that must be run in PyTorch. An error will be thrown if this set is not empty but ``require_full_compilation`` is True
-        torch_executed_modules (Optional[List[str]]): List of modules that must be run in PyTorch. An error will be thrown if this list is not empty but ``require_full_compilation`` is True
+        torch_executed_modules (Optional[Collection[str]]): Collection of modules that must be run in PyTorch. An error will be thrown if this collection is not empty but ``require_full_compilation`` is True
         pass_through_build_failures (bool): Error out if there are issues during compilation (only applicable to torch.compile workflows)
         max_aux_streams (Optional[int]): Maximum streams in the engine
         version_compatible (bool): Build the TensorRT engines compatible with future versions of TensorRT (Restrict to lean runtime operators to provide version forward compatibility for the engines)
@@ -681,12 +678,6 @@ def compile(
 
     engine_capability = EngineCapability._from(engine_capability)
 
-    if torch_executed_modules is not None and torch_executed_modules:
-        logger.warning(
-            f"Detected torch_executed_modules was non-empty: {torch_executed_modules}"
-            "\nThis feature is unimplemented in Torch-TRT Dynamo currently."
-        )
-
     if autocast_low_precision_type is not None:
         if not isinstance(autocast_low_precision_type, (torch.dtype, dtype)):
             raise ValueError(
@@ -749,6 +740,9 @@ def compile(
         "torch_executed_ops": (
             torch_executed_ops if torch_executed_ops is not None else set()
         ),
+        "torch_executed_modules": (
+            torch_executed_modules if torch_executed_modules is not None else set()
+        ),
         "pass_through_build_failures": pass_through_build_failures,
         "max_aux_streams": max_aux_streams,
         "version_compatible": version_compatible,
@@ -1354,6 +1348,7 @@ def preserve_module_specs(
                 gm,
                 min_block_size=settings.min_block_size,
                 torch_executed_ops=settings.torch_executed_ops,
+                torch_executed_modules=settings.torch_executed_modules,
                 require_full_compilation=settings.require_full_compilation,
                 skip_fusion=(num_supported_ops == total_ops),
             )
@@ -1374,6 +1369,7 @@ def preserve_module_specs(
             gm,
             min_block_size=settings.min_block_size,
             torch_executed_ops=settings.torch_executed_ops,
+            torch_executed_modules=settings.torch_executed_modules,
             require_full_compilation=settings.require_full_compilation,
         )
 
@@ -1785,7 +1781,7 @@ def convert_exported_program_to_serialized_trt_engine(
     require_full_compilation: bool = _defaults.REQUIRE_FULL_COMPILATION,
     min_block_size: int = _defaults.MIN_BLOCK_SIZE,
     torch_executed_ops: Optional[Collection[Target]] = None,
-    torch_executed_modules: Optional[List[str]] = None,
+    torch_executed_modules: Optional[Collection[str]] = None,
     pass_through_build_failures: bool = _defaults.PASS_THROUGH_BUILD_FAILURES,
     max_aux_streams: Optional[int] = _defaults.MAX_AUX_STREAMS,
     version_compatible: bool = _defaults.VERSION_COMPATIBLE,
@@ -1887,7 +1883,7 @@ def convert_exported_program_to_serialized_trt_engine(
         require_full_compilation (bool): Require modules to be compiled end to end or return an error as opposed to returning a hybrid graph where operations that cannot be run in TensorRT are run in PyTorch
         min_block_size (int): The minimum number of contiguous TensorRT convertible operations in order to run a set of operations in TensorRT
         torch_executed_ops (Optional[Collection[Target]]): Set of aten operators that must be run in PyTorch. An error will be thrown if this set is not empty but ``require_full_compilation`` is True
-        torch_executed_modules (Optional[List[str]]): List of modules that must be run in PyTorch. An error will be thrown if this list is not empty but ``require_full_compilation`` is True
+        torch_executed_modules (Optional[Collection[str]]): Collection of modules that must be run in PyTorch. An error will be thrown if this collection is not empty but ``require_full_compilation`` is True
         pass_through_build_failures (bool): Error out if there are issues during compilation (only applicable to torch.compile workflows)
         max_aux_streams (Optional[int]): Maximum streams in the engine
         version_compatible (bool): Build the TensorRT engines compatible with future versions of TensorRT (Restrict to lean runtime operators to provide version forward compatibility for the engines)
@@ -2008,12 +2004,6 @@ def convert_exported_program_to_serialized_trt_engine(
 
     engine_capability = EngineCapability._from(engine_capability)
 
-    if torch_executed_modules is not None and torch_executed_modules:
-        logger.warning(
-            f"Detected torch_executed_modules was non-empty: {torch_executed_modules}"
-            "\nThis feature is unimplemented in Torch-TRT Dynamo currently."
-        )
-
     if use_fp32_acc:
         logger.debug(
             "FP32 accumulation for FP16 matmul layers is enabled. If "
@@ -2063,6 +2053,9 @@ def convert_exported_program_to_serialized_trt_engine(
         "torch_executed_ops": (
             torch_executed_ops if torch_executed_ops is not None else set()
         ),
+        "torch_executed_modules": (
+            torch_executed_modules if torch_executed_modules is not None else set()
+        ),
         "pass_through_build_failures": pass_through_build_failures,
         "max_aux_streams": max_aux_streams,
         "version_compatible": version_compatible,

py/torch_tensorrt/dynamo/_settings.py

@@ -82,6 +82,7 @@ class CompilationSettings:
         workspace_size (int): Workspace TRT is allowed to use for the module (0 is default)
         min_block_size (int): Minimum number of operators per TRT-Engine Block
         torch_executed_ops (Collection[Target]): Collection of operations to run in Torch, regardless of converter coverage
+        torch_executed_modules (Collection[str]): Collection of modules to run in Torch
         pass_through_build_failures (bool): Whether to fail on TRT engine build errors (True) or not (False)
         max_aux_streams (Optional[int]): Maximum number of allowed auxiliary TRT streams for each engine
         version_compatible (bool): Provide version forward-compatibility for engine plan files
@@ -147,6 +148,7 @@ class CompilationSettings:
     workspace_size: int = WORKSPACE_SIZE
     min_block_size: int = MIN_BLOCK_SIZE
     torch_executed_ops: Collection[Target] = field(default_factory=set)
+    torch_executed_modules: Collection[str] = field(default_factory=set)
     pass_through_build_failures: bool = PASS_THROUGH_BUILD_FAILURES
     max_aux_streams: Optional[int] = MAX_AUX_STREAMS
     version_compatible: bool = VERSION_COMPATIBLE

py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py

@@ -25,20 +25,26 @@
 from torch_tensorrt.dynamo.partitioning._global_partitioner import (
     TorchTensorRTOperatorSupport,
 )
+from torch_tensorrt.dynamo.partitioning.common import node_in_torch_executed_module
 
 logger = logging.getLogger(__name__)
 
 
 class OpSupportTester(ops.OperatorSupportBase):  # type: ignore
     """Class to determine whether operators within a module are supported"""
 
-    def __init__(self, torch_executed_ops: Collection[Target] = set()) -> None:
+    def __init__(
+        self,
+        torch_executed_ops: Collection[Target] = set(),
+        torch_executed_modules: Collection[str] = set(),
+    ) -> None:
         super().__init__()
 
         # Initialize sets of supported/unsupported operators
         self.supported_operators: Dict[str, int] = {}
         self.unsupported_operators: Dict[str, int] = {}
         self.torch_executed_ops = torch_executed_ops
+        self.torch_executed_modules = torch_executed_modules
 
     def is_node_supported(
         self, submodules: Dict[str, torch.nn.Module], node: torch.fx.Node
@@ -71,6 +77,7 @@ def is_node_supported(
             (node in CONVERTERS or node.op == "get_attr")
             and node_name not in self.torch_executed_ops
             and node.target not in self.torch_executed_ops
+            and not node_in_torch_executed_module(node, self.torch_executed_modules)
         ):
             # If node is a proper, supported computational node, store the operator
             if not node.is_impure() and node.op != "get_attr":
@@ -278,6 +285,7 @@ def partition(
     gm: torch.fx.GraphModule,
     min_block_size: int = MIN_BLOCK_SIZE,
     torch_executed_ops: Collection[Target] = set(),
+    torch_executed_modules: Collection[str] = set(),
     require_full_compilation: bool = REQUIRE_FULL_COMPILATION,
     skip_fusion: bool = False,
 ) -> Tuple[torch.fx.GraphModule, OpSupportTester]:
@@ -288,6 +296,7 @@ def partition(
         gm: FX GraphModule to partition
         min_block_size: Minimum number of operators per TRT-Engine Block
         torch_executed_ops: Collection of operations to run in Torch, regardless of converter coverage
+        torch_executed_modules: Collection of modules to run in Torch
         require_full_compilation: Require that all computational operators be run in TRT
         skip_fusion: Skip fusions found by FxNetAccFusionsFinder
     Returns:
@@ -298,8 +307,24 @@ def partition(
     gm.graph.lint()
     gm.recompile()
 
+    if torch_executed_modules:
+        all_module_types = {
+            module_type
+            for node in gm.graph.nodes
+            for _, module_type in (node.meta.get("nn_module_stack") or {}).values()
+        }
+        unmatched = set(torch_executed_modules) - all_module_types
+        if unmatched:
+            logger.warning(
+                f"The following torch_executed_modules were not found in the graph: "
+                f"{unmatched}. Ensure the module names are fully-qualified class names."
+            )
+
     # Construct
-    supported_ops = OpSupportTester(torch_executed_ops=torch_executed_ops)
+    supported_ops = OpSupportTester(
+        torch_executed_ops=torch_executed_ops,
+        torch_executed_modules=torch_executed_modules,
+    )
     partitioner = TRTPartitioner(
         gm,
         supported_ops,

py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py

@@ -16,6 +16,7 @@
 from torch_tensorrt.dynamo.conversion._ConverterRegistry import (
     ConverterRegistry,
 )
+from torch_tensorrt.dynamo.partitioning.common import node_in_torch_executed_module
 from torch_tensorrt.dynamo.utils import COMPLEX_DTYPES
 
 logger = logging.getLogger(__name__)
@@ -137,13 +138,15 @@ def __init__(
         self,
         support_dict: Optional[SupportDict] = None,
         torch_executed_ops: Collection[Target] = set(),
+        torch_executed_modules: Collection[str] = set(),
     ):
         super().__init__(support_dict)
 
         # Initialize sets of supported/unsupported operators
         self.supported_operators: Dict[str, int] = {}
         self.unsupported_operators: Dict[str, int] = {}
         self.torch_executed_ops: Collection[Target] = torch_executed_ops
+        self.torch_executed_modules: Collection[str] = torch_executed_modules
 
     @staticmethod
     def _has_complex_dtype(node: torch.fx.Node) -> bool:
@@ -208,6 +211,7 @@ def is_node_supported(
             (node in CONVERTERS or node.op == "get_attr")
             and node_name not in self.torch_executed_ops
             and node.target not in self.torch_executed_ops
+            and not node_in_torch_executed_module(node, self.torch_executed_modules)
         ):
             # If node is a proper, supported computational node, store the operator
             if not node.is_impure() and node.op != "get_attr":
@@ -258,6 +262,7 @@ def partition(
     gm: torch.fx.GraphModule,
     min_block_size: int = MIN_BLOCK_SIZE,
     torch_executed_ops: Collection[Target] = set(),
+    torch_executed_modules: Collection[str] = set(),
     require_full_compilation: bool = REQUIRE_FULL_COMPILATION,
 ) -> Tuple[torch.fx.GraphModule, TorchTensorRTOperatorSupport]:
     """Partition an FX GraphModule with aten ops into TRT engines
@@ -267,11 +272,28 @@ def partition(
         gm: FX GraphModule to partition
         min_block_size: Minimum number of operators per TRT-Engine Block
         torch_executed_ops: Collection of operations to run in Torch, regardless of converter coverage
+        torch_executed_modules: Collection of modules to run in Torch
         require_full_compilation: Whether to require that all operators be run in TRT
     Returns:
         torch.fx.GraphModule, TorchTensorRTOperatorSupport
     """
-    supported_ops = TorchTensorRTOperatorSupport(torch_executed_ops=torch_executed_ops)
+    if torch_executed_modules:
+        all_module_types = {
+            module_type
+            for node in gm.graph.nodes
+            for _, module_type in (node.meta.get("nn_module_stack") or {}).values()
+        }
+        unmatched = set(torch_executed_modules) - all_module_types
+        if unmatched:
+            logger.warning(
+                f"The following torch_executed_modules were not found in the graph: "
+                f"{unmatched}. Ensure the module names are fully-qualified class names."
+            )
+
+    supported_ops = TorchTensorRTOperatorSupport(
+        torch_executed_ops=torch_executed_ops,
+        torch_executed_modules=torch_executed_modules,
+    )
     partitioner = TRTPartitioner(
         gm,
         supported_ops,

py/torch_tensorrt/dynamo/partitioning/common.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Dict, List, Optional, Sequence, Set, Tuple
+from typing import Any, Collection, Dict, List, Optional, Sequence, Set, Tuple
 
 import sympy
 import torch
@@ -81,6 +81,33 @@ def _build_submodule_profiles(
     return profiles
 
 
+def node_in_torch_executed_module(
+    node: torch.fx.Node, torch_executed_modules: Collection[str]
+) -> bool:
+    """
+    Determine whether a node traces through a module that should run in Torch.
+    A node matches if any level of its nn_module_stack is one of the torch_executed_modules,
+    matched on fully-qualified class name (e.g. "torchvision.models.resnet.BasicBlock"), which
+    torch.export stores as the second element of each nn_module_stack value.
+    Args:
+        node: FX node to check
+        torch_executed_modules: Collection of fully-qualified module class names to run in Torch
+    Returns:
+        True if the node lies within a torch_executed_module, False otherwise
+    """
+    if not torch_executed_modules:
+        return False
+    stack = node.meta.get("nn_module_stack") or {}
+    for _, module_type in stack.values():
+        if module_type in torch_executed_modules:
+            logger.debug(
+                f"Excluding node {node.name} from TRT because its module type "
+                f"{module_type} is in torch_executed_modules."
+            )
+            return True
+    return False
+
+
 def construct_dynamic_input(
     input_shape: torch.Size,
     input_dtype: torch.dtype,