Merge remote-tracking branch 'dynamo_changes' into sample_backend

Author: gs-olive

.circleci/config.yml

@@ -508,6 +508,7 @@ commands:
       - store_artifacts:
           path: /tmp/testlogs
 
+# =================== FX tests start ======================== #
   test-fx_core:
     description: "Test the fx core"
     steps:
@@ -707,6 +708,167 @@ commands:
       - store_artifacts:
           path: /tmp/testlogs
 
+# =================== FX tests end ======================== #
+
+# =================== Dynamo tests start ======================== #
+  test-dynamo-fx_ts_core:
+    description: "Test the Dynamo core"
+    steps:
+      - run:
+          name: Run Dynamo core tests
+          command: |
+            cd py/torch_tensorrt/dynamo/fx_ts_compat/test
+            pushd core/
+            pytest --junitxml=/tmp/artifacts/test_results/dynamo/core/test_results.xml
+            popd
+
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
+  test-dynamo-fx_ts_converters_acc:
+    description: "Test the Dynamo acc converters"
+    steps:
+      - run:
+          name: Run FX converter tests
+          command: |
+            cd py/torch_tensorrt/dynamo/fx_ts_compat/test
+            pushd converters/acc_op/
+            pytest --junitxml=/tmp/artifacts/test_results/dynamo/converters/acc_op/test_results.xml
+            popd
+
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
+  test-dynamo-fx_ts_converters_aten:
+    description: "Test the dynamo aten converters"
+    steps:
+      - run:
+          name: Run dynamo converter tests
+          command: |
+            cd py/torch_tensorrt/dynamo/fx_ts_compat/test
+            pushd converters/aten_op/
+            pytest --junitxml=/tmp/artifacts/test_results/dynamo/converters/aten_op/test_results.xml
+            popd
+
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
+  test-dynamo-fx_ts_converters_vanilla:
+    description: "Test the dynamo vanilla converters"
+    steps:
+      - run:
+          name: Run dynamo converter tests
+          command: |
+            cd py/torch_tensorrt/dynamo/fx_ts_compat/test
+            pushd converters/vanilla/
+            pytest --junitxml=/tmp/artifacts/test_results/dynamo/converters/vanilla/test_results.xml
+            popd
+
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
+  test-dynamo-fx_ts_passes:
+    description: "Test the dynamo passes"
+    steps:
+      - run:
+          name: Run dynamo passes
+          command: |
+            cd py/torch_tensorrt/dynamo/fx_ts_compat/test
+            pushd passes
+            list_passes=$(ls | grep -v test_setitem*)
+            pytest $list_passes --junitxml=/tmp/artifacts/test_results/dynamo/passes/test_results.xml
+            popd
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
+  test-dynamo-fx_ts_tools:
+    description: "Test the dynamo tools"
+    steps:
+      - run:
+          name: Run dynamo tools
+          command: |
+            cd py/torch_tensorrt/dynamo/fx_ts_compat/test
+            pushd tools
+            pytest --junitxml=/tmp/artifacts/test_results/dynamo/tools/test_results.xml
+            popd
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
+  test-dynamo-fx_ts_trt_lower:
+    description: "Test the dynamo TRT lowering"
+    steps:
+      - run:
+          name: Run dynamo TRT lowering
+          command: |
+            cd py/torch_tensorrt/dynamo/fx_ts_compat/test
+            pushd trt_lower
+            pytest --junitxml=/tmp/artifacts/test_results/dynamo/trt_lower/test_results.xml
+            popd
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
+  test-dynamo-fx_ts_tracer:
+    description: "Test all dynamo tracers"
+    steps:
+      - run:
+          name: Run dynamo tracer
+          command: |
+            cd py/torch_tensorrt/dynamo/fx_ts_compat/test
+            pushd tracer
+            list_tracer=$(ls | grep -v test_dispatch_*)
+            pytest $list_tracer --junitxml=/tmp/artifacts/test_results/fx/tracer/test_results.xml
+            popd
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
+  test-dynamo-fx_ts_tracer_acc:
+    description: "Test the dynamo acc tracer only"
+    steps:
+      - run:
+          name: Run dynamo tracer
+          command: |
+            cd py/torch_tensorrt/dynamo/fx_ts_compat/test
+            pushd tracer
+            list_tracer=$(ls | grep test_acc)
+            pytest $list_tracer --junitxml=/tmp/artifacts/test_results/dynamo/tracer/test_results.xml
+            popd
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
+  test-dynamo-fx_ts_quant:
+    description: "Test the dynamo quant"
+    steps:
+      - run:
+          name: Run dynamo quant tests
+          command: |
+            cd py/torch_tensorrt/dynamo/fx_ts_compat/test
+            pushd quant/
+            pytest --junitxml=/tmp/artifacts/test_results/dynamo/quant/test_results.xml
+            popd
+
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
   test-dynamo-torch_compile:
     description: "Test the Dynamo torch_compile path"
     steps:
