Merge pull request #986 from NVIDIA/unbind

feat (core//conversion) : Add converter for torch.unbind

Author: narendasan

core/conversion/converters/impl/select.cpp

@@ -15,26 +15,32 @@ namespace converters {
 namespace impl {
 namespace {
 
-bool add_split(ConversionCtx* ctx, const torch::jit::Node* n, args& args, bool split_list) {
+bool add_split(ConversionCtx* ctx, const torch::jit::Node* n, args& args, bool split_list, bool unbind) {
   auto in = args[0].ITensor();
-  auto axis = args[2].unwrapToInt();
-  auto inDimSize = in->getDimensions().d[axis];
-  auto numOutputs = 1, numRemainder = 0;
+  auto numOutputs = 1, numRemainder = 0, axis = 0;
   std::vector<int64_t> sizes;
 
-  if (split_list) {
-    sizes = args[1].unwrapToIntList().vec();
-    numOutputs = sizes.size();
+  if (unbind) {
+    axis = args[1].unwrapToInt();
+    numOutputs = in->getDimensions().d[axis];
+    sizes.insert(sizes.end(), numOutputs, 1);
   } else {
-    auto split_size = args[1].unwrapToInt();
-    numOutputs = inDimSize / split_size;
-    numRemainder = inDimSize % split_size;
-    for (int64_t i = 0; i < numOutputs; i++) {
-      sizes.push_back(split_size);
-    }
-    if (numRemainder) {
-      numOutputs += 1;
-      sizes.push_back(numRemainder);
+    axis = args[2].unwrapToInt();
+    auto inDimSize = in->getDimensions().d[axis];
+    if (split_list) {
+      sizes = args[1].unwrapToIntList().vec();
+      numOutputs = sizes.size();
+    } else {
+      auto split_size = args[1].unwrapToInt();
+      numOutputs = inDimSize / split_size;
+      numRemainder = inDimSize % split_size;
+      for (int64_t i = 0; i < numOutputs; i++) {
+        sizes.push_back(split_size);
+      }
+      if (numRemainder) {
+        numOutputs += 1;
+        sizes.push_back(numRemainder);
+      }
     }
   }
 
@@ -340,19 +346,25 @@ auto select_registrations TORCHTRT_UNUSED =
              }})
         .pattern({"aten::split(Tensor self, int[] split_sizes, int dim=0) -> (Tensor[])",
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
-                    add_split(ctx, n, args, true);
+                    add_split(ctx, n, args, true, false);
                     LOG_DEBUG("Converted split op into a list of IValues");
                     return true;
                   }})
         .pattern({"aten::split.Tensor(Tensor(a) self, int split_size, int dim=0) -> (Tensor[])",
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
-                    add_split(ctx, n, args, false);
+                    add_split(ctx, n, args, false, false);
                     LOG_DEBUG("Converted split op into a list of IValues");
                     return true;
                   }})
         .pattern({"aten::split_with_sizes(Tensor(a) self, int[] split_sizes, int dim=0) -> (Tensor[])",
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
-                    add_split(ctx, n, args, true);
+                    add_split(ctx, n, args, true, false);
+                    LOG_DEBUG("Converted split op into a list of IValues");
+                    return true;
+                  }})
+        .pattern({"aten::unbind.int(Tensor(a -> *) self, int dim=0) -> (Tensor[])",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    add_split(ctx, n, args, false, true);
                     LOG_DEBUG("Converted split op into a list of IValues");
                     return true;
                   }})

tests/core/conversion/converters/test_select.cpp

@@ -567,3 +567,30 @@ TEST(Converters, ATenIndexTensorConvertsCorrectly) {
   ASSERT_TRUE(
       torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
 }
+
+TEST(Converters, ATenUnbindConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%x.1 : Tensor):
+        %2 : int = prim::Constant[value=0]()
+        %3 : Tensor[] = aten::unbind(%x.1, %2)
+        %o1.1 : Tensor, %o2.1 : Tensor = prim::ListUnpack(%3)
+        return (%o1.1, %o2.1))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, g.get());
+
+  auto in = at::randint(1, 10, {2, 3, 4, 4}, {at::kCUDA});
+
+  auto jit_in = at::clone(in);
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in});
+
+  auto trt_in = at::clone(in);
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in});
+
+  for (size_t i = 0; i < jit_results.size(); i++) {
+    auto trt = trt_results[i].reshape(jit_results[i].sizes());
+    ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6));
+  }
+}