Merge branch 'main' into anurag.dixit/fix_conv1d

Author: andi4191

.github/workflows/docker_builder.yml

@@ -16,7 +16,7 @@ concurrency:
 
 jobs:
   build:
-    runs-on: linux.2xlarge
+    runs-on: linux.4xlarge.nvidia.gpu
 
     # Define key environment variables
     # Container name is of the form torch_tensorrt:<branch_name>

core/conversion/conversionctx/ConversionCtx.cpp

@@ -164,7 +164,7 @@ void ConversionCtx::RecordNewITensor(const torch::jit::Value* value, nvinfer1::I
 
 std::string ConversionCtx::SerializeEngine() {
 #if NV_TENSORRT_MAJOR > 7
-  auto serialized_network = builder->buildSerializedNetwork(*net, *cfg);
+  auto serialized_network = make_trt(builder->buildSerializedNetwork(*net, *cfg));
   if (!serialized_network) {
     TORCHTRT_THROW_ERROR("Building serialized network failed in TensorRT");
   }

core/conversion/converters/impl/conv_deconv.cpp

@@ -10,6 +10,74 @@ namespace converters {
 namespace impl {
 namespace {
 
+void add_output_padding(nvinfer1::Dims& padding, nvinfer1::Dims& out_padding, bool& has_output_padding) {
+  int nbSpatialDims = out_padding.nbDims;
+  // When there is out_padding, if padding is larger than out_padding, just adjust padding Or reduce out_padding as
+  // minimum as possible.
+  for (int i = 0; i < nbSpatialDims; ++i) {
+    if (padding.d[i] - out_padding.d[i] >= 0) {
+      padding.d[i] -= out_padding.d[i];
+      out_padding.d[i] = 0;
+    } else {
+      // Reduce out_padding as possible.
+      out_padding.d[i] -= padding.d[i];
+      padding.d[i] = 0;
+      has_output_padding = true;
+    }
+  }
+}
+
+nvinfer1::ILayer* add_bias_layer(
+    ConversionCtx* ctx,
+    nvinfer1::ITensor* input_tensor,
+    nvinfer1::Dims& input_dims,
+    nvinfer1::Dims& output_padding,
+    Weights& bias) {
+  nvinfer1::ITensor* input_shape = ctx->net->addShape(*input_tensor)->getOutput(0);
+  // Add padding layer
+  nvinfer1::ITensor* start;
+  nvinfer1::ITensor* totalPadding;
+  auto in_nbDims = input_dims.nbDims;
+  std::vector<int32_t> startVec(in_nbDims, 0);
+  std::vector<int32_t> totalPaddingVec(in_nbDims, 0);
+  int32_t diff = in_nbDims - output_padding.nbDims;
+  for (int32_t i = diff; i < in_nbDims; i++) {
+    int32_t idx = i - diff;
+    startVec[i] = 0; // Don't need begin padding, only post padding
+    totalPaddingVec[i] = output_padding.d[idx];
+  }
+  start = tensor_to_const(ctx, torch::tensor(startVec, torch::kInt32));
+  totalPadding = tensor_to_const(ctx, torch::tensor(totalPaddingVec, torch::kInt32));
+
+  const auto size =
+      ctx->net->addElementWise(*input_shape, *totalPadding, nvinfer1::ElementWiseOperation::kSUM)->getOutput(0);
+
+  nvinfer1::Dims stride;
+  stride.nbDims = in_nbDims;
+  for (int64_t i = 0; i < in_nbDims; i++) {
+    stride.d[i] = 1;
+  }
+  const auto& dummy = stride;
+  auto* sliceLayer = ctx->net->addSlice(*input_tensor, dummy, dummy, stride);
+  sliceLayer->setInput(1, *start);
+  sliceLayer->setInput(2, *size);
+  sliceLayer->setMode(nvinfer1::SliceMode::kFILL);
+  nvinfer1::ITensor* slice_output = sliceLayer->getOutput(0);
+
+  nvinfer1::Dims constantDims;
+  constantDims.nbDims = in_nbDims;
+  for (int64_t i = 0; i < in_nbDims; i++) {
+    constantDims.d[i] = 1;
+  }
+  constantDims.d[diff - 1] =
+      bias.shape.d[0]; // Set C dimension to bias dim and other dimensions to 1 to enable broadcast
+  auto const_layer = ctx->net->addConstant(constantDims, bias.data);
+  auto bias_layer =
+      ctx->net->addElementWise(*slice_output, *const_layer->getOutput(0), nvinfer1::ElementWiseOperation::kSUM);
+
+  return bias_layer;
+}
+
 bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args) {
   // Input to conv/deconv
   auto in = args[0].ITensor();
@@ -93,17 +161,31 @@ bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args)
     nvinfer1::ILayer* layer = nullptr;
     nvinfer1::ITensor* out = nullptr;
     if (transposed) {
-      nvinfer1::IDeconvolutionLayer* deconvLayer =
-          ctx->net->addDeconvolutionNd(*in, num_output_maps, filter_dim, kernel_weights, bias.data);
+      // Fix padding based on output_padding provided
+      nvinfer1::Dims begPadding = padding;
+      bool hasOutputPadding = false;
+      add_output_padding(padding, out_padding, hasOutputPadding);
+
+      nvinfer1::IDeconvolutionLayer* deconvLayer = ctx->net->addDeconvolutionNd(
+          *in, num_output_maps, filter_dim, kernel_weights, hasOutputPadding ? nvinfer1::Weights{} : bias.data);
       deconvLayer->setStrideNd(stride);
       deconvLayer->setDilationNd(dilation);
       deconvLayer->setNbGroups(groups);
-      deconvLayer->setPaddingNd(padding);
+      deconvLayer->setPrePadding(begPadding);
+      deconvLayer->setPostPadding(padding);
+
       // Set deconv kernel weights
       deconvLayer->setInput(1, *kernel);
       TORCHTRT_CHECK(deconvLayer, "Unable to create deconv layer with non-const weights from node: " << *n);
       layer = deconvLayer;
-      out = layer->getOutput(0);
+      out = deconvLayer->getOutput(0);
+      if (hasOutputPadding) {
+        LOG_DEBUG("Padding output deconvolution tensor with:" << out_padding);
+        nvinfer1::ITensor* tensorPtr = deconvLayer->getOutput(0);
+        auto dims = in->getDimensions();
+        layer = add_bias_layer(ctx, tensorPtr, dims, out_padding, bias);
+        out = layer->getOutput(0);
+      }
       if (expand_dims) {
         // Un-expand the expanded dimension
         out = addUnpadding(ctx, n, out, original_dim.nbDims);
@@ -178,20 +260,7 @@ bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args)
     // https://github.com/onnx/onnx-tensorrt/blob/c3cfcbc8248c6bd007e6630af2085df5e4834b42/builtin_op_importers.cpp#L734
     nvinfer1::Dims begPadding = padding;
     bool hasOutputPadding = false;
-    int nbSpatialDims = out_padding.nbDims;
-    // When there is out_padding, if padding is larger than out_padding, just adjust padding Or reduce out_padding as
-    // minimum as possible.
-    for (int i = 0; i < nbSpatialDims; ++i) {
-      if (padding.d[i] - out_padding.d[i] >= 0) {
-        padding.d[i] -= out_padding.d[i];
-        out_padding.d[i] = 0;
-      } else {
-        // Reduce out_padding as possible.
-        out_padding.d[i] -= padding.d[i];
-        padding.d[i] = 0;
-        hasOutputPadding = true;
-      }
-    }
+    add_output_padding(padding, out_padding, hasOutputPadding);
 
     // shape of deconvolution's weight: [in, out/groups, ...]
     // If there is still output padding, remove the bias. Bias will be added below.
@@ -213,51 +282,8 @@ bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args)
 #endif
     if (hasOutputPadding) {
       LOG_DEBUG("Padding output deconvolution tensor with:" << out_padding);
-
-      // Add padding layer
-      nvinfer1::ITensor* start;
-      nvinfer1::ITensor* totalPadding;
-      auto in_nbDims = orig_dims.nbDims;
-      std::vector<int32_t> startVec(in_nbDims, 0);
-      std::vector<int32_t> totalPaddingVec(in_nbDims, 0);
-      int32_t diff = in_nbDims - out_padding.nbDims;
-      for (int32_t i = diff; i < in_nbDims; i++) {
-        int32_t idx = i - diff;
-        startVec[i] = 0; // Don't need begin padding, only post padding
-        totalPaddingVec[i] = out_padding.d[idx];
-      }
-      start = tensor_to_const(ctx, torch::tensor(startVec, torch::kInt32));
-      totalPadding = tensor_to_const(ctx, torch::tensor(totalPaddingVec, torch::kInt32));
-
       nvinfer1::ITensor* tensorPtr = deconv->getOutput(0);
-      nvinfer1::ITensor* deconvOutShape = ctx->net->addShape(*tensorPtr)->getOutput(0);
-      const auto size =
-          ctx->net->addElementWise(*deconvOutShape, *totalPadding, nvinfer1::ElementWiseOperation::kSUM)->getOutput(0);
-
-      nvinfer1::Dims stride;
-      stride.nbDims = in_nbDims;
-      for (int64_t i = 0; i < in_nbDims; i++) {
-        stride.d[i] = 1;
-      }
-      const auto& dummy = stride;
-      auto* sliceLayer = ctx->net->addSlice(*tensorPtr, dummy, dummy, stride);
-      sliceLayer->setInput(1, *start);
-      sliceLayer->setInput(2, *size);
-      sliceLayer->setMode(nvinfer1::SliceMode::kFILL);
-      tensorPtr = sliceLayer->getOutput(0);
-
-      nvinfer1::Dims constantDims;
-      constantDims.nbDims = in_nbDims;
-      for (int64_t i = 0; i < in_nbDims; i++) {
-        constantDims.d[i] = 1;
-      }
-      constantDims.d[diff - 1] =
-          bias.shape.d[0]; // Set C dimension to bias dim and other dimensions to 1 to enable broadcast
-      auto const_layer = ctx->net->addConstant(constantDims, bias.data);
-      auto add_bias_layer =
-          ctx->net->addElementWise(*tensorPtr, *const_layer->getOutput(0), nvinfer1::ElementWiseOperation::kSUM);
-
-      new_layer = add_bias_layer;
+      new_layer = add_bias_layer(ctx, tensorPtr, orig_dims, out_padding, bias);
     } else {
       new_layer = deconv;
     }

core/conversion/converters/impl/matrix_multiply.cpp

@@ -16,20 +16,49 @@ auto mm_registrations TORCHTRT_UNUSED =
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                auto self = args[0].ITensorOrFreeze(ctx);
                auto other = args[1].ITensorOrFreeze(ctx);
+
+               auto selfDims = self->getDimensions().nbDims;
+               auto otherDims = other->getDimensions().nbDims;
+
+               bool squeezeFront = false;
+               bool squeezeBack = false;
+
+               if (selfDims == 1 && selfDims < otherDims) {
+                 squeezeFront = true;
+               } else if (otherDims == 1 && otherDims < selfDims) {
+                 // Append a 1 to the end of the shape before padding front to match self
+                 other = addPadding(ctx, n, other, 2, true, false);
+                 otherDims = other->getDimensions().nbDims;
+                 squeezeBack = true;
+               }
+
                // Ensure self and other tensors have same nbDims by expanding the dimensions (from 0 axis) if
                // necessary.
-               if (self->getDimensions().nbDims < other->getDimensions().nbDims) {
-                 self = addPadding(ctx, n, self, other->getDimensions().nbDims, false, false);
-               } else {
-                 other = addPadding(ctx, n, other, self->getDimensions().nbDims, false, false);
+               if (selfDims < otherDims) {
+                 self = addPadding(ctx, n, self, otherDims, false, false);
+               } else if (otherDims < selfDims) {
+                 other = addPadding(ctx, n, other, selfDims, false, false);
                }
 
                auto mm_layer = ctx->net->addMatrixMultiply(
                    *self, nvinfer1::MatrixOperation::kNONE, *other, nvinfer1::MatrixOperation::kNONE);
 
                TORCHTRT_CHECK(mm_layer, "Unable to create matrix multiplication node: " << *n);
                mm_layer->setName(util::node_info(n).c_str());
-               auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], mm_layer->getOutput(0));
+               auto out = mm_layer->getOutput(0);
+
+               if (squeezeFront || squeezeBack) {
+                 auto squeezeDimOffset = squeezeFront ? 2 : 1;
+                 auto reshapeDims =
+                     util::squeezeDims(out->getDimensions(), out->getDimensions().nbDims - squeezeDimOffset);
+                 auto shuffle_layer = ctx->net->addShuffle(*out);
+                 LOG_DEBUG("Squeezing matmul output for 1d correction: " << reshapeDims);
+                 TORCHTRT_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
+                 shuffle_layer->setReshapeDimensions(reshapeDims);
+                 shuffle_layer->setName((util::node_info(n) + "_squeeze").c_str());
+                 out = shuffle_layer->getOutput(0);
+               }
+               auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out);
 
                LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
                return true;

core/runtime/register_jit_hooks.cpp

@@ -87,7 +87,7 @@ static auto TORCHTRT_UNUSED TRTEngineTSRegistrtion =
         .def_pickle(
             [](const c10::intrusive_ptr<TRTEngine>& self) -> std::vector<std::string> {
               // Serialize TensorRT engine
-              auto serialized_trt_engine = self->cuda_engine->serialize();
+              auto serialized_trt_engine = make_trt(self->cuda_engine->serialize());
 
               // Adding device info related meta data to the serialized file
               auto trt_engine = std::string((const char*)serialized_trt_engine->data(), serialized_trt_engine->size());

docs/_cpp_api/classtorch__tensorrt_1_1DataType.html

@@ -10,7 +10,7 @@
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
 
-  <title>Class DataType &mdash; Torch-TensorRT v2.2.0.dev0+87fc615 documentation</title>
+  <title>Class DataType &mdash; Torch-TensorRT v2.2.0.dev0+d741d2c documentation</title>
 
 
 
@@ -237,7 +237,7 @@
 
 
                 <div class="version">
-                  v2.2.0.dev0+87fc615
+                  v2.2.0.dev0+d741d2c
                 </div>
 
 

docs/_cpp_api/classtorch__tensorrt_1_1Device_1_1DeviceType.html

@@ -10,7 +10,7 @@
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
 
-  <title>Class Device::DeviceType &mdash; Torch-TensorRT v2.2.0.dev0+87fc615 documentation</title>
+  <title>Class Device::DeviceType &mdash; Torch-TensorRT v2.2.0.dev0+d741d2c documentation</title>
 
 
 
@@ -237,7 +237,7 @@
 
 
                 <div class="version">
-                  v2.2.0.dev0+87fc615
+                  v2.2.0.dev0+d741d2c
                 </div>
 
 

docs/_cpp_api/classtorch__tensorrt_1_1TensorFormat.html

@@ -10,7 +10,7 @@
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
 
-  <title>Class TensorFormat &mdash; Torch-TensorRT v2.2.0.dev0+87fc615 documentation</title>
+  <title>Class TensorFormat &mdash; Torch-TensorRT v2.2.0.dev0+d741d2c documentation</title>
 
 
 
@@ -237,7 +237,7 @@
 
 
                 <div class="version">
-                  v2.2.0.dev0+87fc615
+                  v2.2.0.dev0+d741d2c
                 </div>
 
 

docs/_cpp_api/classtorch__tensorrt_1_1ptq_1_1Int8CacheCalibrator.html

@@ -10,7 +10,7 @@
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
 
-  <title>Template Class Int8CacheCalibrator &mdash; Torch-TensorRT v2.2.0.dev0+87fc615 documentation</title>
+  <title>Template Class Int8CacheCalibrator &mdash; Torch-TensorRT v2.2.0.dev0+d741d2c documentation</title>
 
 
 
@@ -237,7 +237,7 @@
 
 
                 <div class="version">
-                  v2.2.0.dev0+87fc615
+                  v2.2.0.dev0+d741d2c
                 </div>
 
 

docs/_cpp_api/classtorch__tensorrt_1_1ptq_1_1Int8Calibrator.html

@@ -10,7 +10,7 @@
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
 
-  <title>Template Class Int8Calibrator &mdash; Torch-TensorRT v2.2.0.dev0+87fc615 documentation</title>
+  <title>Template Class Int8Calibrator &mdash; Torch-TensorRT v2.2.0.dev0+d741d2c documentation</title>
 
 
 
@@ -237,7 +237,7 @@
 
 
                 <div class="version">
-                  v2.2.0.dev0+87fc615
+                  v2.2.0.dev0+d741d2c
                 </div>