add aten::detach lowering pass

Signed-off-by: inocsin <[email protected]>

Author: inocsin

core/lowering/passes/remove_to.cpp

@@ -21,6 +21,7 @@ struct ToRemoval {
 
   void run() {
     findTo(graph_->block());
+    findDetach(graph_->block());
     torch::jit::EliminateDeadCode(graph_);
     LOG_DEBUG(
         "RemoveTo - Note: Removing remaining aten::to operators, if type casts need to be preserved, add a pass before this pass is run");
@@ -39,6 +40,17 @@ struct ToRemoval {
     }
   }
 
+  void findDetach(Block* b) {
+    for (auto it = b->nodes().begin(); it != b->nodes().end(); it++) {
+      auto n = *it;
+      if (n->kind() == c10::Symbol::fromQualString("aten::detach")) {
+        LOG_GRAPH("Found that node " << *n << "  is an detach node (RemoveTo)" << std::endl);
+        n->outputs()[0]->replaceAllUsesWith(n->inputs()[0]);
+        it.destroyCurrent();
+      }
+    }
+  }
+
   std::shared_ptr<Graph> graph_;
 };
 } // namespace

tests/core/lowering/BUILD

@@ -15,10 +15,15 @@ lowering_test(
   name = "test_remove_to",
 )
 
+lowering_test(
+  name = "test_remove_detach_pass",
+)
+
 test_suite(
     name = "lowering_tests",
     tests = [
         ":test_remove_contiguous_pass",
-        ":test_remove_to"
+        ":test_remove_to",
+        ":test_remove_detach_pass"
     ]
 )

tests/core/lowering/test_remove_detach_pass.cpp

@@ -0,0 +1,28 @@
+#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, RemoveDetachCorrectly) {
+  std::string source_graph = R"IR(
+    graph(%input):
+      %2 = aten::detach(%input)
+      %3 = foo::bar(%2)
+      return (%3))IR";
+  std::string target_graph = R"IR(
+    graph(%input):
+      %3 = foo::bar(%input)
+      return (%3))IR";
+
+  auto sg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(source_graph, &*sg);
+  trtorch::core::lowering::passes::RemoveContiguous(sg);
+
+  auto tg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(target_graph, &*tg);
+
+  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
+}