Merge branch 'master' of https://github.com/NVIDIA/TRTorch into bowa_primif

Author: bowang007

.github/scripts/run_cpp_linter.py

@@ -24,9 +24,12 @@
     comment = '''There are some changes that do not conform to C++ style guidelines:\n ```diff\n{}```'''.format(output.stdout.decode("utf-8"))
     approval = 'REQUEST_CHANGES'
 
-pr.create_review(commit, comment, approval)
-
+try:
+    pr.create_review(commit, comment, approval)
+except:
+    print("Unable to submit in depth review, please review logs for linting issues")
+    
 if output.returncode != 0:
     exit(1)
 else:
-    exit(0)
+    exit(0)

.github/scripts/run_py_linter.py

@@ -24,9 +24,12 @@
     comment = '''There are some changes that do not conform to Python style guidelines:\n ```diff\n{}```'''.format(output.stdout.decode("utf-8"))
     approval = 'REQUEST_CHANGES'
 
-pr.create_review(commit, comment, approval)
+try:
+    pr.create_review(commit, comment, approval)
+except:
+    print("Unable to submit in depth review, please review logs for linting issues")
 
 if output.returncode != 0:
     exit(1)
 else:
-    exit(0)
+    exit(0)

BUILD

@@ -15,10 +15,10 @@ pkg_tar(
         "//core/conversion:include",
         "//core/conversion/conversionctx:include",
         "//core/conversion/converters:include",
-        "//core/conversion/converters/impl/plugins:include",
-        "//core/conversion/evaluators:include",
-        "//core/conversion/tensorcontainer:include",
         "//core/conversion/var:include",
+        "//core/conversion/tensorcontainer:include",
+        "//core/conversion/evaluators:include",
+        "//core/plugins:include",
         "//core/lowering:include",
         "//core/lowering/passes:include",
         "//core/runtime:include",
@@ -42,6 +42,7 @@ pkg_tar(
         "//conditions:default": [
             "//cpp/api/lib:libtrtorch.so",
             "//cpp/api/lib:libtrtorchrt.so",
+            "//cpp/api/lib:libtrtorch_plugins.so",
         ],
     }),
     mode = "0755",

core/conversion/conversionctx/ConversionCtx.cpp

@@ -1,9 +1,8 @@
+#include "core/conversion/conversionctx/ConversionCtx.h"
 #include <iostream>
 #include <sstream>
 #include <utility>
 
