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]>

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:

WORKSPACE

@@ -41,27 +41,83 @@ local_repository(
 new_local_repository(
     name = "cuda",
     build_file = "@//third_party/cuda:BUILD",
-    path = "/usr/local/cuda-12.1/",
+    path = "/usr/local/cuda-11.8/",
 )
 
+new_local_repository(
+    name = "cublas",
+    build_file = "@//third_party/cublas:BUILD",
+    path = "/usr",
+)
 #############################################################################################################
 # Tarballs and fetched dependencies (default - use in cases when building from precompiled bin and tarballs)
 #############################################################################################################
 
+# Pytorch 1.13
+# http_archive(
+#     name = "libtorch",
+#     build_file = "@//third_party/libtorch:BUILD",
+#     sha256 = "0a013dceedb252f4965b666a2ad772d962135597db5889bd5d43644697c17dbc",
+#     strip_prefix = "libtorch",
+#     urls = ["https://download.pytorch.org/libtorch/cu117/libtorch-cxx11-abi-shared-with-deps-1.13.0%2Bcu117.zip"],
+# )
+
+# http_archive(
+#     name = "libtorch_pre_cxx11_abi",
+#     build_file = "@//third_party/libtorch:BUILD",
+#     sha256 = "cdbd43985ad9d5886793d5dc455d665cf3fd4b4617ef1094479678ff210ed0af",
+#     strip_prefix = "libtorch",
+#     urls = ["https://download.pytorch.org/libtorch/cu117/libtorch-shared-with-deps-1.13.0%2Bcu117.zip"],
+# )
+
+# Pyt 1.14
+# http_archive(
+#     name = "libtorch",
+#     build_file = "@//third_party/libtorch:BUILD",
+#     sha256 = "28f352d980191492758fbeb6773eda509d3d1917c59a7f52829af1c6b521bbc7",
+#     strip_prefix = "libtorch",
+#     urls = ["https://download.pytorch.org/libtorch/nightly/cu117/libtorch-cxx11-abi-shared-with-deps-1.14.0.dev20221114%2Bcu117.zip"],
+# )
+
+# http_archive(
+#     name = "libtorch_pre_cxx11_abi",
+#     build_file = "@//third_party/libtorch:BUILD",
+#     sha256 = "6e1ee2893cfd2ce3d69604d9988e0f5133aeee45da611edd07cb633f78966ce3",
+#     strip_prefix = "libtorch",
+#     urls = ["https://download.pytorch.org/libtorch/nightly/cu117/libtorch-shared-with-deps-1.14.0.dev20221114%2Bcu117.zip"],
+# )
+
+# Pytorch 2.0
+# http_archive(
+#     name = "libtorch",
+#     build_file = "@//third_party/libtorch:BUILD",
+#     # sha256 = "78f2646bbc46024d6be6db06eac19aea08206b1b84a065af04574a8d92ee3f01",
+#     strip_prefix = "libtorch",
+#     urls = ["https://download.pytorch.org/libtorch/cu117/libtorch-cxx11-abi-shared-with-deps-2.0.0%2Bcu117.zip"],
+# )
+#
+# http_archive(
+#     name = "libtorch_pre_cxx11_abi",
+#     build_file = "@//third_party/libtorch:BUILD",
+#     # sha256 = "d76af42fd10554b0a9c0e50fd13e0988ea8d2d9912d2bc659d0019ed047882cc",
+#     strip_prefix = "libtorch",
+#     urls = ["https://download.pytorch.org/libtorch/cu117/libtorch-shared-with-deps-2.0.0%2Bcu117.zip"],
+# )
+
 http_archive(
     name = "libtorch",
     build_file = "@//third_party/libtorch:BUILD",
-    sha256 = "5ba55259b65e071346a2b547b8d1378595f1467a39aaa923fecb09f134f1bcba",
+    sha256 = "999becce82b73e566d0ffe010cd21fea8cf3a33f90f09dcc6b01150b820ae063",
     strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/nightly/cu121/libtorch-cxx11-abi-shared-with-deps-2.1.0.dev20230619%2Bcu121.zip"],
+    urls = ["https://download.pytorch.org/libtorch/nightly/cu118/libtorch-cxx11-abi-shared-with-deps-2.1.0.dev20230605%2Bcu118.zip"],
 )
 
 http_archive(
     name = "libtorch_pre_cxx11_abi",
     build_file = "@//third_party/libtorch:BUILD",
-    sha256 = "8f6661bfc11597e77400e9e36cc8dd8e5e385ba82361d630d9ccede8518d7c7e",
+    sha256 = "786cc728c63ea69c40bd8fb535cf8e5e1dfff1d43eaad3eb5256b9ed89c1b268",
     strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/nightly/cu121/libtorch-shared-with-deps-2.1.0.dev20230619%2Bcu121.zip"],
+    urls = ["https://download.pytorch.org/libtorch/nightly/cu118/libtorch-shared-with-deps-2.1.0.dev20230605%2Bcu118.zip"],
 )
 
 # Download these tarballs manually from the NVIDIA website
@@ -71,23 +127,57 @@ http_archive(
 http_archive(
     name = "cudnn",
     build_file = "@//third_party/cudnn/archive:BUILD",
-    sha256 = "79d77a769c7e7175abc7b5c2ed5c494148c0618a864138722c887f95c623777c",
-    strip_prefix = "cudnn-linux-x86_64-8.8.1.3_cuda12-archive",
+    strip_prefix = "cudnn-linux-x86_64-8.5.0.96_cuda11-archive",
     urls = [
-        "https://developer.nvidia.com/downloads/compute/cudnn/secure/8.8.1/local_installers/12.0/cudnn-linux-x86_64-8.8.1.3_cuda12-archive.tar.xz",
+        "file:////home/dperi/Downloads/cudnn/cudnn-linux-x86_64-8.5.0.96_cuda11-archive.tar.xz",
     ],
 )
 
+# http_archive(
+#     name = "tensorrt",
+#     urls = ["http://cuda-repo/release-candidates/Libraries/TensorRT/v8.5/8.5.1.7-d95eb1a4/11.8-r520/Linux-x64-agnostic/tar/TensorRT-8.5.1.7.Linux.x86_64-gnu.cuda-11.8.cudnn8.6.tar.gz",],
+#     sha256 = "39cc7f077057d1363794e8ff51c4cf21a5dbeccf1116b0020ba0dae0f3063076",
+#     build_file = "@//third_party/tensorrt/archive:BUILD",
+#     strip_prefix = "TensorRT-8.5.1.7"
+# )
+
 http_archive(
     name = "tensorrt",
+    # urls = ["http://cuda-repo/release-candidates/Libraries/TensorRT/v8.5/8.5.1.7-d95eb1a4/11.8-r520/Linux-x64-agnostic/tar/TensorRT-8.5.1.7.Linux.x86_64-gnu.cuda-11.8.cudnn8.6.tar.gz",],
+    urls = ["file:////home/dperi/Downloads/public_trt/TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz",],
+    # sha256 = "39cc7f077057d1363794e8ff51c4cf21a5dbeccf1116b0020ba0dae0f3063076",
     build_file = "@//third_party/tensorrt/archive:BUILD",
-    sha256 = "0f8157a5fc5329943b338b893591373350afa90ca81239cdadd7580cd1eba254",
-    strip_prefix = "TensorRT-8.6.1.6",
-    urls = [
-        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/8.6.1/tars/TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-12.0.tar.gz",
-    ],
+    strip_prefix = "TensorRT-8.6.1.6"
 )
 
+# http_archive(
+#     name = "tensorrt",
+#     # urls = ["http://cuda-repo/release-candidates/Libraries/TensorRT/v8.5/8.5.1.7-d95eb1a4/11.8-r520/Linux-x64-agnostic/tar/TensorRT-8.5.1.7.Linux.x86_64-gnu.cuda-11.8.cudnn8.6.tar.gz",],
+#     urls = ["file:////home/dperi/Downloads/public_trt/TensorRT-8.5.2.2.Linux.x86_64-gnu.cuda-11.8.cudnn8.6.tar.gz",],
+#     # sha256 = "39cc7f077057d1363794e8ff51c4cf21a5dbeccf1116b0020ba0dae0f3063076",
+#     build_file = "@//third_party/tensorrt/archive:BUILD",
+#     strip_prefix = "TensorRT-8.5.2.2"
+# )
+
+# http_archive(
+#     name = "cudnn",
+#     build_file = "@//third_party/cudnn/archive:BUILD",
+#     # sha256 = "d07e15cf1f59ae822ce9cdb04e051cb936b99b727e586f2b61ed04a5f08s35bd2",
+#     strip_prefix = "cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive",
+#     urls = [
+#         "file:////home/dperi/Downloads/cudnn/cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive.tar.xz",
+#     ],
+# )
+
+# http_archive(
+#     name = "tensorrt",
+#     urls = ["http://cuda-repo/release-candidates/Libraries/TensorRT/v8.4/8.4.3.1-74b5c2dc/11.6-r510/Linux-x64-agnostic/tar/TensorRT-8.4.3.1.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz",],
+#     sha256 = "8d7c2085c1639dcc73875048c23598a8526ce3089136876e31d90258e49e4f61",
+#     build_file = "@//third_party/tensorrt/archive:BUILD",
+#     strip_prefix = "TensorRT-8.4.3.1"
+# )
+
+
 ####################################################################################
 # Locally installed dependencies (use in cases of custom dependencies or aarch64)
 ####################################################################################
@@ -111,11 +201,11 @@ http_archive(
 #    build_file = "third_party/libtorch/BUILD"
 #)
 
-#new_local_repository(
+# new_local_repository(
 #    name = "cudnn",
 #    path = "/usr/",
 #    build_file = "@//third_party/cudnn/local:BUILD"
-#)
+# )
 
 #new_local_repository(
 #   name = "tensorrt",

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()