feat: Improve Dynamo partitioning system

gs-olive · gs-olive · commit 6cfbb597fe14 · 2023-08-04T10:37:16.000-07:00
- Upgrade Dynamo partitioning to use a custom version of the Torch
_SplitterBase for efficiency and optimized usage in the Dynamo case
- Validate existing use cases are still functional, with the same
partitioning schema as before
- Upgrade qualified name checking
diff --git a/py/torch_tensorrt/dynamo/backend/backends.py b/py/torch_tensorrt/dynamo/backend/backends.py
@@ -131,6 +131,11 @@ def _compile_module(
     # Iterate over all components that can be accelerated
     # Generate the corresponding TRT Module for those
     for name, _ in partitioned_module.named_children():
+
+        # Criteria for a module to be convertible to TRT
+        if "_run_on_acc" not in name:
+            continue
+
         submodule = getattr(partitioned_module, name)
 
         # Get submodule inputs
diff --git a/py/torch_tensorrt/dynamo/conversion/converter_registry.py b/py/torch_tensorrt/dynamo/conversion/converter_registry.py
@@ -305,7 +305,8 @@ def unique_targets(self):
         """Returns the set of unique converter targets stored across all registries"""
         return set.union(*[set(registry.keys()) for registry in self.registries])
 
-    def qualified_name_or_str(self, target: Target) -> str:
+    @staticmethod
+    def qualified_name_or_str(target: Target) -> str:
         """Returns string representation of an FX Node target"""
         if isinstance(target, str):
             return target
diff --git a/py/torch_tensorrt/dynamo/lowering/_partition.py b/py/torch_tensorrt/dynamo/lowering/_partition.py
@@ -1,14 +1,23 @@
 import logging
-from typing import Dict, List, Optional, Sequence, Set
+from typing import Dict, List, Optional, Sequence, Set, Tuple
 
 import torch
 
+from torch.fx.passes.splitter_base import (
+    Subgraph,
+    _SplitterBase,
+    _SplitterSettingBase,
+    FxNetAccNodesFinder,
+    FxNetAccFusionsFinder,
+)
+import torch.fx.passes.operator_support as ops
+from torch.fx.passes.tools_common import NodeSet, CALLABLE_NODE_OPS
+from torch.fx.node import Target
+
+from torch_tensorrt.dynamo.conversion.converter_registry import ConverterRegistry
 from torch_tensorrt.dynamo.lowering import SUBSTITUTION_REGISTRY
 from torch_tensorrt.dynamo._defaults import MIN_BLOCK_SIZE
-from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
-from torch.fx.graph_module import GraphModule
 from torch.fx.node import _get_qualified_name
-from torch.fx.passes.operator_support import OperatorSupport
 
 from torch_tensorrt.dynamo import DYNAMO_CONVERTERS as CONVERTERS
 
@@ -21,93 +30,11 @@
 )
 
 
-class TRTPartitioner(CapabilityBasedPartitioner):
-    """Partitioner to split an FX graph into subgraphs based on operator support
-
-    Args:
-        graph_module: FX GraphModule to partition
-        operator_support: OperatorSupport class describing allowed operators
-        non_compute_ops: Operators which are not considered computational (e.g. getattr)
-        allowed_single_node_partition_ops: Nodes which can be included in single-node partitons.
-            Generally useful for module-level exclusion ops which are intensive despite being single functions
-        min_block_size: Minimum number of computational operators per block
-    Returns:
-        torch.fx.GraphModule
-    """
-
-    def __init__(
-        self,
-        graph_module: GraphModule,
-        operator_support: OperatorSupport,
-        *,
-        non_compute_ops: Optional[Sequence[str]] = None,
-        allowed_single_node_partition_ops: Optional[
-            Sequence[str]
-        ] = DEFAULT_SINGLE_NODE_PARTITIONS,
-        min_block_size=MIN_BLOCK_SIZE,
-    ) -> None:
-        super().__init__(
-            graph_module,
-            operator_support,
-            allows_single_node_partition=True,
-            non_compute_ops=non_compute_ops,
-            allowed_single_node_partition_ops=allowed_single_node_partition_ops,
-        )
-
-        self.min_block_size = min_block_size
-
-    def propose_partitions(self) -> List[Partition]:
-        # Propose partitions using the default, then refine the results
-        initial_proposed_partitions = super().propose_partitions()
-        partitions = {i: part for i, part in enumerate(initial_proposed_partitions)}
-
-        # For each partition, determine whether or not the number of computational operators
-        # exceeds the threshold, and if not, remove that partition
-        partitions_to_remove = {}
-        for id, partition in partitions.items():
-            default_non_compute_ops = {"torch.ops.aten.view", "_operator.getitem"}
-            non_compute_ops = default_non_compute_ops.union(set(self.non_compute_ops))
-            exempted_partition = False
-
-            compute_node_count = 0
-            for node in partition.nodes:
-                # 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)
-                    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
-                ):
-                    compute_node_count += 1
-
-            if compute_node_count < self.min_block_size and not exempted_partition:
-                partitions_to_remove[id] = compute_node_count
-
-        # Remove any nodes violating the criteria specified by the user
-        for id, count in partitions_to_remove.items():
-            logger.debug(
-                f"Removing partition which has {count} < {self.min_block_size} computational operators"
-            )
-            del partitions[id]
-
-        return [partitions[k] for k in sorted(partitions.keys())]
-
-    def partition_and_fuse(self) -> GraphModule:
-        partitions = self.propose_partitions()
-        fused_gm = self.fuse_partitions(partitions)
-        return fused_gm
-
-
-class TorchTensorRTOperatorSupport(OperatorSupport):
+class OpSupportTester(ops.OperatorSupportBase):
     """Class to determine whether operators within a module are supported"""
 
