support aten::roll

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

Author: ruoqianguo

core/conversion/converters/impl/select.cpp

@@ -67,6 +67,34 @@ bool add_split(ConversionCtx* ctx, const torch::jit::Node* n, args& args, bool s
   return true;
 }
 
+nvinfer1::ITensor* roll(
+    ConversionCtx* ctx,
+    nvinfer1::ITensor* in,
+    int shift,
+    int dim,
+    const std::vector<int64_t>& in_shape) {
+  auto in_dim = in_shape[dim];
+
+  auto start = (in_dim - shift) % in_dim;
+  // Behavior of % is different in C++ vs Python for negative numbers. This
+  // corrects the difference.
+  if (start < 0) {
+    start = start + in_dim;
+  }
+  at::Tensor index0 = at::arange(start, in_dim, 1, torch::kInt32);
+  at::Tensor index;
+  if (start == 0) {
+    index = index0;
+  } else {
+    at::Tensor index1 = at::arange(start, torch::kInt32);
+    index = at::cat({index0, index1}, 0);
+  }
+  auto index_tensor = tensor_to_const(ctx, index);
+  auto gather_layer = ctx->net->addGather(*in, *index_tensor, dim);
+  auto out = gather_layer->getOutput(0);
+  return out;
+}
+
 auto select_registrations TORCHTRT_UNUSED =
     RegisterNodeConversionPatterns()
         .pattern({"aten::select.int(Tensor(a) self, int dim, int index) -> (Tensor(a))",
@@ -202,6 +230,29 @@ auto select_registrations TORCHTRT_UNUSED =
 
                return true;
              }})
+        .pattern({"aten::roll(Tensor self, int[1] shifts, int[1] dims=[]) -> (Tensor)",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    auto in = args[0].ITensor();
+                    auto shifts = args[1].unwrapToIntList().vec();
+                    auto dims = args[2].unwrapToIntList().vec();
+
+                    TORCHTRT_CHECK(dims.size() == shifts.size(), "dims.size() should be equal to shifts.size()");
+                    if (ctx->input_is_dynamic) {
+                      TORCHTRT_THROW_ERROR("aten::roll is currently not support in dynamic input shape compilation");
+                    } else {
+                      auto in_shape = util::toVec(in->getDimensions());
+                      for (size_t i = 0; i < dims.size(); i++) {
+                        auto dim = dims[i] < 0 ? (in_shape.size() + dims[i]) : dims[i];
+                        TORCHTRT_CHECK(dim < in_shape.size(), "Dimension out of range");
+                        in = roll(ctx, in, shifts[i], dim, in_shape);
+                      }
+                      auto out = ctx->AssociateValueAndTensor(n->outputs()[0], in);
+
+                      LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+
+                      return true;
+                    }
+                  }})
         .pattern(
             {"aten::slice.Tensor(Tensor(a) self, int dim=0, int? start=None, int? end=None, int step=1) -> Tensor(a)",
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {

tests/core/conversion/converters/test_select.cpp

@@ -195,6 +195,84 @@ TEST(Converters, ATenEmbeddingConvertsCorrectly) {
   ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6));
 }
 
+TEST(Converters, ATenRollConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%1 : Tensor):
+            %2 : int[] = prim::Constant[value=[1, 0, 3, 7]]()
+            %3 : int[] = prim::Constant[value=[0, 1, 2, 3]]()
+            %4 : Tensor = aten::roll(%1, %2, %3)
+            return (%4))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, g.get());
+
+  // Run Pytorch
+  auto in = at::randint(1, 10, {2, 3, 4, 5}, {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});
+  auto trt = trt_results[0].reshape(jit_results[0].sizes());
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}
+
+TEST(Converters, ATenRollShiftsNegativeConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%1 : Tensor):
+            %2 : int[] = prim::Constant[value=[0, -3, -3]]()
+            %3 : int[] = prim::Constant[value=[1, 2, 3]]()
+            %4 : Tensor = aten::roll(%1, %2, %3)
+            return (%4))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, g.get());
+
+  // Run Pytorch
+  auto in = at::randint(1, 10, {1, 3, 4, 5}, {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});
+  auto trt = trt_results[0].reshape(jit_results[0].sizes());
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}
+
+TEST(Converters, ATenRollDimsNegativeConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%1 : Tensor):
+            %2 : int[] = prim::Constant[value=[0, -3, -3]]()
+            %3 : int[] = prim::Constant[value=[1, 2, -1]]()
+            %4 : Tensor = aten::roll(%1, %2, %3)
+            return (%4))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, g.get());
+
+  // Run Pytorch
+  auto in = at::randint(1, 10, {1, 3, 4, 5}, {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});
+  auto trt = trt_results[0].reshape(jit_results[0].sizes());
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}
+
 TEST(Converters, ATenSliceConvertsCorrectly) {
   const auto graph = R"IR(
         graph(%x.1 : Tensor):