Add lowering for aten::remainder (with testcase)

Signed-off-by: Cheng Hang <[email protected]>

Author: Njuapp

core/lowering/lowering.cpp

@@ -44,6 +44,7 @@ void LowerGraph(std::shared_ptr<torch::jit::Graph>& g, LowerInfo lower_info) {
   passes::EliminateExceptionOrPassPattern(g);
   passes::ReduceToOperation(g);
   passes::ReduceGelu(g);
+  passes::ReduceRemainder(g);
   passes::RemoveContiguous(g);
   passes::ViewToReshape(g);
   passes::RemoveDropout(g);

core/lowering/passes/BUILD

@@ -18,6 +18,7 @@ cc_library(
         "op_aliasing.cpp",
         "reduce_to.cpp",
         "reduce_gelu.cpp",
+        "reduce_remainder.cpp",
         "remove_bn_dim_check.cpp",
         "remove_contiguous.cpp",
         "view_to_reshape.cpp",

core/lowering/passes/passes.h

@@ -21,6 +21,7 @@ void LinearToAddMM(std::shared_ptr<torch::jit::Graph>& graph);
 void EliminateExceptionOrPassPattern(std::shared_ptr<torch::jit::Graph> graph);
 void ReduceToOperation(std::shared_ptr<torch::jit::Graph>& graph);
 void ReduceGelu(std::shared_ptr<torch::jit::Graph>& graph);
+void ReduceRemainder(std::shared_ptr<torch::jit::Graph>& graph);
 void MarkNodesForFallback(std::shared_ptr<torch::jit::Graph>& g, bool delete_delims);
 void RemoveBNDimCheck(std::shared_ptr<torch::jit::Graph> graph);
 void RemoveContiguous(std::shared_ptr<torch::jit::Graph>& graph);

core/lowering/passes/reduce_remainder.cpp

@@ -0,0 +1,48 @@
+#include <torch/csrc/jit/passes/subgraph_rewrite.h>
+#include "core/util/prelude.h"
+
+namespace torch_tensorrt {
+namespace core {
+namespace lowering {
+namespace passes {
+
+void ReduceRemainder(std::shared_ptr<torch::jit::Graph>& graph) {
+  std::string remainder_pattern = R"IR(
+        graph(%self : Tensor, %other : Tensor):
+            %out : Tensor = aten::remainder(%self, %other)
+            return (%out))IR";
+
+  std::string remainder_reduce_pattern = R"IR(
+        graph(%self : Tensor, %other : Tensor):
+            %alpha : int = prim::Constant[value=1]()
+            %floor: Tensor = aten::floor_divide(%self, %other)
+            %prod: Tensor = aten::mul(%floor, %other)
+            %out: Tensor = aten::sub(%self, %prod, %alpha)
+            return (%out))IR";
+
+  std::string remainder_scalar_pattern = R"IR(
+        graph(%self : Tensor, %other : Scalar):
+            %out : Tensor = aten::remainder(%self, %other)
+            return (%out))IR";
+
+  std::string remainder_scalar_reduce_pattern = R"IR(
+        graph(%self : Tensor, %other : Scalar):
+            %alpha : int = prim::Constant[value=1]()
+            %floor: Tensor = aten::floor_divide(%self, %other)
+            %prod: Tensor = aten::mul(%floor, %other)
+            %out: Tensor = aten::sub(%self, %prod, %alpha)
+            return (%out))IR";
+
+  // replace aten::remainder with pointwise operations
+  torch::jit::SubgraphRewriter map_remainder_to_pointwise_ops;
+  map_remainder_to_pointwise_ops.RegisterRewritePattern(remainder_pattern, remainder_reduce_pattern);
+  map_remainder_to_pointwise_ops.RegisterRewritePattern(remainder_scalar_pattern, remainder_scalar_reduce_pattern);
+  map_remainder_to_pointwise_ops.runOnGraph(graph);
+
+  LOG_GRAPH("Post lowering of [aten::remainder] -> " << *graph);
+}
+
+} // namespace passes
+} // namespace lowering
+} // namespace core
+} // namespace torch_tensorrt

