feat: Moved from converter to lowering pass

Signed-off-by: Anurag Dixit <[email protected]>

Author: Anurag Dixit

core/conversion/converters/impl/expand.cpp

@@ -194,60 +194,6 @@ bool add_expand_dynamic(
   return true;
 }
 
-bool add_repeat(ConversionCtx* ctx, const torch::jit::Node* n, args& args, const std::string& layer) {
-  auto in = args[0].ITensorOrFreeze(ctx);
-  auto input_dims = in->getDimensions();
-  auto repeats = args[1].unwrapToIntList().vec();
-  int repeats_rank = repeats.size();
-  TORCHTRT_CHECK(
-      repeats_rank >= input_dims.nbDims,
-      "Number of repeat dimensions cannot be smaller than number of input dimensions");
-
-  auto num_expand_dims = repeats_rank - input_dims.nbDims;
-
-  if (ctx->input_is_dynamic) {
-    int input_rank = input_dims.nbDims;
-    int output_rank = repeats_rank;
-    auto new_input_shape_tensor = concat(output_rank, input_rank, ctx, in);
-
-    auto shuffle = ctx->net->addShuffle(*in);
-    shuffle->setInput(1, *new_input_shape_tensor);
-    in = shuffle->getOutput(0);
-  } else {
-    if (num_expand_dims > 0) {
-      nvinfer1::Dims reshape_dims;
-      reshape_dims.nbDims = repeats.size();
-      for (int i = 0; i < num_expand_dims; i++) {
-        reshape_dims.d[i] = 1;
-      }
-      for (int i = 0; i < input_dims.nbDims; i++) {
-        reshape_dims.d[num_expand_dims + i] = input_dims.d[i];
-      }
-      // Add a reshape layer to expand dims
-      auto reshape_layer = ctx->net->addShuffle(*in);
-      reshape_layer->setReshapeDimensions(reshape_dims);
-      in = reshape_layer->getOutput(0);
-      LOG_DEBUG("Input reshaped to : " << in->getDimensions() << " from " << input_dims);
-    }
-    LOG_DEBUG("Repeats: " << repeats);
-  }
-
-  // Concat across all repeat axes.
-  for (int i = repeats.size() - 1; i >= 0; --i) {
-    std::vector<nvinfer1::ITensor*> tensors_vec;
-    for (int j = 0; j < repeats[i]; j++) {
-      tensors_vec.push_back(in);
-    }
-    auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
-    concat_layer->setAxis(i);
-    in = concat_layer->getOutput(0);
-  }
-
-  auto out = ctx->AssociateValueAndTensor(n->outputs()[0], in);
-  LOG_DEBUG(layer << " layer output tensor shape: " << out->getDimensions());
-  return true;
-}
-
 auto expand_registrations TORCHTRT_UNUSED =
     RegisterNodeConversionPatterns()
         .pattern(
@@ -284,7 +230,59 @@ auto expand_registrations TORCHTRT_UNUSED =
         .pattern(
             {"aten::repeat(Tensor self, int[] repeats) -> (Tensor)",
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
-               return add_repeat(ctx, n, args, "Repeat");
+               auto in = args[0].ITensorOrFreeze(ctx);
+               auto input_dims = in->getDimensions();
+               auto repeats = args[1].unwrapToIntList().vec();
+               int repeats_rank = repeats.size();
+               TORCHTRT_CHECK(
+                   repeats_rank >= input_dims.nbDims,
+                   "Number of repeat dimensions cannot be smaller than number of input dimensions");
+               auto num_expand_dims = repeats_rank - input_dims.nbDims;
+
+               if (ctx->input_is_dynamic) {
+                 int input_rank = input_dims.nbDims;
+                 int output_rank = repeats_rank;
+                 auto new_input_shape_tensor = concat(output_rank, input_rank, ctx, in);
+
+                 // Add a reshape layer to expand dims
+                 auto shuffle = ctx->net->addShuffle(*in);
+                 shuffle->setInput(1, *new_input_shape_tensor);
+                 in = shuffle->getOutput(0);
+               } else {
+                 if (num_expand_dims > 0) {
+                   nvinfer1::Dims reshape_dims;
+                   reshape_dims.nbDims = repeats.size();
+                   for (int i = 0; i < num_expand_dims; i++) {
+                     reshape_dims.d[i] = 1;
+                   }
+                   for (int i = 0; i < input_dims.nbDims; i++) {
+                     reshape_dims.d[num_expand_dims + i] = input_dims.d[i];
+                   }
+                   // Add a reshape layer to expand dims
+                   auto reshape_layer = ctx->net->addShuffle(*in);
+                   reshape_layer->setReshapeDimensions(reshape_dims);
+                   in = reshape_layer->getOutput(0);
+                   LOG_DEBUG("Input reshaped to : " << in->getDimensions() << " from " << input_dims);
+                 }
+                 LOG_DEBUG("Repeats: " << repeats);
+               }
+
+               // Concat across all repeat axes.
+               // TODO: Implementation might not be performant. Explore other strategies to improve performance.
+               for (int i = repeats.size() - 1; i >= 0; --i) {
+                 std::vector<nvinfer1::ITensor*> tensors_vec;
+                 for (int j = 0; j < repeats[i]; j++) {
+                   tensors_vec.push_back(in);
+                 }
+                 auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
+                 concat_layer->setAxis(i);
+                 in = concat_layer->getOutput(0);
+               }
+
+               auto out = ctx->AssociateValueAndTensor(n->outputs()[0], in);
+
+               LOG_DEBUG("Repeat layer output tensor shape: " << out->getDimensions());
+               return true;
              }})
         .pattern(
             {"aten::repeat_interleave.self_int(Tensor self, int repeats, int? dim=None, *, int? output_size=None) -> (Tensor)",
@@ -397,11 +395,6 @@ auto expand_registrations TORCHTRT_UNUSED =
 
                return true;
              }})
