final revision

Author: cehongwang

.github/workflows/build-test-windows.yml

@@ -135,7 +135,7 @@ jobs:
         pushd .
         cd tests/py/dynamo
         ../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_core_runtime_tests_results.xml runtime/test_000_*
-        ../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_core_partitioning_tests_results.xml partitioning/
+        ../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_core_partitioning_tests_results.xml partitioning/test_000_*
         ../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_core_lowering_tests_results.xml lowering/
         popd
 
@@ -219,6 +219,7 @@ jobs:
         pushd .
         cd tests/py/dynamo
         ../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l1_dynamo_core_tests_results.xml runtime/test_001_*
+        ../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l1_dynamo_core_partitioning_tests_results.xml partitioning/test_001_*
         popd
 
   L1-dynamo-compile-tests:

py/torch_tensorrt/dynamo/partitioning/_resource_partitioner.py

@@ -59,7 +59,7 @@
 
 logger = logging.getLogger(__name__)
 
-MAX_NUM_OF_ENGINES = 40
+MAX_NUM_OF_ENGINES = 50
 ENGINE_COMPILATION_MEMORY_USAGE_MULTIPLIER = 4
 
 
@@ -127,6 +127,9 @@ def partition_graph(self) -> torch.fx.GraphModule:
             subgraphs, subgraph_size_budget=self.calculate_size_budget()
         )
 
+        if len(subgraphs) == 1:
+            return self.module
+
         # Set the number of TRT engines to be generated
         self.num_trt_accelerated_subgraphs = len([s for s in subgraphs if s.is_acc])
 
@@ -315,7 +318,7 @@ def break_subgraph_by_size(
             )
             new_subgraphs = self.step_and_validate(new_subgraphs, step_size)
             size_0, size_1 = self.size_of_subgraphs(new_subgraphs)
-            if size_0 > size_to_break:
+            if size_0 > size_to_break or size_0 > size_1:
                 break
 
         return new_subgraphs, size_0, size_1