feat: Inform user if no valid partitions

- Fall back to global partitioner if fast partitioner fails

Author: gs-olive

py/torch_tensorrt/dynamo/backend/backends.py

@@ -108,15 +108,47 @@ def _compile_module(
     Returns:
         Compiled FX GraphModule
     """
-    # Partition module into components that can be TRT-accelerated
-    if settings.use_fast_partitioner:
-        partitioned_module = partitioning.fast_partition(
-            gm,
-            verbose=settings.debug,
-            min_block_size=settings.min_block_size,
-            torch_executed_ops=settings.torch_executed_ops,
+    # Check the number of supported operations in the graph
+    num_supported_ops, total_ops = partitioning.get_graph_converter_support(
+        gm, settings.debug, settings.torch_executed_ops
+    )
+
+    # If the number of supported operations is 0 or less than the block size, skip the subgraph
+    # TODO: Add condition to second expression below when require_full_compilation is added
+    if num_supported_ops == 0 or (num_supported_ops < settings.min_block_size):
+        logger.warning(
+            f"{num_supported_ops} supported operations detected in subgraph containing {total_ops} computational nodes. "
+            f"Skipping this subgraph, since min_block_size was detected to be {settings.min_block_size}"
         )
+        return gm
     else:
+        logger.debug(
+            f"Detected support for {num_supported_ops} operators out of {total_ops} in subgraph."
+        )
+
+    # Partition module into components that can be TRT-accelerated
+    fast_partitioner_failed = False
+
+    # If specified, try using the fast partitioner and fall back to the global one on failure
+    if settings.use_fast_partitioner:
+        try:
+            partitioned_module = partitioning.fast_partition(
+                gm,
+                verbose=settings.debug,
+                min_block_size=settings.min_block_size,
+                torch_executed_ops=settings.torch_executed_ops,
+            )
+        except torch.fx.passes.splitter_base.FxNetSplitterInternalError:
+            logger.error(
+                "Partitioning failed on the subgraph with fast partition. See trace above. "
+                + "Retrying with global partition.",
+                exc_info=True,
+            )
+
+            fast_partitioner_failed = True
+            settings.use_fast_partitioner = False
+
+    if not settings.use_fast_partitioner:
         partitioned_module = partitioning.global_partition(
             gm,
             verbose=settings.debug,
@@ -162,4 +194,8 @@ def _compile_module(
     for name, trt_mod in trt_modules.items():
         setattr(partitioned_module, name, trt_mod)
 
+    # Reset settings object to user specification after fallback to global partitioning mode
+    if fast_partitioner_failed:
+        settings.use_fast_partitioner = True
+
     return partitioned_module

py/torch_tensorrt/dynamo/partitioning/__init__.py

@@ -1,3 +1,3 @@
 from ._adjacency_partitioner import partition as fast_partition
 from ._global_partitioner import partition as global_partition
-from .common import get_submod_inputs
+from .common import get_graph_converter_support, get_submod_inputs

py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py

@@ -12,8 +12,10 @@
     _SplitterSettingBase,
 )
 from torch.fx.passes.tools_common import CALLABLE_NODE_OPS, NodeSet
-from torch_tensorrt.dynamo import DYNAMO_CONVERTERS as CONVERTERS
-from torch_tensorrt.dynamo._defaults import MIN_BLOCK_SIZE
+from torch_tensorrt.dynamo._defaults import DEBUG, MIN_BLOCK_SIZE
+from torch_tensorrt.dynamo.conversion.converter_registry import (
+    DYNAMO_CONVERTERS as CONVERTERS,
+)
 from torch_tensorrt.dynamo.conversion.converter_registry import ConverterRegistry
 
 from .common import DEFAULT_SINGLE_NODE_PARTITIONS
@@ -212,7 +214,7 @@ def starter_nodes(self) -> Tuple[NodeSet, NodeSet]:
 
 def partition(
     gm: torch.fx.GraphModule,
-    verbose: bool = True,
+    verbose: bool = DEBUG,
     min_block_size: int = MIN_BLOCK_SIZE,
     torch_executed_ops: Collection[Target] = set(),
 ) -> torch.fx.GraphModule:

py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py

@@ -3,13 +3,13 @@
 
 import torch
 from torch.fx.graph_module import GraphModule
-from torch.fx.node import _get_qualified_name
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.passes.operator_support import OperatorSupport, SupportDict
-from torch_tensorrt.dynamo._defaults import MIN_BLOCK_SIZE
+from torch_tensorrt.dynamo._defaults import DEBUG, MIN_BLOCK_SIZE
 from torch_tensorrt.dynamo.conversion.converter_registry import (
     DYNAMO_CONVERTERS as CONVERTERS,
 )
+from torch_tensorrt.dynamo.conversion.converter_registry import ConverterRegistry
 
 from .common import DEFAULT_SINGLE_NODE_PARTITIONS
 
@@ -69,14 +69,15 @@ def propose_partitions(self) -> List[Partition]:
                 # Partitions are exempted from min_block_size if they contain an allowed single-node op
                 if (
                     node.op == "call_function"
-                    and _get_qualified_name(node.target)
+                    and ConverterRegistry.qualified_name_or_str(node.target)
                     in self.allowed_single_node_partition_ops
                 ):
                     exempted_partition = True
                     break
                 elif (
                     node.op == "call_function"
-                    and _get_qualified_name(node.target) not in non_compute_ops
+                    and ConverterRegistry.qualified_name_or_str(node.target)
+                    not in non_compute_ops
                 ):
                     compute_node_count += 1
 
@@ -118,11 +119,7 @@ def __init__(
     def is_node_supported(
         self, submodules: Mapping[str, torch.nn.Module], node: torch.fx.Node
     ) -> bool:
-        node_name = (
-            _get_qualified_name(node.target)
-            if not isinstance(node.target, str)
-            else node.target
-        )
+        node_name = ConverterRegistry.qualified_name_or_str(node.target)
 
         if node in CONVERTERS and node_name not in self.torch_executed_ops:
             # If node is a proper, supported computational node, store the operator
@@ -142,32 +139,37 @@ def is_node_supported(
 
             return False
 
-    def print_support_overview(self, num_trt_blocks: Optional[int] = None) -> None:
+    def print_support_overview(
+        self, num_trt_blocks: Optional[int] = None, print_node_support: bool = False
+    ) -> None:
         if num_trt_blocks is not None:
             logger.debug(
                 f"\nNumber of TensorRT-Accelerated Engines Generated: {num_trt_blocks}"
             )
 
-        # Reformat support messages for debugger to print node overview as a single string
-        supported_nodes_str = "\nSupported Nodes:\n"
-        for node_name, count in self.supported_operators.items():
-            supported_nodes_str += f"- {node_name} + Operator Count: {count}\n"
+        if print_node_support:
+            # Reformat support messages for debugger to print node overview as a single string
+            supported_nodes_str = "\nSupported Nodes:\n"
+            for node_name, count in self.supported_operators.items():
+                supported_nodes_str += f"- {node_name} + Operator Count: {count}\n"
 
-        logger.debug(supported_nodes_str)
+            logger.debug(supported_nodes_str)
 
-        if self.unsupported_operators:
-            unsupported_nodes_str = "\nUnsupported or Excluded Nodes:\n"
-            for node_name, count in self.unsupported_operators.items():
-                unsupported_nodes_str += f"- {node_name} + Operator Count: {count}\n"
+            if self.unsupported_operators:
+                unsupported_nodes_str = "\nUnsupported or Excluded Nodes:\n"
+                for node_name, count in self.unsupported_operators.items():
+                    unsupported_nodes_str += (
+                        f"- {node_name} + Operator Count: {count}\n"
+                    )
 
-            logger.debug(unsupported_nodes_str)
-        else:
-            logger.debug("\nAll Nodes Supported\n")
+                logger.debug(unsupported_nodes_str)
+            else:
+                logger.debug("\nAll Nodes Supported\n")
 
 
 def partition(
     gm: torch.fx.GraphModule,
-    verbose: bool = True,
+    verbose: bool = DEBUG,
     min_block_size: int = MIN_BLOCK_SIZE,
     torch_executed_ops: Optional[Set[str]] = None,
 ) -> torch.fx.GraphModule:

py/torch_tensorrt/dynamo/partitioning/common.py

@@ -1,8 +1,8 @@
-import logging
-from typing import Any, Optional, Sequence, Set
+from typing import Any, Optional, Sequence, Set, Tuple
 
 import torch
 from torch.fx.node import _get_qualified_name
+from torch_tensorrt.dynamo._defaults import DEBUG
 from torch_tensorrt.dynamo.lowering import SUBSTITUTION_REGISTRY
 
 DEFAULT_SINGLE_NODE_PARTITIONS: Set[str] = {
@@ -11,9 +11,6 @@
 }
 
 
-logger = logging.getLogger(__name__)
-
-
 def get_submod_inputs(
     mod: torch.fx.GraphModule,
     submod: torch.fx.GraphModule,
@@ -39,3 +36,41 @@ def get_input(self: Any, inputs: Sequence[torch.Tensor]) -> None:
     mod(*inputs)
     handle.remove()
     return acc_inputs
+
+
+def get_graph_converter_support(
+    graph_module: torch.fx.GraphModule,
+    verbose: bool = DEBUG,
+    torch_executed_ops: Optional[Set[str]] = None,
+) -> Tuple[int, int]:
+    """Helper function to get converter support overview pre-partitioning
+
+    Args:
+        graph_module: FX GraphModule to determine support for
+        verbose: Bool representing whether to print operator support
+        torch_executed_ops: Collection of operations to run in Torch, regardless of converter coverage
+    Returns:
+        The number of supported call_function nodes in the graph
+    """
+    from ._global_partitioner import TorchTensorRTOperatorSupport
+
+    # Instantiate operator support object and module dictionary
+    op_support = TorchTensorRTOperatorSupport(torch_executed_ops=torch_executed_ops)
+    module_dict = dict(graph_module.named_modules())
+
+    number_of_supported_nodes = 0
+    total_functional_nodes = 0
+
+    # Iterate over all nodes in the graph, enumerating call_function nodes
+    for node in graph_module.graph.nodes:
+        if node.op == "call_function":
+            total_functional_nodes += 1
+
+            if op_support.is_node_supported(module_dict, node):
+                number_of_supported_nodes += 1
+
+    # Print node support overview prior to partitioning
+    if verbose:
+        op_support.print_support_overview(print_node_support=True)
+
+    return number_of_supported_nodes, total_functional_nodes