-    def __init__(self, support_dict=None, torch_executed_ops=set()):
-        super().__init__(support_dict)
+    def __init__(self, torch_executed_ops: Sequence[Target] = set()) -> None:
+        super().__init__()
 
         # Initialize sets of supported/unsupported operators
         self.supported_operators = {}
@@ -117,11 +44,7 @@ def __init__(self, support_dict=None, torch_executed_ops=set()):
     def is_node_supported(
         self, submodules: Dict[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
@@ -164,11 +87,139 @@ def print_support_overview(self, num_trt_blocks: Optional[int] = None):
             logger.debug("\nAll Nodes Supported\n")
 
 
+class TRTPartitioner(_SplitterBase):
+    """Partitioner to split an FX graph into subgraphs based on operator support
+
+    Adapted from, and modified for the Torch-TensorRT Dynamo case:
+    https://github.com/pytorch/pytorch/blob/93f538db355ea10c684a57f7a632ed03292ef98f/torch/fx/passes/splitter_base.py#L256C9-L871
+
+    Args:
+        graph_module: FX GraphModule to partition
+        operator_support: OperatorSupport class describing allowed operators
+        allowed_single_node_partition_ops: Nodes which can be included in single-node partitons.
+            Generally useful for module-level exclusion ops which are intensive despite being single functions
+        min_block_size: Minimum number of computational operators per block
+    Returns:
+        torch.fx.GraphModule
+    """
+
+    def __init__(
+        self,
+        graph_module: torch.fx.GraphModule,
+        operator_support: ops.OperatorSupportBase,
+        allowed_single_node_partition_ops: Optional[
+            Sequence[str]
+        ] = DEFAULT_SINGLE_NODE_PARTITIONS,
+        min_block_size: int = MIN_BLOCK_SIZE,
+    ):
+        """
+        Preprocesses graph before splitting:
+        - finds nodes supported by ACC,
+        - finds fusion groups for ACC nodes having non-tensor IO,
+        - builds a graph of direct dependencies,
+        - builds a map of fused nodes to their fusions.
+        As a result we get self.acc_nodes, self.deps and self.fusions.
+        """
+        assert isinstance(graph_module, torch.fx.GraphModule)
+
+        self.graph_module = graph_module
+
+        self.settings = _SplitterSettingBase(
+            min_acc_module_size=min_block_size, allow_non_tensor=True
+        )
+        self.operator_support = operator_support
+
+        # Get all accelerated nodes based on operator support conditions
+        self.acc_nodes = FxNetAccNodesFinder(
+            self.graph_module, self.operator_support, self.settings.allow_non_tensor
+        )()
+
+        if self.settings.skip_fusion:
+            self.fusions = {}
+        else:
+            self.fusions = FxNetAccFusionsFinder(graph_module, self.acc_nodes)()
+
+        # Modify deps to add more deps for fused nodes
+        self.deps = self.find_deps()
+        self.update_deps_for_fusions()
+
+        self.non_acc_submodule_name = "_run_on_gpu_"
+        self._node_submodule_map: Dict[str, str] = {}
+
+        self.num_trt_accelerated_subgraphs = None
+        self.allowed_single_node_partition_ops = allowed_single_node_partition_ops
+
+    def remove_small_acc_subgraphs(self, subgraphs: List[Subgraph]) -> List[Subgraph]:
+        """
+        This pass finds ACC submodules with less than specified size and merges
+        them with adjacent GPU submodules.
+        """
+        result: List[Subgraph] = []
+        for subgraph in subgraphs:
+            if subgraph.is_acc:
+                if len(subgraph.nodes) >= self.settings.min_acc_module_size or any(
+                    ConverterRegistry.qualified_name_or_str(node.target)
+                    in self.allowed_single_node_partition_ops
+                    for node in subgraph.nodes
+                ):
+                    result.append(subgraph)
+                else:
+                    logger.debug(
+                        "Eliminating acc subgraph because it's smaller than the threshold: "
+                        f"{len(subgraph.nodes)} < {self.settings.min_acc_module_size}"
+                    )
+                    if result:
+                        result[-1].nodes.extend(subgraph.nodes)
+                    else:
+                        subgraph.is_acc = False
+                        result.append(subgraph)
+            else:
+                if result and not result[-1].is_acc:
+                    result[-1].nodes.extend(subgraph.nodes)
+                else:
+                    result.append(subgraph)
+        return result
+
+    def partition_graph(self) -> torch.fx.GraphModule:
+        """Partitions the GraphModule into subgraphs based on operator support
+
+        Returns a GraphModule with submodules for each segment
+        """
+        # Delegate nodes based on operator coverage
+        subgraphs = self.put_nodes_into_subgraphs()
+
+        # Remove segments smaller than the block size (with exceptions)
+        subgraphs = self.remove_small_acc_subgraphs(subgraphs)
+
+        # Set the number of TRT engines to be generated
+        self.num_trt_accelerated_subgraphs = len([s for s in subgraphs if s.is_acc])
+
+        # Tag the accelerated nodes and split the graph accordingly
+        self.tag(subgraphs)
+        return self.split()
+
+    def starter_nodes(self) -> Tuple[NodeSet, NodeSet]:
+        """Generates starter nodes for partitioning + segmentation"""
+        # Starter accelerated nodes are all callable accelerated ops
+        starter_acc_nodes = {
+            node for node in self.acc_nodes if node.op in CALLABLE_NODE_OPS
+        }
+
+        # Started non-accelerated nodes are the rest of the callable nodes
+        starter_non_acc_nodes = {
+            node
+            for node in self.graph_module.graph.nodes
+            if (node not in starter_acc_nodes and node.op in CALLABLE_NODE_OPS)
+        }
+
+        return starter_non_acc_nodes, starter_acc_nodes
+
+
 def partition(
     gm: torch.fx.GraphModule,
     verbose: bool = True,
     min_block_size: int = MIN_BLOCK_SIZE,
-    torch_executed_ops: Sequence[str] = set(),
+    torch_executed_ops: Sequence[Target] = set(),
 ) -> torch.fx.GraphModule:
     """Partition an FX GraphModule with aten ops into TRT engines
     Partitioning is based on converter operator support
@@ -181,18 +232,21 @@ def partition(
     Returns:
         torch.fx.GraphModule
     """
-    supported_ops = TorchTensorRTOperatorSupport(torch_executed_ops=torch_executed_ops)
+    # Ensure graph is clean prior to partitioning
+    gm.graph.eliminate_dead_code()
+    gm.graph.lint()
+    gm.recompile()
+
+    # Construct
+    supported_ops = OpSupportTester(torch_executed_ops=torch_executed_ops)
     partitioner = TRTPartitioner(gm, supported_ops, min_block_size=min_block_size)
 
-    # Determine partitions based on user specifications and operator support
-    # Then, fuse partitions and display overview of supported/unsupported operators
-    partitions = partitioner.propose_partitions()
-    fused_graph = partitioner.fuse_partitions(partitions)
+    partitioned_graph = partitioner.partition_graph()
 
     if verbose:
-        supported_ops.print_support_overview(len(partitions))
+        supported_ops.print_support_overview(partitioner.num_trt_accelerated_subgraphs)
 
-    return fused_graph
+    return partitioned_graph
 
 
 def get_submod_inputs(
diff --git a/py/torch_tensorrt/dynamo/lowering/_partition_old.py b/py/torch_tensorrt/dynamo/lowering/_partition_old.py
