Merge branch 'main' into aten_size_fix

Author: peri044

.circleci/config.yml

@@ -263,7 +263,7 @@ commands:
     parameters:
       torch-build:
         type: string
-        default: "2.0.0.dev20230120+cu117"
+        default: "2.0.0.dev20230129+cu117"
       torch-build-index:
         type: string
         default: "https://download.pytorch.org/whl/nightly/cu117"
@@ -1026,7 +1026,7 @@ parameters:
   # Nightly platform config
   torch-build:
     type: string
-    default: "2.0.0.dev20230120+cu117"
+    default: "2.0.0.dev20230129+cu117"
   torch-build-index:
     type: string
     default: "https://download.pytorch.org/whl/nightly/cu117"

.github/code-owners.yml

@@ -109,5 +109,4 @@
 "component: fx":
   - "frank-wei"
   - "yinghai"
-  - "842974287"
   - "wushirong"

core/conversion/converters/converter_util.cpp

@@ -85,10 +85,10 @@ nvinfer1::ILayer* add_elementwise(
     const std::string& name) {
   if (self->getType() == nvinfer1::DataType::kFLOAT && other->getType() == nvinfer1::DataType::kINT32) {
     LOG_DEBUG("Type mismatch, casting other to " << self->getType());
-    other = castITensor(ctx, other, self->getType());
+    other = castITensor(ctx, other, self->getType(), name);
   } else if (self->getType() == nvinfer1::DataType::kINT32 && other->getType() == nvinfer1::DataType::kFLOAT) {
     LOG_DEBUG("Type mismatch, casting self to " << other->getType());
-    self = castITensor(ctx, self, other->getType());
+    self = castITensor(ctx, self, other->getType(), name);
   }
   // ensure self to have larger number of dimension
   bool swapSelfOther = false;
@@ -106,13 +106,13 @@ nvinfer1::ILayer* add_elementwise(
       LOG_DEBUG(
           "Element-wise op type promotion adding cast from " << self->getType() << " to " << promo_type << " for layer "
                                                              << name);
-      self = castITensor(ctx, self, promo_type);
+      self = castITensor(ctx, self, promo_type, name);
     }
     if (other->getType() != promo_type) {
       LOG_DEBUG(
           "Element-wise op type promotion adding cast from " << other->getType() << " to " << promo_type
                                                              << " for layer " << name);
-      other = castITensor(ctx, other, promo_type);
+      other = castITensor(ctx, other, promo_type, name);
     }
   }
 

core/conversion/converters/impl/conv_deconv.cpp

@@ -13,7 +13,12 @@ namespace {
 bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args) {
   // Input to conv/deconv
   auto in = args[0].ITensor();
-
+  if (in->getType() == nvinfer1::DataType::kINT32) {
+    LOG_WARNING(
+        "Found type  " << in->getType() << "in aten::convolution, casting to" << nvinfer1::DataType::kFLOAT
+                       << " for compatibility.");
+    in = castITensor(ctx, in, nvinfer1::DataType::kFLOAT);
+  }
   // Conv /deconv parameters
   auto stride = util::toDims(args[3].unwrapToIntList());
   auto padding = util::toDims(args[4].unwrapToIntList());

core/conversion/converters/impl/select.cpp

@@ -149,8 +149,26 @@ auto select_registrations TORCHTRT_UNUSED =
                  // IShuffleLayer removes redundant dimensions
                  auto shuffle_layer = ctx->net->addShuffle(*out);
                  TORCHTRT_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
-                 shuffle_layer->setReshapeDimensions(
-                     util::squeezeDims(out->getDimensions(), dim, !ctx->input_is_dynamic));
+
+                 auto num_zero_dimensions =
+                     util::validateInputDimsForShuffle(out->getDimensions(), ctx->input_is_dynamic);
+                 TORCHTRT_CHECK(
+                     num_zero_dimensions >= 0,
+                     "Detected multiple zero dimensions and dynamic shape in aten::select, "
+                         << "which is not currently supported in TensorRT");
+
+                 // If the input is not dynamic, and the tensor is empty (has some dimension 0)
+                 // Then 0 is no longer a placeholder for inherited dimensions
+                 if (!ctx->input_is_dynamic && (num_zero_dimensions > 0)) {
+                   LOG_DEBUG("Setting zero as a true dimension (not placeholder) in aten::select");
+                   shuffle_layer->setZeroIsPlaceholder(false);
+                 }
+
+                 shuffle_layer->setReshapeDimensions(util::squeezeDims(
+                     out->getDimensions(),
+                     dim,
+                     ctx->input_is_dynamic,
+                     ctx->input_is_dynamic && (num_zero_dimensions > 0)));
                  shuffle_layer->setName(util::node_info(n).c_str());
                  out = shuffle_layer->getOutput(0);
                }

core/ir/Input.cpp

