fix: Replace dropout lowering pass

- Remove existing dropout removal lowering pass implementation due to
bug
- Use Torch JIT dropout removal lowering pass to resolve bug where
nested dropouts resulted in invalid graph
- Existing removal process left artifacts in graph which caused an
internal assertion error
- Add regression test to catch nested dropout bug
- Update tests to remove testing for `feature_alpha_dropout` and
`feature_alpha_dropout_`, which are not removed by the JIT lowering pass
and can be added in later

Author: gs-olive

core/lowering/passes/remove_dropout.cpp

@@ -1,3 +1,4 @@
+#include "torch/csrc/jit/passes/remove_dropout.h"
 #include <torch/csrc/jit/passes/subgraph_rewrite.h>
 
 #include "core/util/prelude.h"
@@ -8,85 +9,7 @@ namespace lowering {
 namespace passes {
 
 void RemoveDropout(std::shared_ptr<torch::jit::Graph>& graph) {
-  std::string dropout_pattern = R"IR(
-        graph(%input, %4, %5):
-            %6 = aten::dropout(%input, %4, %5)
-            return (%6))IR";
-  std::string no_dropout_pattern = R"IR(
-        graph(%input, %4, %5):
-            return (%input))IR";
-
-  torch::jit::SubgraphRewriter remove_dropout;
-  remove_dropout.RegisterRewritePattern(dropout_pattern, no_dropout_pattern);
-  remove_dropout.runOnGraph(graph);
-
-  std::string dropout_inplace_pattern = R"IR(
-        graph(%input, %4, %5):
-            %6 = aten::dropout_(%input, %4, %5)
-            return (%6))IR";
-  std::string no_dropout_inplace_pattern = R"IR(
-        graph(%input, %4, %5):
-            return (%input))IR";
-
-  torch::jit::SubgraphRewriter remove_dropout_inplace_pattern;
-  remove_dropout_inplace_pattern.RegisterRewritePattern(dropout_inplace_pattern, no_dropout_inplace_pattern);
-  remove_dropout_inplace_pattern.runOnGraph(graph);
-
-  // remove feature_dropout
-  std::string feature_dropout_pattern = R"IR(
-        graph(%input, %4, %5):
-            %6 = aten::feature_dropout(%input, %4, %5)
-            return (%6))IR";
-  std::string no_feature_dropout_pattern = R"IR(
-        graph(%input, %4, %5):
-            return (%input))IR";
-
-  torch::jit::SubgraphRewriter remove_feature_dropout_pattern;
-  remove_feature_dropout_pattern.RegisterRewritePattern(feature_dropout_pattern, no_feature_dropout_pattern);
-  remove_feature_dropout_pattern.runOnGraph(graph);
-
-  // remove feature_dropout inplace
-  std::string feature_dropout_inplace_pattern = R"IR(
-        graph(%input, %4, %5):
-            %6 = aten::feature_dropout_(%input, %4, %5)
-            return (%6))IR";
-  std::string no_feature_dropout_inplace_pattern = R"IR(
-        graph(%input, %4, %5):
-            return (%input))IR";
-
-  torch::jit::SubgraphRewriter remove_feature_dropout_inplace_pattern;
-  remove_feature_dropout_inplace_pattern.RegisterRewritePattern(
-      feature_dropout_inplace_pattern, no_feature_dropout_inplace_pattern);
-  remove_feature_dropout_inplace_pattern.runOnGraph(graph);
-
-  // remove feature_alpha_dropout
-  std::string feature_alpha_dropout_pattern = R"IR(
-        graph(%input, %4, %5):
-            %6 = aten::feature_alpha_dropout(%input, %4, %5)
-            return (%6))IR";
-  std::string no_feature_alpha_dropout_pattern = R"IR(
-        graph(%input, %4, %5):
-            return (%input))IR";
-
-  torch::jit::SubgraphRewriter remove_feature_alpha_dropout_pattern;
-  remove_feature_alpha_dropout_pattern.RegisterRewritePattern(
-      feature_alpha_dropout_pattern, no_feature_alpha_dropout_pattern);
-  remove_feature_alpha_dropout_pattern.runOnGraph(graph);
-
-  // remove feature_alpha_dropout inplace
-  std::string feature_alpha_dropout_inplace_pattern = R"IR(
-        graph(%input, %4, %5):
-            %6 = aten::feature_alpha_dropout_(%input, %4, %5)
-            return (%6))IR";
-  std::string no_feature_alpha_dropout_inplace_pattern = R"IR(
-        graph(%input, %4, %5):
-            return (%input))IR";
-
-  torch::jit::SubgraphRewriter remove_feature_alpha_dropout_inplace_pattern;
-  remove_feature_alpha_dropout_inplace_pattern.RegisterRewritePattern(
-      feature_alpha_dropout_inplace_pattern, no_feature_alpha_dropout_inplace_pattern);
-  remove_feature_alpha_dropout_inplace_pattern.runOnGraph(graph);
-
+  torch::jit::removeDropout(graph);
   LOG_GRAPH("Post remove dropout: " << *graph);
 }
 

