Merge branch 'main' into bose_fx2trt_converters_slice_select

Author: apbose

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

@@ -68,7 +68,7 @@ c10::optional<torch::jit::IValue> EvaluateNode(ConversionCtx* ctx, const torch::
       return {};
     }
   }
-  auto eval = evaluators::EvalNode(n, eval_args);
+  auto eval = evaluators::EvalNode(ctx, n, eval_args);
   return eval;
 }
 
@@ -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);

core/conversion/converters/impl/matrix_multiply.cpp

@@ -72,6 +72,85 @@ auto mm_registrations TORCHTRT_UNUSED =
                mm_layer->setName(util::node_info(n).c_str());
                auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], mm_layer->getOutput(0));
 
+               LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+               return true;
+             }})
+        .pattern(
+            {"aten::baddbmm(Tensor self, Tensor batch1, Tensor batch2, *, Scalar beta=1, Scalar alpha=1) -> Tensor",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               auto self = args[0].ITensorOrFreeze(ctx);
+               auto bat1 = args[1].ITensorOrFreeze(ctx);
+               auto bat2 = args[2].ITensorOrFreeze(ctx);
+               nvinfer1::Dims batch1Dims = bat1->getDimensions();
+               nvinfer1::Dims batch2Dims = bat2->getDimensions();
+
+               // check dimensions
+               TORCHTRT_CHECK(
+                   batch1Dims.nbDims == 3,
+                   "Expected 3-dimensional tensor, but got "
+                       << batch1Dims.nbDims
+                       << "-dimensional tensor for argument 'batch1' (while checking arguments for baddbmm)");
+               TORCHTRT_CHECK(
+                   batch2Dims.nbDims == 3,
+                   "Expected 3-dimensional tensor, but got "
+                       << batch2Dims.nbDims
+                       << "-dimensional tensor for argument 'batch2' (while checking arguments for baddbmm)");
+               TORCHTRT_CHECK(
+                   batch1Dims.d[0] == batch2Dims.d[0],
+                   "Expected tensor to have size " << batch1Dims.d[0] << " at dimension 0, but got size "
+                                                   << batch2Dims.d[0]
+                                                   << " for argument 'batch2' (while checking arguments for baddbmm)");
+               TORCHTRT_CHECK(
+                   batch1Dims.d[2] == batch2Dims.d[1],
+                   "Expected tensor to have size " << batch1Dims.d[2] << " at dimension 1, but got size "
+                                                   << batch2Dims.d[1]
+                                                   << " for argument 'batch2' (while checking arguments for baddbmm)");
+
+               auto mm_layer = ctx->net->addMatrixMultiply(
+                   *bat1, nvinfer1::MatrixOperation::kNONE, *bat2, nvinfer1::MatrixOperation::kNONE);
+               TORCHTRT_CHECK(mm_layer, "Unable to create matrix multiplication for node: " << *n);
+               mm_layer->setName((util::node_info(n) + "_matmul").c_str());
+
+               auto mm_out = mm_layer->getOutput(0);
+
+               auto alpha = args[4].unwrapToScalar();
+               if (alpha.to<float>() != 1.) {
+                 auto alpha_tensor = scalar_to_tensor(ctx, alpha);
+                 auto alpha_layer = add_elementwise(
+                     ctx,
+                     nvinfer1::ElementWiseOperation::kPROD,
+                     mm_out,
+                     alpha_tensor,
+                     util::node_info(n) + std::string("_alpha_mul"));
+                 TORCHTRT_CHECK(alpha_layer, "Unable to create alpha_mul layer from node: " << *n);
+                 mm_out = alpha_layer->getOutput(0);
+               }
+
+               auto beta = args[3].unwrapToScalar();
+               // If beta is 0, then input will be ignored, and nan and inf in it will not be propagated.
+               if (beta.to<float>() != 0.) {
+                 if (beta.to<float>() != 1.) {
+                   auto beta_tensor = scalar_to_tensor(ctx, beta);
+                   auto beta_layer = add_elementwise(
+                       ctx,
+                       nvinfer1::ElementWiseOperation::kPROD,
+                       self,
+                       beta_tensor,
+                       util::node_info(n) + std::string("_beta_mul"));
+                   TORCHTRT_CHECK(beta_layer, "Unable to create beta_mul layer from node: " << *n);
+                   self = beta_layer->getOutput(0);
+                 }
+                 auto self_add_layer = add_elementwise(
+                     ctx,
+                     nvinfer1::ElementWiseOperation::kSUM,
+                     self,
+                     mm_out,
+                     util::node_info(n) + std::string("_self_add"));
+                 TORCHTRT_CHECK(self_add_layer, "Unable to create self_add layer from node: " << *n);
+                 mm_out = self_add_layer->getOutput(0);
+               }
+
+               auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], mm_out);
                LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
                return true;
              }});

