Merge pull request #943 from p517332051/netease_fix_stack_bug

fix bug stack when dim -1

Author: narendasan

core/conversion/converters/impl/stack.cpp

@@ -19,9 +19,16 @@ auto stack_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().patt
      [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
        auto in = args[0].IValue()->toListRef();
        auto dim = args[1].unwrapToInt();
+       if (-1 == dim) {
+         auto first_in = in[0];
+         if (first_in.isTensor()) {
+           dim = first_in.toTensor().ndimension();
+         } else {
+           dim = first_in.toCustomClass<TensorContainer>()->tensor()->getDimensions().nbDims;
+         }
+       }
 
        std::vector<nvinfer1::ITensor*> tensors;
-
        for (auto t : in) {
          nvinfer1::ITensor* itensor;
 

tests/core/conversion/converters/test_stack.cpp

@@ -5,51 +5,72 @@
 #include "torch/csrc/jit/ir/irparser.h"
 
 TEST(Converters, ATenStackPureTensorConvertsCorrectly) {
+  auto TestATenStackPureTensorConvertsCorrectly = [](const std::string& graph) {
+    auto g = std::make_shared<torch::jit::Graph>();
+    torch::jit::parseIR(graph, g.get());
+
+    auto in1 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+    auto in2 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+
+    auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+    auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, in2});
+
+    params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+    auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, in2});
+
+    ASSERT_TRUE(
+        torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+  };
   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";
+  const auto graph2 = R"IR(
+      graph(%0 : Tensor,
+            %1 : Tensor):
+        %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.get());
+  TestATenStackPureTensorConvertsCorrectly(graph);
+  TestATenStackPureTensorConvertsCorrectly(graph2);
+}
 
-  auto in1 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
-  auto in2 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+TEST(Converters, ATenStackDiffTensorConvertsCorrectly) {
+  auto TestATenStackDiffTensorConvertsCorrectly = [](const std::string& graph) {
+    auto g = std::make_shared<torch::jit::Graph>();
+    torch::jit::parseIR(graph, g.get());
 
-  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
-  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, in2});
+    auto in1 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+    auto in2 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
 
-  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
-  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, in2});
+    auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {in2});
+    auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1});
 
-  ASSERT_TRUE(
-      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
-}
+    params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {in2});
+    auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1});
 
-TEST(Converters, ATenStackDiffTensorConvertsCorrectly) {
+    ASSERT_TRUE(
+        torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+  };
   const auto graph = R"IR(
       graph(%0 : Tensor,
             %1 : Float(4, 4, 4, strides=[16, 4, 1])):
         %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.get());
-
-  auto in1 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
-  auto in2 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
-
-  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {in2});
-  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1});
-
-  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {in2});
-  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1});
-
-  ASSERT_TRUE(
-      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
-}
+  const auto graph2 = R"IR(
+      graph(%0 : Tensor,
+            %1 : Float(4, 4, 4, strides=[16, 4, 1])):
+        %2 : Tensor[] = prim::ListConstruct(%0, %1)
+        %3 : int = prim::Constant[value=-1]()
+        %4 : Tensor = aten::stack(%2, %3)
+        return (%4))IR";
+  TestATenStackDiffTensorConvertsCorrectly(graph);
+  TestATenStackDiffTensorConvertsCorrectly(graph2);
+}