Merge pull request #116 from abhi-iyer/master

Add support for aten::select.int and aten::stack

Author: narendasan

core/conversion/converters/BUILD

@@ -28,7 +28,9 @@ cc_library(
         "impl/shuffle.cpp",
         "impl/softmax.cpp",
         "impl/unary.cpp",
-        "impl/interpolate.cpp"
+        "impl/interpolate.cpp",
+        "impl/select.cpp",
+        "impl/stack.cpp"
     ],
     deps = [
         "@tensorrt//:nvinfer",

core/conversion/converters/impl/select.cpp

@@ -0,0 +1,58 @@
+#include "torch/torch.h"
+#include "core/util/prelude.h"
+#include "core/conversion/converters/converters.h"
+#include "NvInfer.h"
+
+#include <ATen/ATen.h>
+#include <vector>
+
+namespace trtorch {
+namespace core {
+namespace conversion {
+namespace converters {
+namespace impl {
+namespace {
+
+auto select_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
+    .pattern({
+        "aten::select.int(Tensor(a) self, int dim, int index) -> (Tensor(a))",
+        [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+            auto in = args[0].ITensor();
+            auto axis  = args[1].unwrapToInt();
+            auto ind = (int32_t) args[2].unwrapToInt();
+
+            // index to access needs to be an at::Tensor
+            at::Tensor indices = torch::tensor({ind}).to(torch::kI32);
+            auto weights = Weights(ctx, indices);
+
+            // IConstantLayer to convert indices from Weights to ITensor
+            auto const_layer = ctx->net->addConstant(weights.shape, weights.data);
+            TRTORCH_CHECK(const_layer, "Unable to create constant layer from node: " << *n);
+            auto const_out = const_layer->getOutput(0);
+            
+            // IGatherLayer takes in input tensor, the indices, and the axis of input tensor to take indices from
+            auto gather_layer = ctx->net->addGather(*in, *const_out, axis);
+            TRTORCH_CHECK(gather_layer, "Unable to create gather layer from node: " << *n);
+            auto gather_out = gather_layer->getOutput(0);
+
+            // IShuffleLayer removes redundant dimensions
+            auto shuffle_layer = ctx->net->addShuffle(*gather_out);
+            TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
+            shuffle_layer->setReshapeDimensions(util::unpadDims(gather_out->getDimensions()));
+            shuffle_layer->setName(util::node_info(n).c_str());
+            auto shuffle_out = shuffle_layer->getOutput(0);
+
+            auto out = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle_out);
+
+            LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+
+            return true;
+        }
+    });
+
+} // namespace
+} // namespace impl
+} // namespace converters
+} // namespace conversion
+} // namespace core
+} // namespace trtorch

core/conversion/converters/impl/stack.cpp

@@ -0,0 +1,64 @@
+#include "torch/torch.h"
+#include "core/util/prelude.h"
+#include "core/conversion/converters/converters.h"
+#include "core/conversion/tensorcontainer/TensorContainer.h"
+#include "NvInfer.h"
+
+#include <ATen/ATen.h>
+#include <vector>
+
+namespace trtorch {
+namespace core {
+namespace conversion {
+namespace converters {
+namespace impl {
+namespace {
+
+auto stack_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
+    .pattern({
+        "aten::stack(Tensor[] tensors, int dim=0) -> (Tensor)",
+        [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+            auto in = args[0].IValue()->toListRef();
+            auto dim = args[1].unwrapToInt();
+
+            std::vector<nvinfer1::ITensor*> tensors; 
+            
+            for (auto t : in) {
+                nvinfer1::ITensor* itensor;
+
+                if (t.isTensor()) {
+                    auto weight = Weights(ctx, t.toTensor());
+
+                    auto const_layer = ctx->net->addConstant(weight.shape, weight.data);
+                    TRTORCH_CHECK(const_layer, "Unable to create constant layer from node: " << *n);
+
+                    itensor = const_layer->getOutput(0);
+                } else {
+                    auto cont = t.toCustomClass<TensorContainer>();
+                    itensor = cont->tensor();
+                }
+
+                auto shuffle_layer = ctx->net->addShuffle(*itensor);
+                TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
+                shuffle_layer->setReshapeDimensions(util::unsqueezeDims(itensor->getDimensions(), dim));
+                
+                tensors.push_back(shuffle_layer->getOutput(0));
+            }
+
+            auto concat_layer = ctx->net->addConcatenation(tensors.data(), tensors.size());
+            TRTORCH_CHECK(concat_layer, "Unable to create concatenation layer from node: " << *n);
+            concat_layer->setAxis(static_cast<int>(dim));
+            auto out = ctx->AssociateValueAndTensor(n->outputs()[0], concat_layer->getOutput(0));
+
+            LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+
+            return true;
+        }
+    });
+
+} // namespace
+} // namespace impl
+} // namespace converters
+} // namespace conversion
+} // namespace core
+} // namespace trtorch

core/util/trt_util.cpp

@@ -82,6 +82,56 @@ nvinfer1::Dims toDimsPad(c10::List<int64_t> l, uint64_t pad_to) {
     return dims;
 }
 