-        .pattern(
-            {"aten::tile(Tensor self, int[] dims) -> (Tensor)",
-             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
-               return add_repeat(ctx, n, args, "Tile");
-             }})
         .pattern(
             {"aten::meshgrid(Tensor[] tensors) -> (Tensor[])",
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
@@ -491,4 +484,4 @@ auto expand_registrations TORCHTRT_UNUSED =
 } // namespace converters
 } // namespace conversion
 } // namespace core
-} // namespace torch_tensorrt
+} // namespace torch_tensorrt

core/lowering/lowering.cpp

@@ -153,6 +153,7 @@ void LowerGraph(std::shared_ptr<torch::jit::Graph>& g, std::vector<torch::jit::I
   passes::ReplaceScalarImplicit(g);
   passes::RewriteInputsWithParams(g, params);
   passes::ReplaceAtenPad(g);
+  passes::ReplaceTileWithRepeat(g);
   LOG_GRAPH(*g);
 }
 

core/lowering/passes/BUILD

@@ -31,6 +31,7 @@ cc_library(
         "replace_aten_pad.cpp",
         "rewrite_inputs_with_params.cpp",
         "silu_to_sigmoid_multiplication.cpp",
+        "tile_to_repeat.cpp",
         "unpack_addmm.cpp",
         "unpack_batch_norm.cpp",
         "unpack_hardsigmoid.cpp",

core/lowering/passes/CMakeLists.txt

@@ -6,6 +6,7 @@ target_sources(${lib_name}
             "${CMAKE_CURRENT_SOURCE_DIR}/linear_to_addmm.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/module_fallback.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/op_aliasing.cpp"
+	    "${CMAKE_CURRENT_SOURCE_DIR}/tile_to_repeat.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/reduce_gelu.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/reduce_remainder.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/reduce_to.cpp"

core/lowering/passes/passes.h

@@ -50,6 +50,7 @@ void UnpackAndCastNumToTensor(std::shared_ptr<torch::jit::Graph>& graph, std::st
 void UnpackAndCastFull(std::shared_ptr<torch::jit::Graph>& graph, std::string target_device_name);
 void ReplaceScalarImplicit(std::shared_ptr<torch::jit::Graph>& graph);
 void ReplaceAtenPad(std::shared_ptr<torch::jit::Graph>& graph);
+void ReplaceTileWithRepeat(std::shared_ptr<torch::jit::Graph>& graph);
 
 // utility functions exposed for testing
 std::string unmangle_cls_name(const std::string& name);

core/lowering/passes/tile_to_repeat.cpp

@@ -0,0 +1,25 @@
+#include "core/util/prelude.h"
+#include "torch/csrc/jit/passes/subgraph_rewrite.h"
+
+namespace torch_tensorrt {
+namespace core {
+namespace lowering {
+namespace passes {
+void ReplaceTileWithRepeat(std::shared_ptr<torch::jit::Graph>& graph) {
+  std::string tile_pattern = R"IR(
+                        graph(%input, %1):
+                            %2 = aten::tile(%input, %1)
+                            return (%2))IR";
+  std::string repeat_pattern = R"IR(
+                        graph(%input, %1):
+                            %2 = aten::repeat(%input, %1)
+                            return (%2))IR";
+  torch::jit::SubgraphRewriter tile_to_repeat;
+  tile_to_repeat.RegisterRewritePattern(tile_pattern, repeat_pattern);
+  tile_to_repeat.runOnGraph(graph);
+  LOG_GRAPH("Mapping tile -> repeat: " << *graph);
+}
+} // namespace passes
+} // namespace lowering
+} // namespace core
+} // namespace torch_tensorrt

docsrc/contributors/lowering.rst

@@ -205,3 +205,10 @@ Unroll Loops
     `torch/csrc/jit/passes/loop_unrolling.h <https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/passes/loop_unrolling.h>`_
 
 Unrolls the operations of compatable loops (e.g. sufficently short) so that you only have to go through the loop once.
+
+Replace Tile with Repeat
+***************************************
+
+    `Torch-TensorRT/core/lowering/passes/tile_to_repeat.cpp <https://github.com/pytorch/TensorRT/blob/master/core/lowering/passes/tile_to_repeat.cpp>`_
+
+Removes dropout operators since we are doing inference.

tests/core/conversion/converters/test_expand.cpp

@@ -1,6 +1,7 @@
 #include <torch/torch.h>
 #include <string>
 #include "core/compiler.h"
+#include "core/lowering/passes/passes.h"
 #include "gtest/gtest.h"
 #include "tests/util/util.h"
 #include "torch/csrc/jit/ir/irparser.h"
@@ -680,6 +681,7 @@ TEST(Converters, ATenTileConvertsCorrectly) {
   auto g = std::make_shared<torch::jit::Graph>();
 
   torch::jit::parseIR(graph, g.get());
+  torch_tensorrt::core::lowering::passes::ReplaceTileWithRepeat(g);
 
   auto in = at::randint(1, 10, {1, 3}, {at::kCUDA});
 
@@ -704,6 +706,7 @@ TEST(Converters, ATenTileRepeatRankConvertsCorrectly) {
   auto g = std::make_shared<torch::jit::Graph>();
 
   torch::jit::parseIR(graph, g.get());
+  torch_tensorrt::core::lowering::passes::ReplaceTileWithRepeat(g);
 
   auto in = at::randint(1, 10, {1, 3}, {at::kCUDA});
 
@@ -728,6 +731,7 @@ TEST(Converters, ATenTileConvertsCorrectlyWithDynamicInput) {
   auto g = std::make_shared<torch::jit::Graph>();
 
   torch::jit::parseIR(graph, g.get());
+  torch_tensorrt::core::lowering::passes::ReplaceTileWithRepeat(g);
 
   auto in = at::randint(1, 10, {1, 3}, {at::kCUDA});
 
@@ -752,6 +756,7 @@ TEST(Converters, ATenTile3dConvertsCorrectly) {
   auto g = std::make_shared<torch::jit::Graph>();
 
   torch::jit::parseIR(graph, g.get());
+  torch_tensorrt::core::lowering::passes::ReplaceTileWithRepeat(g);
 
   auto in = at::randint(1, 10, {2, 3, 2}, {at::kCUDA});
 
@@ -776,6 +781,7 @@ TEST(Converters, ATenTile3dConvertsCorrectlyWithDynamicInput) {
   auto g = std::make_shared<torch::jit::Graph>();
 
   torch::jit::parseIR(graph, g.get());
+  torch_tensorrt::core::lowering::passes::ReplaceTileWithRepeat(g);
 
   auto in = at::randint(1, 10, {2, 3, 2}, {at::kCUDA});
 

tests/core/lowering/BUILD

@@ -103,6 +103,10 @@ lowering_test(
     name = "test_replace_aten_pad_pass",
 )
 
+lowering_test(
+    name = "test_tile_to_repeat_pass",
+)
+
 test_suite(
     name = "lowering_tests",
     tests = [
@@ -122,6 +126,7 @@ test_suite(
         ":test_remove_unnecessary_casts",
         ":test_replace_aten_pad_pass",
         ":test_rewrite_inputs_with_params",
+        ":test_tile_to_repeat_pass",
         ":test_unpack_hardsigmoid",
         ":test_unpack_hardswish",
         ":test_unpack_reduce_ops",

tests/core/lowering/test_tile_to_repeat_pass.cpp

@@ -0,0 +1,26 @@
+#include <string>
+#include "core/compiler.h"
+#include "core/lowering/passes/passes.h"
+#include "gtest/gtest.h"
+#include "tests/util/util.h"
+#include "torch/csrc/jit/ir/irparser.h"
+#include "torch/csrc/jit/ir/subgraph_matcher.h"
+
+TEST(LoweringPasses, TileToRepeatCorrectly) {
+  std::string source_graph = R"IR(
+        graph(%input, %dim):
+            %o : Tensor = aten::tile(%input, %dim)
+            return (%o))IR";
+  std::string target_graph = R"IR(
+        graph(%input, %dim):
+            %o : Tensor = aten::repeat(%input, %dim)
+            return (%o))IR";
+  auto sg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(source_graph, sg.get());
+  torch_tensorrt::core::lowering::passes::ReplaceTileWithRepeat(sg);
+
+  auto tg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(target_graph, tg.get());
+
+  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
+}