refactor: Move dynamo tests to tests/py/dynamo

Signed-off-by: Dheeraj Peri <[email protected]>

chore: Apply linting

Signed-off-by: Dheeraj Peri <[email protected]>

chore: Reset WORKSPACE

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

.circleci/config.yml

@@ -519,7 +519,7 @@ commands:
           command: |
             set -e
             mkdir -p /tmp/artifacts/test_results
-            cd tests/py
+            cd tests/py/ts/
             pytest --junitxml=/tmp/artifacts/test_results/api/api_test_results.xml api/
             pytest --junitxml=/tmp/artifacts/test_results/models/models_test_results.xml models/
             pytest --junitxml=/tmp/artifacts/test_results/integrations/integrations_test_results.xml integrations/
@@ -733,49 +733,49 @@ commands:
 # =================== FX tests end ======================== #
 
 # =================== Dynamo tests start ======================== #
-  test-dynamo-fx_ts:
-    description: "Test the Dynamo fx_ts_compat path"
+
+  test-dynamo-torch_compile:
+    description: "Test Dynamo torch_compile tests"
     steps:
       - run:
-          name: Run Dynamo fx_ts_compat core tests
+          name: Run Dynamo torch_compile tests
           command: |
-            cd py/torch_tensorrt/dynamo/fx_ts_compat/test
-            pushd core/
-            pytest --junitxml=/tmp/artifacts/test_results/dynamo/fx_ts_compat/test_results.xml
+            cd tests/py/dynamo/backend/
+            pytest --junitxml=/tmp/artifacts/test_results/dynamo/torch_compile/test_results.xml
             popd
 
       - store_test_results:
           path: /tmp/artifacts
       - store_artifacts:
           path: /tmp/testlogs
 
-  test-dynamo-compile-core:
-    description: "Test the Dynamo compile path"
+  test-dynamo-models_torch_compile:
+    description: "Test the Dynamo models via torch_compile path"
     steps:
       - run:
-          name: Run Dynamo compile core tests
+          name: Run Dynamo models via torch_compile path
           command: |
-            cd py/torch_tensorrt/dynamo/backend
-            pushd test/
-            pytest --junitxml=/tmp/artifacts/test_results/dynamo/backend/test_results.xml
+            cd tests/py/dynamo/models
+            pip3 install timm
+            pip3 install transformers
+            pytest test_models.py --junitxml=/tmp/artifacts/test_results/dynamo/backend/test_results.xml --ir torch_compile
             popd
 
       - store_test_results:
           path: /tmp/artifacts
       - store_artifacts:
           path: /tmp/testlogs
 
-  test-dynamo-compile:
-    description: "Test the Dynamo compile path"
+  test-dynamo-models_torch_export:
+    description: "Test the Dynamo models via torch_export path"
     steps:
       - run:
-          name: Run Dynamo compile E2E tests
+          name: Run Dynamo models via torch_export path
           command: |
-            cd py/torch_tensorrt/dynamo/
-            pushd test/
+            cd tests/py/dynamo/models
             pip3 install timm
             pip3 install transformers
-            pytest --junitxml=/tmp/artifacts/test_results/dynamo/backend/test_results.xml --ir dynamo_compile
+            pytest test_models_export.py --junitxml=/tmp/artifacts/test_results/dynamo/backend/test_results.xml --ir dynamo
             popd
 
       - store_test_results:
