remove checkshape_plugin and move vec2Tensor to util

Signed-off-by: Ruoqian Guo <[email protected]>

Author: ruoqianguo

core/conversion/converters/BUILD

@@ -58,6 +58,7 @@ cc_library(
     deps = [
         "@tensorrt//:nvinfer",
         "//core/util:prelude",
+        "//core/util:converter_util",
         "//core/conversion/var",
         "//core/conversion/tensorcontainer",
         "//core/conversion/conversionctx",

core/conversion/converters/impl/expand.cpp

@@ -1,9 +1,9 @@
 #include "NvInfer.h"
 #include "core/conversion/converters/converters.h"
 #include "core/conversion/tensorcontainer/TensorContainer.h"
+#include "core/util/converter_util.h"
 #include "core/util/prelude.h"
 #include "core/util/trt_util.h"
-#include "plugins/checkshape_plugin.h"
 #include "torch/torch.h"
 
 #include <ATen/ATen.h>
@@ -16,33 +16,6 @@ namespace converters {
 namespace impl {
 namespace {
 
-nvinfer1::ILayer* create_plugin(
-    ConversionCtx* ctx,
-    const torch::jit::Node* n,
-    nvinfer1::ITensor* inShape,
-    nvinfer1::ITensor* expandShape,
-    int32_t in_rank,
-    int32_t expand_rank,
-    const char* name) {
-  auto creator = new plugins::CheckShapePluginCreator();
-  std::vector<nvinfer1::PluginField> fields;
-  nvinfer1::PluginField input_rank("input_rank", &in_rank, nvinfer1::PluginFieldType::kINT32, 1);
-  nvinfer1::PluginField output_rank("expand_rank", &expand_rank, nvinfer1::PluginFieldType::kINT32, 1);
-  fields.push_back(input_rank);
-  fields.push_back(output_rank);
-  nvinfer1::PluginFieldCollection collection;
-  collection.nbFields = fields.size();
-  collection.fields = fields.data();
-  auto plugin = creator->createPlugin(name, &collection);
-
-  nvinfer1::ITensor* inputs[] = {inShape, expandShape};
-  auto expandShape_layer = ctx->net->addPluginV2(inputs, 2, *plugin);
-  TRTORCH_CHECK(expandShape_layer, "Unable to create interpolation plugin from node" << *n);
-
-  expandShape_layer->setName("CheckShapePlugin");
-  return expandShape_layer;
-}
-
 void addSliceInput(nvinfer1::Dims& dims, int idx, ConversionCtx* ctx, nvinfer1::ISliceLayer* slice) {
   int32_t rank = static_cast<int32_t>(dims.nbDims);
   int32_t* tmp = new int32_t[rank];
@@ -54,18 +27,12 @@ void addSliceInput(nvinfer1::Dims& dims, int idx, ConversionCtx* ctx, nvinfer1::
   slice->setInput(idx, *t);
 }
 
-nvinfer1::ITensor* vec2Tensor(int32_t* dim, int rank, ConversionCtx* ctx) {
-  const nvinfer1::Dims d{1, {static_cast<int32_t>(rank)}};
-  const nvinfer1::Weights w{nvinfer1::DataType::kINT32, dim, rank};
-  return ctx->net->addConstant(d, w)->getOutput(0);
-}
-
 nvinfer1::ITensor* concat(int max_rank, int old_rank, ConversionCtx* ctx, nvinfer1::ITensor* tensor) {
   if (max_rank - old_rank > 0) {
     int32_t* tmp = new int32_t[max_rank - old_rank];
     for (int i = 0; i < (max_rank - old_rank); i++)
       tmp[i] = 1;
-    auto max_rank_tensor = vec2Tensor(tmp, max_rank - old_rank, ctx);
+    auto max_rank_tensor = util::arrToTensor(tmp, max_rank - old_rank, ctx);
     auto in_shape_tensor = ctx->net->addShape(*tensor)->getOutput(0);
     nvinfer1::ITensor* const args[2] = {max_rank_tensor, in_shape_tensor};
     return ctx->net->addConcatenation(args, 2)->getOutput(0);
@@ -81,11 +48,12 @@ bool add_expand(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor
       "Number of dimensions of the desired expansion must be greater than or equal to the number of input dimensions");
 
   // Validate the expansion. Eg: an input of [3, 1] can be expanded to [1, 3, 4] but not [3, 4, 1]
-  for (int i = expandedDims.nbDims - 1; i >= 0; --i) {
+  for (int64_t i = expandedDims.nbDims - 1; i >= 0; --i) {
     int64_t offset = expandedDims.nbDims - 1 - i;
     int64_t dim = input_dims.nbDims - 1 - offset;
     int64_t size = (dim >= 0) ? input_dims.d[dim] : 1;
     int64_t targetSize = expandedDims.d[i];
+    // In expand layer passing -1 as the size for a dimension means not changing the size of that dimension.
     if (targetSize != -1) {
       if (size != targetSize) {
         if (size != 1) {
@@ -96,6 +64,8 @@ bool add_expand(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor
         }
       }
     } else {
+      // For the new dimensions, the size cannot be set to -1. Eg: an input of [3, 1] can be expanded to [3, -1, 4] but
+      // not [-1, 3, 4].
       if (dim < 0) {
         TRTORCH_THROW_ERROR(
             "The expanded size of the tensor (" << targetSize << ") isn't allowed in a leading, non-existing dimension "
@@ -111,10 +81,10 @@ bool add_expand(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor
   if (num_expand_dims > 0) {
     nvinfer1::Dims reshape_dims;
     reshape_dims.nbDims = expandedDims.nbDims;
-    for (int i = 0; i < num_expand_dims; i++) {
+    for (int64_t i = 0; i < num_expand_dims; i++) {
       reshape_dims.d[i] = 1;
     }
-    for (int i = 0; i < input_dims.nbDims; i++) {
+    for (int64_t i = 0; i < input_dims.nbDims; i++) {
       reshape_dims.d[num_expand_dims + i] = input_dims.d[i];
     }
     // Add a reshape layer to expand dims
@@ -130,7 +100,7 @@ bool add_expand(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor
 
   // Set the stride of non singleton dimension to 1
   std::vector<int64_t> strides_vec(expandedDims.nbDims, 0);
-  for (int i = 0; i < expandedDims.nbDims; i++) {
+  for (int64_t i = 0; i < expandedDims.nbDims; i++) {
     strides_vec[i] = (in->getDimensions().d[i] != 1);
   }
 
@@ -150,22 +120,51 @@ bool add_expand_dynamic(
     ConversionCtx* ctx,
     const torch::jit::Node* n,
     nvinfer1::ITensor* in,
-    nvinfer1::ITensor* expandedDimsTensor) {
-  auto input_shape_tensor = ctx->net->addShape(*in)->getOutput(0);
+    nvinfer1::ITensor* expandedDimsTensor,
+    nvinfer1::Dims expandedDims,
+    bool is_expand_layer) {
+  auto input_dims = in->getDimensions();
   auto input_rank = in->getDimensions().nbDims;
   auto output_rank = expandedDimsTensor->getDimensions().d[0];
   TRTORCH_CHECK(
       input_rank <= output_rank,
       "Number of dimensions of the desired expansion must be greater than or equal to the number of input dimensions");
 
-  // add a plugin to check expandedDimsTensor whether match input_shape_tensor
-  auto expandShape_layer =
-      create_plugin(ctx, n, input_shape_tensor, expandedDimsTensor, input_rank, output_rank, "expandShape");
-  auto _tensor = expandShape_layer->getOutput(0);
+  /* TODO: When the inputs are dynamic, some dimensions of the inputs are indeterminate before setBindingDimensions. For
+     these indeterminate dimensions, we don't validate the expansion. Eg: For an input of [3, -1], we omit the
+     validation of the second dimension. Need to explore a better way to validate the expansion.
+  */
+  // Validate the expansion. Eg: an input of [3, 1] can be expanded to [1, 3, 4] but not [3, 4, 1]
+  for (int64_t i = expandedDims.nbDims - 1; i >= 0; --i) {
+    int64_t offset = expandedDims.nbDims - 1 - i;
+    int64_t dim = input_dims.nbDims - 1 - offset;
+    int64_t size = (dim >= 0) ? input_dims.d[dim] : 1;
+    int64_t targetSize = expandedDims.d[i];
+    // Passing -1 as the size for a dimension means not changing the size of that dimension in expand layer.
+    if (targetSize != -1) {
+      if (size != targetSize) {
+        // if size == -1, we can't validate the expansion before setBindingDimensions.
+        if (!(size == -1 || size == 1)) {
+          TRTORCH_THROW_ERROR(
+              "The expanded size of tensor (" << targetSize << ")"
+                                              << " must match the existing size (" << size << ")"
+                                              << " at dimension " << i);
+        }
+      }
+    } else {
+      // In dynamic expand layer, for the new dimensions, the size cannot be set to -1. Eg: an input of [3, 1] can be
+      // expanded to [3, -1, 4] but not [-1, 3, 4].
+      if (is_expand_layer && dim < 0) {
+        TRTORCH_THROW_ERROR(
+            "The expanded size of the tensor (" << targetSize << ") isn't allowed in a leading, non-existing dimension "
+                                                << i);
+      }
+    }
+  }
 
   size_t max_rank = std::max(input_rank, output_rank);
 
-  // Dimensions are right alignment
+  // Dimensions are right alignment. Eg: an input of [3, 1] and max_rank = 4, the result of concat is [1, 1, 3, 1]
   auto new_input_shape_tensor = concat(max_rank, input_rank, ctx, in);
   // LOG_DEBUG("Expand layer output tensor shape: " << new_output_shape_tensor->getDimensions());
   auto new_output_shape_tensor = expandedDimsTensor;
@@ -189,7 +188,7 @@ bool add_expand_dynamic(
   // min(1, sub(input_shape, 1))
   int32_t* one_vector_tmp = new int32_t[1];
   one_vector_tmp[0] = 1;
-  auto one_vector = vec2Tensor(one_vector_tmp, 1, ctx);
+  auto one_vector = util::arrToTensor(one_vector_tmp, 1, ctx);
   auto x_sub_one = ctx->net->addElementWise(*new_input_shape_tensor, *one_vector, nvinfer1::ElementWiseOperation::kSUB)
                        ->getOutput(0);
   auto strides = ctx->net->addElementWise(*one_vector, *x_sub_one, nvinfer1::ElementWiseOperation::kMIN)->getOutput(0);
@@ -224,8 +223,8 @@ auto expand_registrations TRTORCH_UNUSED =
                  int32_t* tmp = new int32_t[expanded_size_rank];
                  for (int i = 0; i < expanded_size_rank; i++)
                    tmp[i] = expanded_size[i];
-                 auto expandedDimsTensor = vec2Tensor(tmp, expanded_size_rank, ctx);
-                 return add_expand_dynamic(ctx, n, in, expandedDimsTensor);
+                 auto expandedDimsTensor = util::arrToTensor(tmp, expanded_size_rank, ctx);
+                 return add_expand_dynamic(ctx, n, in, expandedDimsTensor, expandedDims, true);
                } else {
                  return add_expand(ctx, n, in, expandedDims);
                }
@@ -239,7 +238,8 @@ auto expand_registrations TRTORCH_UNUSED =
                auto targetDims = targetTensor->getDimensions();
                LOG_DEBUG("(expand_as layer) Expand input from " << input_dims << " to " << targetDims);
                if (ctx->input_is_dynamic) {
-                 return add_expand_dynamic(ctx, n, in, ctx->net->addShape(*targetTensor)->getOutput(0));
+                 return add_expand_dynamic(
+                     ctx, n, in, ctx->net->addShape(*targetTensor)->getOutput(0), targetDims, false);
                } else {
                  return add_expand(ctx, n, in, targetDims);
                }

core/conversion/converters/impl/plugins/BUILD

@@ -10,12 +10,10 @@ config_setting(
 cc_library(
     name = "plugins",
     hdrs = [
-        "interpolate_plugin.h",
-        "checkshape_plugin.h"
+        "interpolate_plugin.h"
     ],
     srcs = [
-        "interpolate_plugin.cpp",
-        "checkshape_plugin.cpp"
+        "interpolate_plugin.cpp"
     ],
     deps = [
         "@tensorrt//:nvinfer",
@@ -39,5 +37,5 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar")
 pkg_tar(
     name = "include",
     package_dir = "core/conversion/converters/impl/plugins",
-    srcs = ["interpolate_plugin.h", "checkshape_plugin.h"],
+    srcs = ["interpolate_plugin.h"],
 )