tests/core/lowering/test_remove_dropout_pass.cpp

@@ -32,39 +32,15 @@ TEST(LoweringPasses, RemoveDropoutLowersCorrectly) {
   ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
 }
 
-TEST(LoweringPasses, RemoveDropoutInplaceLowersCorrectly) {
-  std::string source_graph = R"IR(
-    graph(%x.1):
-      %3 : float = prim::Constant[value=0.5]()
-      %4 : bool = prim::Constant[value=0]()
-      %y.1 : Tensor = aten::dropout_(%x.1, %3, %4)
-      %11 : Tensor = aten::relu(%y.1)
-      return (%11))IR";
-  std::string target_graph = R"IR(
-    graph(%x.1):
-      %11 : Tensor = aten::relu(%x.1)
-      return (%11))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.get());
-  torch_tensorrt::core::lowering::passes::RemoveDropout(sg);
-
-  auto tg = std::make_shared<torch::jit::Graph>();
-  torch::jit::parseIR(target_graph, tg.get());
-
-  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
-}
-
-TEST(LoweringPasses, RemoveFeatureDropoutLowersCorrectly) {
+TEST(LoweringPasses, RemoveDropoutNestedLowersCorrectly) {
   std::string source_graph = R"IR(
     graph(%x.1):
       %3 : float = prim::Constant[value=0.5]()
       %4 : bool = prim::Constant[value=0]()
-      %y.1 : Tensor = aten::feature_dropout(%x.1, %3, %4)
-      %11 : Tensor = aten::relu(%y.1)
-      return (%11))IR";
+      %y.1 : Tensor = aten::dropout(%x.1, %3, %4)
+      %z.1 : Tensor = aten::dropout(%y.1, %3, %4)
+      %12 : Tensor = aten::relu(%z.1)
+      return (%12))IR";
   std::string target_graph = R"IR(
     graph(%x.1):
       %11 : Tensor = aten::relu(%x.1)
@@ -82,12 +58,12 @@ TEST(LoweringPasses, RemoveFeatureDropoutLowersCorrectly) {
   ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
 }
 
-TEST(LoweringPasses, RemoveFeatureDropoutInplaceLowersCorrectly) {
+TEST(LoweringPasses, RemoveDropoutInplaceLowersCorrectly) {
   std::string source_graph = R"IR(
     graph(%x.1):
       %3 : float = prim::Constant[value=0.5]()
       %4 : bool = prim::Constant[value=0]()
-      %y.1 : Tensor = aten::feature_dropout_(%x.1, %3, %4)
+      %y.1 : Tensor = aten::dropout_(%x.1, %3, %4)
       %11 : Tensor = aten::relu(%y.1)
       return (%11))IR";
   std::string target_graph = R"IR(
@@ -107,12 +83,12 @@ TEST(LoweringPasses, RemoveFeatureDropoutInplaceLowersCorrectly) {
   ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
 }
 
-TEST(LoweringPasses, RemoveFeatureAlphaDropoutLowersCorrectly) {
+TEST(LoweringPasses, RemoveFeatureDropoutLowersCorrectly) {
   std::string source_graph = R"IR(
     graph(%x.1):
       %3 : float = prim::Constant[value=0.5]()
       %4 : bool = prim::Constant[value=0]()
-      %y.1 : Tensor = aten::feature_alpha_dropout(%x.1, %3, %4)
+      %y.1 : Tensor = aten::feature_dropout(%x.1, %3, %4)
       %11 : Tensor = aten::relu(%y.1)
       return (%11))IR";
   std::string target_graph = R"IR(
@@ -132,12 +108,12 @@ TEST(LoweringPasses, RemoveFeatureAlphaDropoutLowersCorrectly) {
   ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
 }
 
-TEST(LoweringPasses, RemoveFeatureAlphaDropoutInplaceLowersCorrectly) {
+TEST(LoweringPasses, RemoveFeatureDropoutInplaceLowersCorrectly) {
   std::string source_graph = R"IR(
     graph(%x.1):
       %3 : float = prim::Constant[value=0.5]()
       %4 : bool = prim::Constant[value=0]()
-      %y.1 : Tensor = aten::feature_alpha_dropout_(%x.1, %3, %4)
+      %y.1 : Tensor = aten::feature_dropout_(%x.1, %3, %4)
       %11 : Tensor = aten::relu(%y.1)
       return (%11))IR";
   std::string target_graph = R"IR(