@@ -69,11 +69,16 @@ bool valid_input_dtype(nvinfer1::DataType dtype) {
   }
 }
 
+bool valid_input_domain(std::vector<double> domain) {
+  return (domain.size() == 2) && (domain[0] < domain[1]);
+}
+
 Input::Input(
     std::vector<int64_t> shape,
     at::ScalarType dtype,
     nvinfer1::TensorFormat format,
-    bool dtype_is_user_defined) {
+    bool dtype_is_user_defined,
+    std::vector<double> tensor_domain) {
   if (shape.size() > 5) {
     LOG_WARNING("Verify that this dim size is accepted");
   }
@@ -93,6 +98,11 @@ Input::Input(
           << "), Torch-TensorRT only supports contiguous format (NCHW) except with input type Float32 where channel last (NHWC) is also supported");
   this->format = format;
   this->dtype_is_user_defined = dtype_is_user_defined;
+
+  TORCHTRT_CHECK(
+      valid_input_domain(tensor_domain),
+      "Unsupported tensor domain: [" << tensor_domain[0] << ", " << tensor_domain[1] << ")");
+  this->tensor_domain = tensor_domain;
 }
 
 Input::Input(
@@ -101,7 +111,8 @@ Input::Input(
     std::vector<int64_t> max_shape,
     at::ScalarType dtype,
     nvinfer1::TensorFormat format,
-    bool dtype_is_user_defined) {
+    bool dtype_is_user_defined,
+    std::vector<double> tensor_domain) {
   if (min_shape.size() > 5 || opt_shape.size() > 5 || max_shape.size() > 5) {
     LOG_WARNING("Verify that this dim size is accepted");
   }
@@ -146,6 +157,10 @@ Input::Input(
           << "), Torch-TensorRT only supports contiguous format (NCHW) except with input type Float32 where channel last (NHWC) is also supported");
   this->format = format;
   this->dtype_is_user_defined = dtype_is_user_defined;
+  TORCHTRT_CHECK(
+      valid_input_domain(tensor_domain),
+      "Unsupported tensor domain: [" << tensor_domain[0] << ", " << tensor_domain[1] << ")");
+  this->tensor_domain = tensor_domain;
 }
 
 std::ostream& operator<<(std::ostream& os, const Input& input) {

core/ir/ir.h

@@ -31,19 +31,22 @@ struct Input : torch::CustomClassHolder {
       std::vector<int64_t> shape,
       at::ScalarType dtype = at::kFloat,
       nvinfer1::TensorFormat format = nvinfer1::TensorFormat::kLINEAR,
-      bool dtype_is_user_defined = false);
+      bool dtype_is_user_defined = false,
+      std::vector<double> tensor_domain = std::vector<double>{0, 2});
   Input(
       std::vector<int64_t> min_shape,
       std::vector<int64_t> opt_shape,
       std::vector<int64_t> max_shape,
       at::ScalarType dtype = at::kFloat,
       nvinfer1::TensorFormat format = nvinfer1::TensorFormat::kLINEAR,
-      bool dtype_is_used_defined = false);
+      bool dtype_is_user_defined = false,
+      std::vector<double> tensor_domain = std::vector<double>{0, 2});
 
   friend std::ostream& operator<<(std::ostream& os, const Input& input);
 
   bool input_is_dynamic = false;
   bool dtype_is_user_defined = false;
+  std::vector<double> tensor_domain;
   nvinfer1::Dims input_shape;
   nvinfer1::Dims min;
   nvinfer1::Dims max;

core/lowering/lowering.cpp

@@ -144,6 +144,7 @@ void LowerGraph(std::shared_ptr<torch::jit::Graph>& g, std::vector<torch::jit::I
   passes::UnpackAndCastFull(g, lower_info.getGPUDeviceString());
   passes::ReplaceScalarImplicit(g);
   passes::RewriteInputsWithParams(g, params);
+  passes::ReplaceAtenPad(g);
   LOG_GRAPH(*g);
 }
 

core/lowering/passes/BUILD

@@ -28,6 +28,7 @@ cc_library(
         "remove_dropout.cpp",
         "remove_nops.cpp",
         "remove_unnecessary_casts.cpp",
+        "replace_aten_pad.cpp",
         "rewrite_inputs_with_params.cpp",
         "silu_to_sigmoid_multiplication.cpp",
         "unpack_addmm.cpp",

core/lowering/passes/CMakeLists.txt

@@ -15,6 +15,7 @@ target_sources(${lib_name}
             "${CMAKE_CURRENT_SOURCE_DIR}/remove_nops.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/remove_set_attrs.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/remove_unnecessary_casts.cpp"
+            "${CMAKE_CURRENT_SOURCE_DIR}/replace_aten_pad.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/silu_to_sigmoid_multiplication.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/unpack_addmm.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/unpack_batch_norm.cpp"