-#include "core/conversion/conversionctx/ConversionCtx.h"
-
 namespace trtorch {
 namespace core {
 namespace conversion {

core/conversion/converters/BUILD

@@ -36,12 +36,15 @@ cc_library(
         "impl/concat.cpp",
         "impl/constant.cpp",
         "impl/conv_deconv.cpp",
+        "impl/cumsum.cpp",
         "impl/element_wise.cpp",
         "impl/expand.cpp",
         "impl/interpolate.cpp",
+        "impl/layer_norm.cpp",
         "impl/linear.cpp",
         "impl/lstm_cell.cpp",
         "impl/matrix_multiply.cpp",
+        "impl/normalize.cpp",
         "impl/pooling.cpp",
         "impl/reduce.cpp",
         "impl/replication_pad.cpp",
@@ -64,7 +67,7 @@ cc_library(
         "//core/conversion/var",
         "//core/conversion/tensorcontainer",
         "//core/conversion/conversionctx",
-        "//core/conversion/converters/impl/plugins",
+        "//core/plugins:trtorch_plugins",
     ] + select({
         ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
         "//conditions:default": ["@libtorch//:libtorch"],

core/conversion/converters/converter_util.cpp

@@ -1,6 +1,7 @@
 #include "core/conversion/converters/converter_util.h"
 #include "core/conversion/converters/converters.h"
 #include "core/util/prelude.h"
+#include "torch/torch.h"
 
 namespace trtorch {
 namespace core {
@@ -59,6 +60,68 @@ nvinfer1::ITensor* addUnpadding(
   }
 }
 
+nvinfer1::ILayer* add_elementwise(
+    ConversionCtx* ctx,
+    nvinfer1::ElementWiseOperation op,
+    nvinfer1::ITensor* self,
+    nvinfer1::ITensor* other,
+    const std::string& name) {
+  // ensure self to have larger number of dimension
+  bool swapSelfOther = false;
+  if (self->getDimensions().nbDims < other->getDimensions().nbDims) {
+    std::swap(self, other);
+    swapSelfOther = true;
+  }
+  auto selfDim = util::toVec(self->getDimensions());
+  auto otherDim = util::toVec(other->getDimensions());
+  if (selfDim.size() != otherDim.size()) {
+    // other is with dynamic shape, need to expand its dimension now and get its
+    // shape at runtime
+    if (otherDim.end() != std::find(otherDim.begin(), otherDim.end(), -1)) {
+      auto thOtherStaticShapeMask = torch::ones(selfDim.size(), torch::kInt32);
+      auto thOtherDynamicShapeMask = torch::zeros(selfDim.size(), torch::kInt32);
+      for (size_t start = selfDim.size() - otherDim.size(), idx = 0; idx < otherDim.size(); ++idx) {
+        if (-1 != otherDim[idx]) {
+          thOtherStaticShapeMask[start + idx] = otherDim[idx];
+        } else {
+          thOtherStaticShapeMask[start + idx] = 0;
+          thOtherDynamicShapeMask[start + idx] = 1;
+        }
+      }
+      auto otherStaticShapeMask = tensor_to_const(ctx, thOtherStaticShapeMask);
+      auto otherDynamicShapeMask = tensor_to_const(ctx, thOtherDynamicShapeMask);
+      auto selfShape = ctx->net->addShape(*self)->getOutput(0);
+      // size of dynamic dimension of other need to the same as that of
+      // corresponding dimension of self
+      auto otherDynamicShape =
+          ctx->net->addElementWise(*selfShape, *otherDynamicShapeMask, nvinfer1::ElementWiseOperation::kPROD)
+              ->getOutput(0);
+      auto targetOtherShape =
+          ctx->net->addElementWise(*otherDynamicShape, *otherStaticShapeMask, nvinfer1::ElementWiseOperation::kSUM)
+              ->getOutput(0);
+
+      auto otherShuffle = ctx->net->addShuffle(*other);
+      otherShuffle->setName(std::string("Reshape other tensor to have the same nDim as self for " + name).c_str());
+      otherShuffle->setInput(1, *targetOtherShape);
+      other = otherShuffle->getOutput(0);
+    } else {
+      // other is with static shape, expand dimension to make tow tensor have
+      // the same number of dimension
+      auto otherShuffle = ctx->net->addShuffle(*other);
+      otherShuffle->setReshapeDimensions(util::toDimsPad(otherDim, selfDim.size()));
+      other = otherShuffle->getOutput(0);
+    }
+  }
+  if (swapSelfOther) {
+    // swap back
+    std::swap(self, other);
+    swapSelfOther = false;
+  }
+  auto ele = ctx->net->addElementWise(*self, *other, op);
+  ele->setName(name.c_str());
+  return ele;
+}
+
 } // namespace converters
 } // namespace conversion
 } // namespace core

core/conversion/converters/converter_util.h

@@ -35,6 +35,13 @@ nvinfer1::ITensor* addUnpadding(
     bool trailing = true,
     bool use_zeros = true);
 
+nvinfer1::ILayer* add_elementwise(
+    ConversionCtx* ctx,
+    nvinfer1::ElementWiseOperation op,
+    nvinfer1::ITensor* self,
+    nvinfer1::ITensor* other,
+    const std::string& name);
+
 } // namespace converters
 } // namespace conversion
 } // namespace core

core/conversion/converters/impl/activation.cpp

@@ -146,7 +146,6 @@ auto acthardtanh TRTORCH_UNUSED =
 
                     auto new_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kLEAKY_RELU);
                     new_layer->setAlpha(negative_slopeScalar);
-
                     new_layer->setName(util::node_info(n).c_str());
                     auto out_tensor = new_layer->getOutput(0);
                     out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
@@ -167,6 +166,35 @@ auto acthardtanh TRTORCH_UNUSED =
                     auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
                     LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
                     return true;
