replace cumsum_plugin with native TRT

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

Author: ruoqianguo

core/conversion/converters/impl/cumsum.cpp

@@ -3,7 +3,6 @@
 #include "core/conversion/tensorcontainer/TensorContainer.h"
 #include "core/util/prelude.h"
 #include "core/util/trt_util.h"
-#include "plugins/cumsum_plugin.h"
 #include "torch/torch.h"
 
 #include <ATen/ATen.h>
@@ -16,26 +15,10 @@ namespace converters {
 namespace impl {
 namespace {
 
-void create_plugin(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* in, const char* name, int dim) {
-  LOG_WARNING("Cumsum layer will be run through ATen, not TensorRT. Performance may be lower than expected");
-
-  auto creator = new plugins::CumsumPluginCreator();
-  auto plugin = creator->createPlugin(name, dim);
-
-  auto cumsum_layer = ctx->net->addPluginV2(reinterpret_cast<nvinfer1::ITensor* const*>(&in), 1, *plugin);
-  TRTORCH_CHECK(cumsum_layer, "Unable to create cumsum plugin from node" << *n);
-
-  cumsum_layer->setName(util::node_info(n).c_str());
-
-  auto layer_output = ctx->AssociateValueAndTensor(n->outputs()[0], cumsum_layer->getOutput(0));
-
-  LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions());
-}
-
 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].ITensor();
+       auto in = args[0].ITensorOrFreeze(ctx);
        auto input_dims = in->getDimensions();
        int dim = args[1].unwrapToInt();
        TRTORCH_CHECK(
@@ -45,7 +28,41 @@ auto cumsum_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().patt
        if (dim < 0) {
          dim += input_dims.nbDims;
        }
-       create_plugin(ctx, n, in, "Cumsum", dim);
+
+       // 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;
      }});
 

core/conversion/converters/impl/plugins/BUILD

@@ -10,12 +10,10 @@ config_setting(
 cc_library(
     name = "plugins",
     hdrs = [
-        "interpolate_plugin.h",
-        "cumsum_plugin.h"
+        "interpolate_plugin.h"
     ],
     srcs = [
-        "interpolate_plugin.cpp",
-        "cumsum_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", "cumsum_plugin.h"],
+    srcs = ["interpolate_plugin.h"],
 )