@@ -1039,9 +1039,9 @@ jobs:
           command: pip3 install --pre /tmp/dist/x86_64-linux/*cp39-cp39*.whl
       # We install torch after torch-trt because pip automatically enforces the version constraint otherwise
       - dump-test-env
-      - test-dynamo-compile
-      - test-dynamo-compile-core
-      - test-dynamo-fx_ts
+      - test-dynamo-torch_compile
+      - test-dynamo-models_torch_compile
+      - test-dynamo-models_torch_export
 
   package-x86_64-linux:
     parameters:

py/torch_tensorrt/dynamo/backend/lowering/_fusers.py

@@ -69,4 +69,4 @@ def fuse_permute_linear(gm: torch.fx.GraphModule):
     gm.graph.eliminate_dead_code()
     gm.graph.lint()
     gm.recompile()
-    return gm
+    return gm

py/torch_tensorrt/dynamo/compile.py

@@ -74,11 +74,6 @@ def compile(
         + "torch_executed_ops, pass_through_build_failures}"
     )
 
-    if "use_experimental_fx_rt" in kwargs:
-        use_experimental_rt = kwargs["use_experimental_fx_rt"]
-
-    logger.info(f"Using {'C++' if use_experimental_rt else 'Python'} TRT Runtime")
-
     if not isinstance(inputs, collections.abc.Sequence):
         inputs = [inputs]
 

py/torch_tensorrt/dynamo/utils.py

@@ -10,6 +10,7 @@
 
 COSINE_THRESHOLD = 0.99
 
+
 def use_python_runtime_parser(use_python_runtime: Optional[bool] = None) -> bool:
     """Parses a user-provided input argument regarding Python runtime
 
@@ -40,6 +41,7 @@ def use_python_runtime_parser(use_python_runtime: Optional[bool] = None) -> bool
 
     return using_python_runtime
 
+
 def cosine_similarity(gt_tensor, pred_tensor):
     gt_tensor = gt_tensor.flatten().to(torch.float32)
     pred_tensor = pred_tensor.flatten().to(torch.float32)
@@ -51,6 +53,7 @@ def cosine_similarity(gt_tensor, pred_tensor):
 
     return res
 
+
 def prepare_inputs(
     inputs: Union[_Input.Input, torch.Tensor, Sequence, Dict],
     device: torch.device = torch.device("cuda"),

py/torch_tensorrt/dynamo/backend/test/test_backend_compiler.py renamed to tests/py/dynamo/backend/test_backend_compiler.py

@@ -1,8 +1,8 @@
+import torch
+import torch_tensorrt
 from torch_tensorrt.dynamo.backend.lowering import partition
 from torch.testing._internal.common_utils import run_tests, TestCase
-import torch
 from copy import deepcopy
-from torch_tensorrt.dynamo.backend import compile
 from utils import lower_graph_testing, DECIMALS_OF_AGREEMENT
 
 
@@ -33,8 +33,9 @@ def forward(self, x, y):
         torch._dynamo.reset()
 
         # Validate that the results between Torch and Torch-TRT are similar
-        optimized_model = compile(
+        optimized_model = torch_tensorrt.compile(
             fx_graph,
+            "torch_compile",
             inputs,
             min_block_size=1,
             pass_through_build_failures=True,
@@ -101,8 +102,9 @@ def forward(self, x, y):
         torch._dynamo.reset()
 
         # Validate that the results between Torch and Torch-TRT are similar
-        optimized_model = compile(
+        optimized_model = torch_tensorrt.compile(
             fx_graph,
+            "torch_compile",
             inputs,
             min_block_size=1,
             pass_through_build_failures=True,
@@ -143,8 +145,9 @@ def forward(self, x, y):
         ]
 
         # Validate that the results between Torch and Torch-TRT are similar
-        optimized_model = compile(
+        optimized_model = torch_tensorrt.compile(
             fx_graph,
+            "torch_compile",
             inputs,
             min_block_size=1,
             pass_through_build_failures=True,

py/torch_tensorrt/dynamo/backend/test/test_compiler_utils.py renamed to tests/py/dynamo/backend/test_compiler_utils.py

@@ -2,7 +2,7 @@
 from utils import lower_graph_testing, DECIMALS_OF_AGREEMENT
 from torch.testing._internal.common_utils import run_tests, TestCase
 import torch
-from torch_tensorrt.dynamo.backend import compile
+import torch_tensorrt
 
 
 class TestLowering(TestCase):
@@ -162,8 +162,12 @@ def forward(self, x, y, z):
         torch._dynamo.reset()
 
         # Validate that the results between Torch and Torch-TRT are similar
-        optimized_model = compile(
-            fx_graph, inputs, min_block_size=1, pass_through_build_failures=True
+        optimized_model = torch_tensorrt.compile(
+            fx_graph,
+            "torch_compile",
+            inputs,
+            min_block_size=1,
+            pass_through_build_failures=True,
         )
         optimized_model_results = optimized_model(*inputs).detach().cpu()
         torch_model_results = fx_graph(*inputs).detach().cpu()

py/torch_tensorrt/dynamo/backend/test/test_decompositions.py renamed to tests/py/dynamo/backend/test_decompositions.py

@@ -1,7 +1,7 @@
 import torch
+import torch_tensorrt
 from utils import lower_graph_testing
 from torch.testing._internal.common_utils import run_tests, TestCase
-from torch_tensorrt.dynamo.backend import compile
 
 
 class TestMaxPool1D(TestCase):
@@ -39,8 +39,12 @@ def forward(self, x):
         torch._dynamo.reset()
 
         # Validate that the results between Torch and Torch-TRT are similar
-        optimized_model = compile(
-            fx_graph, inputs, min_block_size=1, pass_through_build_failures=True
+        optimized_model = torch_tensorrt.compile(
+            fx_graph,
+            "torch_compile",
+            inputs,
+            min_block_size=1,
+            pass_through_build_failures=True,
         )
         optimized_model_results = optimized_model(*inputs).detach().cpu()
         torch_model_results = fx_graph(*inputs).detach().cpu()
@@ -85,8 +89,12 @@ def forward(self, x, y):
         torch._dynamo.reset()
 
         # Validate that the results between Torch and Torch-TRT are similar
-        optimized_model = compile(
-            fx_graph, inputs, min_block_size=1, pass_through_build_failures=True
+        optimized_model = torch_tensorrt.compile(
+            fx_graph,
+            "torch_compile",
+            inputs,
+            min_block_size=1,
+            pass_through_build_failures=True,
         )
         optimized_model_results = optimized_model(*inputs).detach().cpu()
         torch_model_results = fx_graph(*inputs).detach().cpu()

py/torch_tensorrt/dynamo/backend/test/test_partitioning.py renamed to tests/py/dynamo/backend/test_partitioning.py

@@ -1,7 +1,7 @@
 from utils import lower_graph_testing
 from torch.testing._internal.common_utils import run_tests, TestCase
 import torch
-from torch_tensorrt.dynamo.backend import compile
+import torch_tensorrt
 
 
 class TestFakeTensors(TestCase):
@@ -40,8 +40,12 @@ def forward(self, x):
         torch._dynamo.reset()
 
         # Validate that the results between Torch and Torch-TRT are similar
-        optimized_model = compile(
-            fx_graph, inputs, min_block_size=1, pass_through_build_failures=True
+        optimized_model = torch_tensorrt.compile(
+            fx_graph,
+            "torch_compile",
+            inputs,
+            min_block_size=1,
+            pass_through_build_failures=True,
         )
         optimized_model_results = optimized_model(*inputs).detach().cpu()
         torch_model_results = fx_graph(*inputs).detach().cpu()
@@ -92,8 +96,12 @@ def forward(self, x):
         torch._dynamo.reset()
 
         # Validate that the results between Torch and Torch-TRT are similar
-        optimized_model = compile(
-            fx_graph, inputs, min_block_size=1, pass_through_build_failures=True
+        optimized_model = torch_tensorrt.compile(
+            fx_graph,
+            "torch_compile",
+            inputs,
+            min_block_size=1,
+            pass_through_build_failures=True,
         )
         optimized_model_results = optimized_model(*inputs).detach().cpu()
         torch_model_results = fx_graph(*inputs).detach().cpu()