+                  }})
+        .pattern({"aten::gelu(Tensor self) -> (Tensor)",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    auto in = args[0].ITensorOrFreeze(ctx);
+                    nvinfer1::DataType type = in->getType();
+                    TRTORCH_CHECK(
+                        type == nvinfer1::DataType::kFLOAT || type == nvinfer1::DataType::kHALF,
+                        "gelu only supports kFLOAT and kHALF");
+                    std::string pluginName = "CustomGeluPluginDynamic";
+                    nvinfer1::PluginFieldCollection fc;
+                    std::vector<nvinfer1::PluginField> f;
+                    int type_id = ctx->settings.op_precision == nvinfer1::DataType::kFLOAT
+                        ? 0
+                        : 1; // Integer encoding the DataType (0: FP32, 1: FP16)
+                    f.emplace_back(nvinfer1::PluginField("type_id", &type_id, nvinfer1::PluginFieldType::kINT32, 1));
+                    fc.nbFields = f.size();
+                    fc.fields = f.data();
+
+                    auto creator = getPluginRegistry()->getPluginCreator("CustomGeluPluginDynamic", "1", "");
+                    auto gelu_plugin = creator->createPlugin("gelu", &fc);
+
+                    TRTORCH_CHECK(gelu_plugin, "Unable to create gelu plugin from TensorRT plugin registry" << *n);
+                    auto new_layer =
+                        ctx->net->addPluginV2(reinterpret_cast<nvinfer1::ITensor* const*>(&in), 1, *gelu_plugin);
+                    new_layer->setName("gelu");
+                    auto out_tensor = new_layer->getOutput(0);
+                    out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
+                    LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
+                    return true;
                   }});
 
 } // namespace

core/conversion/converters/impl/cumsum.cpp

@@ -0,0 +1,74 @@
+#include "NvInfer.h"
+#include "core/conversion/converters/converters.h"
+#include "core/conversion/tensorcontainer/TensorContainer.h"
+#include "core/util/prelude.h"
+#include "core/util/trt_util.h"
+#include "torch/torch.h"
+
+#include <ATen/ATen.h>
+#include <vector>
+
+namespace trtorch {
+namespace core {
+namespace conversion {
+namespace converters {
+namespace impl {
+namespace {
+
+auto cumsum_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().pattern(
+    {"aten::cumsum(Tensor self, int dim, *, int? dtype=None) -> (Tensor)",
+     [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+       auto in = args[0].ITensorOrFreeze(ctx);
+       auto input_dims = in->getDimensions();
+       int dim = args[1].unwrapToInt();
+       TRTORCH_CHECK(
+           (dim >= 0 && dim < input_dims.nbDims) || (dim < 0 && (input_dims.nbDims + dim >= 0)),
+           "Dimension out of range (expected to be in range of [" << -input_dims.nbDims << ", " << input_dims.nbDims - 1
+                                                                  << "], but got " << dim << ")");
+       if (dim < 0) {
+         dim += input_dims.nbDims;
+       }
+
+       // Scan through each slice across summation axis and add it to the running sum
+       auto loop = ctx->net->addLoop();
+       nvinfer1::ITensor* tripLimit = NULL;
+       if (input_dims.d[dim] > 0) {
+         torch::Tensor axis = torch::tensor(input_dims.d[dim], torch::kInt32);
+         tripLimit = tensor_to_const(ctx, axis);
+       } else {
+         nvinfer1::ITensor* inpShape = ctx->net->addShape(*in)->getOutput(0);
+         torch::Tensor dimValue = torch::tensor(dim, torch::kInt32);
+         nvinfer1::ITensor* axis = tensor_to_const(ctx, dimValue);
+         tripLimit = ctx->net->addGather(*inpShape, *axis, 0)->getOutput(0);
+       }
+
+       loop->addTripLimit(*tripLimit, nvinfer1::TripLimit::kCOUNT);
+
+       auto iterator = loop->addIterator(*in, dim, false);
+       auto data = iterator->getOutput(0);
+       auto newDims = data->getDimensions();
+
+       torch::Tensor zeroValue = at::full(util::toVec(newDims), 0, torch::kFloat32);
+       auto zeroTensor = tensor_to_const(ctx, zeroValue);
+       auto runningSum = loop->addRecurrence(*zeroTensor);
+       auto runningSumTensor = runningSum->getOutput(0);
+
+       auto curSum = ctx->net->addElementWise(*data, *runningSumTensor, nvinfer1::ElementWiseOperation::kSUM);
+       runningSum->setInput(1, *curSum->getOutput(0));
+
+       nvinfer1::ILoopOutputLayer* loopOut =
+           loop->addLoopOutput(*curSum->getOutput(0), nvinfer1::LoopOutput::kCONCATENATE, dim);
+       loopOut->setInput(1, *tripLimit);
+
+       auto layer_output = ctx->AssociateValueAndTensor(n->outputs()[0], loopOut->getOutput(0));
+
+       LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions());
+       return true;
+     }});
+
+} // namespace
+} // namespace impl
+} // namespace converters
+} // namespace conversion
+} // namespace core
+} // namespace trtorch

core/conversion/converters/impl/element_wise.cpp

@@ -1,4 +1,5 @@
 #include <torch/torch.h>
+#include "core/conversion/converters/converter_util.h"
 #include "core/conversion/converters/converters.h"
 #include "core/util/prelude.h"
 
@@ -9,65 +10,6 @@ namespace converters {
 namespace impl {
 namespace {
 
-nvinfer1::ILayer* add_elementwise(
-    ConversionCtx* ctx,
-    nvinfer1::ElementWiseOperation op,
-    nvinfer1::ITensor* self,
-    nvinfer1::ITensor* other,
-    const std::string& name) {
-  // ensure self to have larger number of dimension
-  bool swapSelfOther = false;
-  if (self->getDimensions().nbDims < other->getDimensions().nbDims) {
-    std::swap(self, other);
-    swapSelfOther = true;
-  }
-  auto selfDim = util::toVec(self->getDimensions());
-  auto otherDim = util::toVec(other->getDimensions());
-  if (selfDim.size() != otherDim.size()) {
-    // other is with dynamic shape, need to expand its dimension now and get its shape at runtime
-    if (otherDim.end() != std::find(otherDim.begin(), otherDim.end(), -1)) {
-      auto thOtherStaticShapeMask = torch::ones(selfDim.size(), torch::kInt32);
-      auto thOtherDynamicShapeMask = torch::zeros(selfDim.size(), torch::kInt32);
-      for (size_t start = selfDim.size() - otherDim.size(), idx = 0; idx < otherDim.size(); ++idx) {
-        if (-1 != otherDim[idx]) {
-          thOtherStaticShapeMask[start + idx] = otherDim[idx];
-        } else {
-          thOtherStaticShapeMask[start + idx] = 0;
-          thOtherDynamicShapeMask[start + idx] = 1;
-        }
-      }
-      auto otherStaticShapeMask = tensor_to_const(ctx, thOtherStaticShapeMask);
-      auto otherDynamicShapeMask = tensor_to_const(ctx, thOtherDynamicShapeMask);
-      auto selfShape = ctx->net->addShape(*self)->getOutput(0);
-      // size of dynamic dimension of other need to the same as that of corresponding dimension of self
-      auto otherDynamicShape =
-          ctx->net->addElementWise(*selfShape, *otherDynamicShapeMask, nvinfer1::ElementWiseOperation::kPROD)
-              ->getOutput(0);
-      auto targetOtherShape =
-          ctx->net->addElementWise(*otherDynamicShape, *otherStaticShapeMask, nvinfer1::ElementWiseOperation::kSUM)
-              ->getOutput(0);
-
-      auto otherShuffle = ctx->net->addShuffle(*other);
-      otherShuffle->setName(std::string("Reshape other tensor to have the same nDim as self for " + name).c_str());
-      otherShuffle->setInput(1, *targetOtherShape);
-      other = otherShuffle->getOutput(0);
-    } else {
-      // other is with static shape, expand dimension to make tow tensor have the same number of dimension
-      auto otherShuffle = ctx->net->addShuffle(*other);
-      otherShuffle->setReshapeDimensions(util::toDimsPad(otherDim, selfDim.size()));
-      other = otherShuffle->getOutput(0);
-    }
-  }
-  if (swapSelfOther) {
-    // swap back
-    std::swap(self, other);
-    swapSelfOther = false;
-  }
-  auto ele = ctx->net->addElementWise(*self, *other, op);
-  ele->setName(name.c_str());
-  return ele;
-}
-
 nvinfer1::ITensor* clamp_util(
     ConversionCtx* ctx,
     const torch::jit::Node* n,