Merge pull request #298 from NVIDIA/fix_remove_to_pass

Fix remove to pass

Author: narendasan

core/lowering/passes/remove_to.cpp

@@ -33,7 +33,7 @@ struct ToRemoval {
       auto n = *it;
       if (n->kind() == c10::Symbol::fromQualString("aten::to")) {
         LOG_GRAPH("Found that node " << *n << "  is an to node (RemoveTo)" << std::endl);
-        n->outputs()[0]->replaceAllUsesWith(n->inputs()[1]);
+        n->outputs()[0]->replaceAllUsesWith(n->inputs()[0]);
         it.destroyCurrent();
       }
     }

tests/core/lowering/BUILD

@@ -11,9 +11,14 @@ lowering_test(
   name = "test_remove_contiguous_pass",
 )
 
+lowering_test(
+  name = "test_remove_to",
+)
+
 test_suite(
     name = "lowering_tests",
     tests = [
-        ":test_remove_contiguous_pass"
+        ":test_remove_contiguous_pass",
+        ":test_remove_to"
     ]
 )

tests/core/lowering/test_remove_to.cpp

@@ -0,0 +1,31 @@
+#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, RemoveContiguousLowersCorrectly) {
+  std::string source_graph = R"IR(
+    graph(%x.1):
+      %6 : None = prim::Constant()
+      %4 : bool = prim::Constant[value=0]()
+      %3 : int = prim::Constant[value=5]() # experiments/test.py:8:17
+      %y.1 : Tensor = aten::to(%x.1, %3, %4, %4, %6)
+      %11 : Tensor = aten::relu(%y.1)
+      return (%3))IR";
+  std::string target_graph = R"IR(
+    graph(%x.1):
+      %11 : Tensor = aten::relu(%x.1)
+      return (%11))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());
+}