@@ -719,11 +881,55 @@ commands:
             pip3 install transformers
             pytest --junitxml=/tmp/artifacts/test_results/dynamo/test_results.xml --ir torch_compile
             popd
+
       - store_test_results:
           path: /tmp/artifacts
       - store_artifacts:
           path: /tmp/testlogs
 
+  test-dynamo-fx_ts:
+    description: "Test the dynamo backend"
+    steps:
+      - run:
+          name: Run dynamo tests
+          command: |
+            mkdir -p /tmp/artifacts/test_results
+      - test-dynamo-fx_ts_converters_acc
+      - test-dynamo-fx_ts_converters_aten
+      - test-dynamo-fx_ts_converters_vanilla
+      - test-dynamo-fx_ts_passes
+      - test-dynamo-fx_ts_tools
+      - test-dynamo-fx_ts_trt_lower
+      - test-dynamo-fx_ts_tracer
+      - test-dynamo-fx_ts_core
+      - test-dynamo-fx_ts_quant
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
+  test-dynamo-fx_ts-no-aten:
+    description: "Test the dynamo backend without aten operators"
+    steps:
+      - run:
+          name: Run dynamo tests without aten ops
+          command: |
+            mkdir -p /tmp/artifacts/test_results
+      - test-dynamo-fx_ts_converters_acc
+      - test-dynamo-fx_ts_converters_vanilla
+      - test-dynamo-fx_ts_passes
+      - test-dynamo-fx_ts_tools
+      - test-dynamo-fx_ts_trt_lower
+      - test-dynamo-fx_ts_tracer_acc
+      - test-dynamo-fx_ts_core
+      - test-dynamo-fx_ts_quant
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
+# =================== Dynamo tests end ======================== #
+
 # Define a job to be invoked later in a workflow.
 # See: https://circleci.com/docs/2.0/configuration-reference/#jobs
 jobs:
@@ -930,6 +1136,7 @@ jobs:
       # We install torch after torch-trt because pip automatically enforces the version constraint otherwise
       - dump-test-env
       - test-dynamo-torch_compile
+      - test-dynamo-fx_ts
 
   test-py-dynamo-x86_64-linux-no-aten:
     parameters:
@@ -960,7 +1167,7 @@ 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-torch_compile
+      - test-dynamo-fx_ts-no-aten
 
   package-x86_64-linux:
     parameters:

README.md

