feat: Add support for exempting full-support blocks

- When a graph is fully supported, we can ignore the minimum block size
argument, which is primarily helpful in reducing segmentation. If the
minimum block size is above the number of total operators in the graph,
and we support all of those, the whole graph will run in Torch
regardless. As a result, we can exempt fully supported graphs from the
min block size requirement
- Alternatively, if preferable, we can display a warning in such a case,
but still respect the minimum block size argument

Author: gs-olive

py/torch_tensorrt/dynamo/lowering/_partition.py

@@ -61,6 +61,11 @@ def propose_partitions(self) -> List[Partition]:
         initial_proposed_partitions = super().propose_partitions()
         partitions = {i: part for i, part in enumerate(initial_proposed_partitions)}
 
+        # A graph is fully supported if there is a single partition and all operators are supported/convertible
+        full_support = len(partitions) == 1 and not getattr(
+            self.operator_support, "unsupported_operators", True
+        )
+
         # For each partition, determine whether or not the number of computational operators
         # exceeds the threshold, and if not, remove that partition
         partitions_to_remove = {}
@@ -85,7 +90,11 @@ def propose_partitions(self) -> List[Partition]:
                 ):
                     compute_node_count += 1
 
-            if compute_node_count < self.min_block_size and not exempted_partition:
+            if (
+                compute_node_count < self.min_block_size
+                and not exempted_partition
+                and not full_support
+            ):
                 partitions_to_remove[id] = compute_node_count
 
         # Remove any nodes violating the criteria specified by the user

tests/py/dynamo/backend/test_partitioning.py

@@ -7,22 +7,6 @@
 
 
 class TestPartitioning(TestCase):
-    def test_partition_fully_supported_one_op(self):
-        class FullySupportedOneOp(torch.nn.Module):
-            def __init__(self, *args, **kwargs) -> None:
-                super().__init__(*args, **kwargs)
-
-            def forward(self, x, y):
-                return torch.ops.aten.add.Tensor(x, y)
-
-        fx_graph = torch.fx.symbolic_trace(FullySupportedOneOp())
-        partitioned_graph = partition(deepcopy(fx_graph))
-        self.assertEquals(
-            len(list(partitioned_graph.named_children())),
-            0,
-            "Single operators should not be segmented",
-        )
-
     def test_partition_fully_supported_multi_op(self):
         class FullySupportedMultiOp(torch.nn.Module):
             def __init__(self, *args, **kwargs) -> None: