Merge pull request #318 from NVIDIA/true_divide_lowering

True divide lowering

Author: narendasan

core/conversion/converters/impl/element_wise.cpp

@@ -340,7 +340,7 @@ auto element_wise_registrations TRTORCH_UNUSED =
                     LOG_DEBUG("Output tensor shape: " << out->getDimensions());
                     return true;
                   }})
-        .pattern({"aten::div_.Scalar(Tensor self, Scalar other) -> (Tensor)",
+        .pattern({"aten::div_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)",
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                     auto self = args[0].ITensorOrFreeze(ctx);
                     auto otherScalar = args[1].unwrapToScalar().to<float>();

core/lowering/lowering.cpp

@@ -48,6 +48,7 @@ void LowerGraph(std::shared_ptr<torch::jit::Graph>& g) {
   // passes::UnpackBatchNorm(g);
   passes::UnpackLogSoftmax(g);
   passes::RemoveNOPs(g);
+  passes::AliasOperators(g);
   torch::jit::EliminateDeadCode(g);
   LOG_GRAPH(*g);
 }

core/lowering/passes/BUILD

@@ -25,6 +25,7 @@ cc_library(
         "unpack_addmm.cpp",
         "unpack_batch_norm.cpp",
         "unpack_log_softmax.cpp",
+        "op_aliasing.cpp"
     ],
     deps = [
         "//core/util:prelude",

core/lowering/passes/op_aliasing.cpp

@@ -0,0 +1,33 @@
+#include <torch/csrc/jit/passes/subgraph_rewrite.h>
+
+#include "core/util/prelude.h"
+
+namespace trtorch {
+namespace core {
+namespace lowering {
+namespace passes {
+
+void AliasOperators(std::shared_ptr<torch::jit::Graph>& graph) {
+  std::string true_divide_pattern = R"IR(
+        graph(%s, %o):
+            %1 : Tensor = aten::true_divide(%s, %o)
+            return (%1))IR";
+  std::string div_pattern = R"IR(
+        graph(%s, %o):
+            %1 : Tensor = aten::div(%s, %o)
+            return (%1))IR";
+  ;
+
+  // TODO
+  // complete other element wise pass
+
+  torch::jit::SubgraphRewriter true_divide_to_div;
+  true_divide_to_div.RegisterRewritePattern(true_divide_pattern, div_pattern);
+  true_divide_to_div.runOnGraph(graph);
+  LOG_GRAPH("Post map true_divide -> div: " << *graph);
+}
+
+} // namespace passes
+} // namespace lowering
+} // namespace core
+} // namespace trtorch

core/lowering/passes/passes.h

@@ -19,6 +19,7 @@ void RemoveNOPs(std::shared_ptr<torch::jit::Graph> graph);
 void UnpackAddMM(std::shared_ptr<torch::jit::Graph>& graph);
 void UnpackBatchNorm(std::shared_ptr<torch::jit::Graph>& graph);
 void UnpackLogSoftmax(std::shared_ptr<torch::jit::Graph>& graph);
+void AliasOperators(std::shared_ptr<torch::jit::Graph>& graph);
 
 } // namespace passes
 } // namespace lowering

tests/core/lowering/BUILD

@@ -19,11 +19,16 @@ lowering_test(
   name = "test_remove_detach_pass",
 )
 
+lowering_test(
+  name = "test_operator_aliasing_pass",
+)
+
 test_suite(
     name = "lowering_tests",
     tests = [
         ":test_remove_contiguous_pass",
         ":test_remove_to",
-        ":test_remove_detach_pass"
+        ":test_remove_detach_pass",
+        ":test_operator_aliasing_pass"
     ]
 )

tests/core/lowering/test_operator_aliasing_pass.cpp

@@ -0,0 +1,27 @@
+#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, LoweringTrueDivideCorrectly) {
+  std::string source_graph = R"IR(
+    graph(%s, %o):
+      %2 = aten::true_divide(%s, %o)
+      return (%2))IR";
+  std::string target_graph = R"IR(
+    graph(%s, %o):
+      %2 = aten::div(%s, %o)
+      return (%2))IR";
+
+  auto sg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(source_graph, &*sg);
+  trtorch::core::lowering::passes::ElementWisePass(sg);
+
+  auto tg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(target_graph, &*tg);
+
+  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
+}