@@ -73,6 +73,7 @@ import torch_tensorrt
 ...
 
 trt_ts_module = torch_tensorrt.compile(torch_script_module,
+    # If the inputs to the module are plain Tensors, specify them via the `inputs` argument:
     inputs = [example_tensor, # Provide example tensor for input shape or...
         torch_tensorrt.Input( # Specify input object with shape and dtype
             min_shape=[1, 3, 224, 224],
@@ -81,6 +82,12 @@ trt_ts_module = torch_tensorrt.compile(torch_script_module,
             # For static size shape=[1, 3, 224, 224]
             dtype=torch.half) # Datatype of input tensor. Allowed options torch.(float|half|int8|int32|bool)
     ],
+
+    # For inputs containing tuples or lists of tensors, use the `input_signature` argument:
+    # Below, we have an input consisting of a Tuple of two Tensors (Tuple[Tensor, Tensor])
+    # input_signature = ( (torch_tensorrt.Input(shape=[1, 3, 224, 224], dtype=torch.half),
+    #                      torch_tensorrt.Input(shape=[1, 3, 224, 224], dtype=torch.half)), ),
+
     enabled_precisions = {torch.half}, # Run with FP16
 )
 
@@ -114,7 +121,7 @@ torch.jit.save(trt_ts_module, "trt_torchscript_module.ts") # save the TRT embedd
 These are the following dependencies used to verify the testcases. Torch-TensorRT can work with other versions, but the tests are not guaranteed to pass.
 
 - Bazel 5.2.0
-- Libtorch 2.0.0.dev20230103 (built with CUDA 11.7)
+- Libtorch 2.1.0.dev20230314 (built with CUDA 11.7)
 - CUDA 11.7
 - cuDNN 8.5.0
 - TensorRT 8.5.1.7
@@ -124,7 +131,7 @@ These are the following dependencies used to verify the testcases. Torch-TensorR
 Releases: https://github.com/pytorch/TensorRT/releases
 
 ```
-pip install torch-tensorrt==1.2.0 --find-links https://github.com/pytorch/TensorRT/releases/expanded_assets/v1.2.0
+pip install torch-tensorrt
 ```
 
 ## Compiling Torch-TensorRT

WORKSPACE

@@ -56,17 +56,17 @@ new_local_repository(
 http_archive(
     name = "libtorch",
     build_file = "@//third_party/libtorch:BUILD",
-    sha256 = "8b3b48615169c83c1b643c0efade078ea080b1da598e15fcf01bc59421f3095e",
+    sha256 = "7c4b8754830fef23ec19c5eaf414794cee9597b435df055f5c1d0471d3e81568",
     strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/nightly/cu117/libtorch-cxx11-abi-shared-with-deps-2.0.0.dev20230219%2Bcu117.zip"],
+    urls = ["https://download.pytorch.org/libtorch/nightly/cu117/libtorch-cxx11-abi-shared-with-deps-2.1.0.dev20230314%2Bcu117.zip"],
 )
 
 http_archive(
     name = "libtorch_pre_cxx11_abi",
     build_file = "@//third_party/libtorch:BUILD",
-    sha256 = "aa7fd06079d260ff83c344d043fb84fbd9cf831cf375ed8b5a1b62416817af31",
+    sha256 = "f1e64a75dd12d0ba4c8c1f61947299e0a9c50684dff64f0cfbf355aa7a13e8cf",
     strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/nightly/cu117/libtorch-shared-with-deps-2.0.0.dev20230219%2Bcu117.zip"],
+    urls = ["https://download.pytorch.org/libtorch/nightly/cu117/libtorch-shared-with-deps-2.1.0.dev20230314%2Bcu117.zip"],
 )
 
 # Download these tarballs manually from the NVIDIA website

core/compiler.cpp

@@ -352,8 +352,9 @@ torch::jit::Module CompileGraph(const torch::jit::Module& mod, CompileSpec cfg)
       // Determine if the block is convertible/has collection output, and based on the result,
       // whether full compilation can be expected
       auto isBlockConvertible = conversion::VerifyConverterSupportForBlock(g->block(), true);
+      auto inputIsCollection = conversion::InputIsCollection(g->block());
       auto outputIsCollection = conversion::OutputIsCollection(g->block());
-      auto requires_collection_handling = (isBlockConvertible && outputIsCollection);
+      auto requires_collection_handling = (isBlockConvertible && (inputIsCollection || outputIsCollection));
 
       // Determine whether user specifications necessitate partitioning
       auto isFallbackRequested = userRequestedFallback(cfg);

core/conversion/conversion.cpp

@@ -556,10 +556,20 @@ std::set<std::string> ConvertableOpsInBlock(const torch::jit::Block* b) {
   return convertable_ops;
 }
 
+bool InputIsCollection(const torch::jit::Block* b) {
+  for (auto in : b->inputs()) {
+    if (in->type()->kind() == torch::jit::TypeKind::TupleType || in->type()->kind() == torch::jit::TypeKind::ListType) {
+      return true;
+    }
+  }
+  return false;
+}
+
 bool OutputIsCollection(const torch::jit::Block* b) {
   for (auto out : b->outputs()) {
     if (out->type()->kind() == torch::jit::TypeKind::TupleType ||
-        out->type()->kind() == torch::jit::TypeKind::ListType) {
+        out->type()->kind() == torch::jit::TypeKind::ListType ||
+        out->type()->kind() == torch::jit::TypeKind::DictType) {
       return true;
     }
   }

core/conversion/conversion.h

@@ -26,6 +26,8 @@ std::string ConvertBlockToEngine(
 
 bool OpSupported(const torch::jit::Node* n);
 
+bool InputIsCollection(const torch::jit::Block* b);
+
 bool OutputIsCollection(const torch::jit::Block* b);
 
 bool VerifyConverterSupportForBlock(const torch::jit::Block* b, bool suppress_errors = false);