+nvinfer1::Dims unpadDims(const nvinfer1::Dims& d) {
+    nvinfer1::Dims dims;
+
+    int j = 0;
+    bool pad_dims_done = false;
+
+    for (int i = 0; i < d.nbDims; i++) {
+        if (d.d[i] == 1 && !pad_dims_done) {
+            // skip over unecessary dimension
+            continue;
+        } else {
+            dims.d[j] = d.d[i];
+            j++;
+
+            // keep all other dimensions (don't skip over them)
+            pad_dims_done = true;
+        }
+    }
+
+    dims.nbDims = j;
+
+    return dims;
+}
+
+nvinfer1::Dims unsqueezeDims(const nvinfer1::Dims& d, int pos) {
+    // acceptable range for pos is [0, d.nbDims]
+    TRTORCH_ASSERT(pos >= 0 && pos <= d.nbDims, "ERROR: Index to unsqueeze is out of bounds.");
+    
+    nvinfer1::Dims dims;
+
+    int i = 0;
+    int j = 0;
+
+    while (i <= d.nbDims) {
+        if (j != pos) {
+            dims.d[j] = d.d[i];
+            i++;
+        } else {
+            // add new dimension at pos
+            dims.d[j] = 1;
+        }
+
+        j++;
+    }
+
+    dims.nbDims = d.nbDims+1;
+
+    return dims;
+}
+
 std::vector<int64_t> toVec(nvinfer1::Dims d) {
     std::vector<int64_t> dims;
     for (int i = 0; i < d.nbDims; i++) {

core/util/trt_util.h

@@ -79,6 +79,8 @@ int64_t volume(const nvinfer1::Dims& d);
 
 nvinfer1::Dims toDimsPad(c10::IntArrayRef l, uint64_t pad_to);
 nvinfer1::Dims toDimsPad(c10::List<int64_t> l, uint64_t pad_to);
+nvinfer1::Dims unpadDims(const nvinfer1::Dims& d);
+nvinfer1::Dims unsqueezeDims(const nvinfer1::Dims& d, int pos);
 nvinfer1::Dims toDims(c10::IntArrayRef l);
 nvinfer1::Dims toDims(c10::List<int64_t> l);
 nvinfer1::DimsHW toDimsHW(c10::List<int64_t> l);

tests/core/converters/BUILD

@@ -59,6 +59,14 @@ converter_test(
   name = "test_interpolate"
 )
 
+converter_test(
+  name = "test_select"
+)
+
+converter_test(
+  name = "test_stack"
+)
+
 test_suite(
   name = "test_converters",
   tests = [
@@ -74,6 +82,8 @@ test_suite(
     ":test_softmax",
     ":test_unary",
     ":test_interpolate",
+    ":test_select",
+    ":test_stack"
   ]
 )
 

tests/core/converters/test_select.cpp

@@ -0,0 +1,59 @@
+#include <string>
+#include "gtest/gtest.h"
+#include "torch/csrc/jit/ir/irparser.h"
+#include "tests/util/util.h"
+#include "core/compiler.h"
+
+TEST(Converters, ATenSelectIntConvertsCorrectly) {
+    const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=0]()
+        %3 : Tensor = aten::select(%0, %2, %2)
+        return (%3))IR";
+    
+    auto g = std::make_shared<torch::jit::Graph>();
+
+    torch::jit::parseIR(graph, &*g);
+
+    auto in = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+
+    auto jit_in = at::clone(in);
+    auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto jit_results = trtorch::tests::util::RunGraph(g, params, {jit_in});
+
+    auto trt_in = at::clone(in);
+    params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {trt_in});
+
+    auto trt = trt_results[0].reshape(jit_results[0].sizes());
+
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}
+
+TEST(Converters, ATenSelectIntTwiceConvertsCorrectly) {
+    const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=0]()
+        %3 : int = prim::Constant[value=3]()
+        %4 : Tensor = aten::select(%0, %2, %2)
+        %5 : Tensor = aten::select(%4, %2, %3)
+        return (%5))IR";
+    
+    auto g = std::make_shared<torch::jit::Graph>();
+
+    torch::jit::parseIR(graph, &*g);
+
+    auto in = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+
+    auto jit_in = at::clone(in);
+    auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto jit_results = trtorch::tests::util::RunGraph(g, params, {jit_in});
+
+    auto trt_in = at::clone(in);
+    params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {trt_in});
+
+    auto trt = trt_results[0].reshape(jit_results[0].sizes());
+
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}

tests/core/converters/test_stack.cpp

@@ -0,0 +1,53 @@
+#include <string>
+#include "gtest/gtest.h"
+#include "torch/csrc/jit/ir/irparser.h"
+#include "tests/util/util.h"
+#include "core/compiler.h"
+
+TEST(Converters, ATenStackPureTensorConvertsCorrectly) {
+    const auto graph = R"IR(
+      graph(%0 : Tensor,
+            %1 : Tensor):
+        %2 : Tensor[] = prim::ListConstruct(%0, %1)
+        %3 : int = prim::Constant[value=3]()
+        %4 : Tensor = aten::stack(%2, %3)
+        return (%4))IR";
+
+    auto g = std::make_shared<torch::jit::Graph>();
+    torch::jit::parseIR(graph, &*g);
+
+    auto in1 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+    auto in2 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+
+    auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto jit_results = trtorch::tests::util::RunGraph(g, params, {in1, in2});
+
+    params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in1, in2});
+
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}
+
+TEST(Converters, ATenStackDiffTensorConvertsCorrectly) {
+    const auto graph = R"IR(
+      graph(%0 : Tensor,
+            %1 : Float(4, 4, 4)):
+        %2 : Tensor[] = prim::ListConstruct(%0, %1)
+        %3 : int = prim::Constant[value=1]()
+        %4 : Tensor = aten::stack(%2, %3)
+        return (%4))IR";
+
+    auto g = std::make_shared<torch::jit::Graph>();
+    torch::jit::parseIR(graph, &*g);
+
+    auto in1 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+    auto in2 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+
+    auto params = trtorch::core::conversion::get_named_params(g->inputs(), {in2});
+    auto jit_results = trtorch::tests::util::RunGraph(g, params, {in1});
+
+    params = trtorch::core::conversion::get_named_params(g->inputs(), {in2});
+    auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in1});
+
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}