tests/core/conversion/converters/test_element_wise.cpp

@@ -1,5 +1,6 @@
 #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"
@@ -423,3 +424,52 @@ TEST(Converters, ATenLEScalarConvertsCorrectly) {
       return (%2))IR";
   pointwise_test_helper(graph, true, false, {5, 5});
 }
+
+TEST(Converters, ATenRemainderConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor, %1 : Tensor):
+        %2 : Tensor = aten::remainder(%0, %1)
+        return (%2))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto input1 = at::randint(-5, 5, {4, 5}, {at::kCUDA});
+  auto input2 = at::randint(1, 5, {5}, {at::kCUDA});
+
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {input1, input2});
+
+  torch_tensorrt::core::lowering::passes::ReduceRemainder(g);
+
+  input1 = at::clone(input1);
+  input2 = at::clone(input2);
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {input1, input2});
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 5e-2));
+}
+
+TEST(Converters, ATenRemainderWithScalarConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %scalar : float = prim::Constant[value=2.4]()
+        %1 : Tensor = aten::remainder(%0, %scalar)
+        return (%1))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(-5, 5, {5}, {at::kCUDA});
+
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in});
+
+  torch_tensorrt::core::lowering::passes::ReduceRemainder(g);
+
+  in = at::clone(in);
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in});
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
+}

tests/core/lowering/BUILD

@@ -46,6 +46,10 @@ lowering_test(
     name = "test_reduce_gelu",
 )
 
+lowering_test(
+    name = "test_reduce_remainder",
+)
+
 lowering_test(
     name = "test_remove_detach_pass",
 )
@@ -82,6 +86,7 @@ test_suite(
         ":test_view_to_reshape_pass",
         ":test_remove_dropout_pass",
         ":test_reduce_to_pass",
+        ":test_reduce_remainder",
         ":test_reduce_gelu",
         ":test_unpack_hardswish",
         ":test_unpack_reduce_ops"

tests/core/lowering/test_reduce_remainder.cpp

@@ -0,0 +1,57 @@
+#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, ReduceRemainderCorrectly) {
+  std::string source_graph = R"IR(
+        graph(%self : Tensor, %other : Tensor):
+            %out : Tensor = aten::remainder(%self, %other)
+            return (%out))IR";
+  std::string target_graph = R"IR(
+        graph(%self : Tensor, %other : Tensor):
+                %alpha : int = prim::Constant[value=1]()
+                %floor: Tensor = aten::floor_divide(%self, %other)
+                %prod: Tensor = aten::mul(%floor, %other)
+                %out: Tensor = aten::sub(%self, %prod, %alpha)
+                return (%out))IR";
+
+  torch_tensorrt::core::util::logging::get_logger().set_reportable_log_level(
+      torch_tensorrt::core::util::logging::LogLevel::kGRAPH);
+  auto sg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(source_graph, &*sg);
+  torch_tensorrt::core::lowering::passes::ReduceRemainder(sg);
+
+  auto tg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(target_graph, &*tg);
+
+  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
+}
+
+TEST(LoweringPasses, ReduceRemainderScalarCorrectly) {
+  std::string source_graph = R"IR(
+        graph(%self : Tensor, %other : Scalar):
+            %out : Tensor = aten::remainder(%self, %other)
+            return (%out))IR";
+  std::string target_graph = R"IR(
+        graph(%self : Tensor, %other : Scalar):
+            %alpha : int = prim::Constant[value=1]()
+            %floor: Tensor = aten::floor_divide(%self, %other)
+            %prod: Tensor = aten::mul(%floor, %other)
+            %out: Tensor = aten::sub(%self, %prod, %alpha)
+            return (%out))IR";
+
+  torch_tensorrt::core::util::logging::get_logger().set_reportable_log_level(
+      torch_tensorrt::core::util::logging::LogLevel::kGRAPH);
+  auto sg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(source_graph, &*sg);
+  torch_tensorrt::core::lowering::passes::ReduceRemainder(sg);
+
+  auto tg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(target_graph, &*tg);
+
+  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
+}