Seperate the resource partitioning from adjacency partitioning

Author: cehongwang

py/torch_tensorrt/dynamo/_compiler.py

@@ -860,6 +860,20 @@ def preserve_module_specs(
             require_full_compilation=settings.require_full_compilation,
         )
 
+    from torch_tensorrt.dynamo.partitioning._resource_partitioner import (
+        resource_partition as resource_partitioner_partition,
+    )
+
+    partitioned_module, _ = resource_partitioner_partition(
+        gm,
+        partitioned_module,
+        min_block_size=settings.min_block_size,
+        torch_executed_ops=settings.torch_executed_ops,
+        require_full_compilation=settings.require_full_compilation,
+        skip_fusion=(num_supported_ops == total_ops),
+        cpu_memory_budget=settings.cpu_memory_budget,
+    )
+
     dryrun_tracker.unsupported_ops = supported_ops.unsupported_operators
 
     # The global partitioner leaves non-TRT nodes as-is

py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py

@@ -234,7 +234,7 @@ def partition_graph(self) -> torch.fx.GraphModule:
 
         subgraphs = self.break_subgraphs(
             subgraphs,
-            subgraph_size_budget=500 * 1024 * 1024,  # self.calculate_size_budget()
+            subgraph_size_budget=1000 * 1024 * 1024,  # self.calculate_size_budget()
         )
 
         # Set the number of TRT engines to be generated