Merge pull request #352 from NVIDIA/bowa_silu

Lowering pass for SiLU

Author: narendasan

core/lowering/lowering.cpp

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

core/lowering/passes/BUILD

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

core/lowering/passes/passes.h

@@ -20,6 +20,7 @@ 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);
+void SiluToSigmoidMultipication(std::shared_ptr<torch::jit::Graph>& graph);
 
 } // namespace passes
 } // namespace lowering

core/lowering/passes/silu_to_sigmoid_multiplication.cpp

@@ -0,0 +1,31 @@
+#include <torch/csrc/jit/passes/subgraph_rewrite.h>
+
+#include "core/util/prelude.h"
+
+namespace trtorch {
+namespace core {
+namespace lowering {
+namespace passes {
+
+void SiluToSigmoidMultipication(std::shared_ptr<torch::jit::Graph>& graph) {
+  std::string silu_pattern = R"IR(
+        graph(%x):
+            %1 : Tensor = aten::silu(%x)
+            return (%1))IR";
+  std::string sigmoid_multiplication_pattern = R"IR(
+        graph(%x):
+            %1 : Tensor = aten::sigmoid(%x)
+            %2 : Tensor = aten::mul(%x, %1)
+            return (%2))IR";
+  ;
+
+  torch::jit::SubgraphRewriter map_silu_to_sigmoid_multiplication;
+  map_silu_to_sigmoid_multiplication.RegisterRewritePattern(silu_pattern, sigmoid_multiplication_pattern);
+  map_silu_to_sigmoid_multiplication.runOnGraph(graph);
+  LOG_GRAPH("Post map silu -> x * sigmoid(x): " << *graph);
+}
+
+} // namespace passes
+} // namespace lowering
+} // namespace core
+} // namespace trtorch

tests/core/lowering/BUILD

@@ -23,6 +23,10 @@ lowering_test(
   name = "test_operator_aliasing_pass",
 )
 
+lowering_test(
+  name = "test_silu_to_sigmoid_multiplication",
+)
+
 test_suite(
     name = "lowering_tests",
     tests = [

tests/core/lowering/test_silu_to_sigmoid_multiplication.cpp

@@ -0,0 +1,29 @@
+#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, RemoveSiluLowersCorrectly) {
+  std::string source_graph = R"IR(
+      graph(%x.1 : Tensor):
+        %2 : Tensor = aten::silu(%x.1)
+        return (%2))IR";
+  std::string target_graph = R"IR(
+      graph(%x.1):
+        %2 : Tensor = aten::sigmoid(%x.1)
+        %3 : Tensor = aten::mul(%x.1, %2)
+        return (%3))IR";
+
+  trtorch::core::util::logging::get_logger().set_reportable_log_level(trtorch::core::util::logging::LogLevel::kGRAPH);
+  auto sg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(source_graph, &*sg);
+  trtorch::core::lowering::passes::SiluToSigmoidMultipication(sg);
+
+  auto tg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(target_graph, &*tg);
+
+  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
+}