core/conversion/converters/impl/shuffle.cpp

@@ -70,25 +70,37 @@ static auto shuffle_registrations TORCHTRT_UNUSED =
                auto in = args[0].ITensorOrFreeze(ctx);
                auto in_shape = util::toVec(in->getDimensions());
                std::vector<int64_t> new_shape;
+               nvinfer1::ITensor* shape_tensor;
                if (ctx->input_is_dynamic) {
-                 new_shape = util::toVec(args[1].unwrapToIntList().vec());
-                 int nbDynamicDims = 0;
-                 for (size_t i = 0; i < new_shape.size(); i++) {
-                   if (in_shape[i] == -1)
-                     nbDynamicDims++;
-                 }
-                 if (nbDynamicDims > 1) {
-                   TORCHTRT_THROW_ERROR(
-                       "Resize is currently not supported when target shape contains more than one dynamic dimension");
+                 LOG_DEBUG("Using dynamic version of reshape layer");
+                 if (args[1].isITensorList()) {
+                   LOG_DEBUG("Shape tensor is an ITensorList");
+                   auto new_shape = args[1].unwrapToITensorList();
+                   auto concat_layer = ctx->net->addConcatenation(new_shape.data(), new_shape.size());
+                   TORCHTRT_CHECK(concat_layer, "Unable to create concatenation layer from node: " << *n);
+                   concat_layer->setAxis(static_cast<int32_t>(0));
+                   shape_tensor = concat_layer->getOutput(0);
+                 } else if (args[1].isIntList()) {
+                   LOG_DEBUG("Shape tensor is an IntList");
+                   auto shape_vec = args[1].unwrapToIntList().vec();
+                   shape_tensor = tensor_to_const(ctx, torch::tensor(shape_vec).to(torch::kI32));
+                 } else {
+                   LOG_ERROR(
+                       "Invalid IValue type of " << args[1].IValue()->type()
+                                                 << " detected for shape tensor from node: " << *n);
                  }
                } else {
                  new_shape = torch::reshape(torch::rand(in_shape), args[1].unwrapToIntList().vec()).sizes().vec();
                }
-
                auto shuffle = ctx->net->addShuffle(*in);
-               TORCHTRT_CHECK(shuffle, "Unable to create shuffle layer from node: " << *n);
-               shuffle->setReshapeDimensions(util::toDims(new_shape));
                shuffle->setName(util::node_info(n).c_str());
+               TORCHTRT_CHECK(shuffle, "Unable to create shuffle layer from node: " << *n);
+
+               if (ctx->input_is_dynamic) {
+                 shuffle->setInput(1, *shape_tensor);
+               } else {
+                 shuffle->setReshapeDimensions(util::toDims(new_shape));
+               }
 
                auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
                LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());

core/conversion/converters/impl/stack.cpp

@@ -43,10 +43,9 @@ auto stack_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().patt
            auto cont = t.toCustomClass<TensorContainer>();
            itensor = cont->tensor();
          }
-
          auto shuffle_layer = ctx->net->addShuffle(*itensor);
          TORCHTRT_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
-         shuffle_layer->setReshapeDimensions(util::unsqueezeDims(itensor->getDimensions(), dim));
+         shuffle_layer->setReshapeDimensions(util::unsqueezeDims(itensor->getDimensions(), dim, 1, false));
 
          tensors.push_back(shuffle_layer->getOutput(0));
        }

core/conversion/evaluators/NodeEvaluatorRegistry.cpp

@@ -114,9 +114,9 @@ std::vector<std::string> getEvaluatorList() {
   return get_evaluator_registry().GetRegisteredEvaluatorList();
 }
 
-c10::optional<torch::jit::IValue> EvalNode(const torch::jit::Node* n, kwargs& args) {
+c10::optional<torch::jit::IValue> EvalNode(ConversionCtx* ctx, const torch::jit::Node* n, kwargs& args) {
   auto evaluator = get_evaluator_registry().GetEvaluator(n);
-  return evaluator(n, args);
+  return evaluator(ctx, n, args);
 }
 
 void register_node_evaluator(torch::jit::NodeKind node_kind, EvalRegistration eval_reg) {