refactor(RemoveNOPs): Rename RemoveTo to RemoveNOPs since it now covers

both to and detach. We should be able to make this generic for a set of
operators that have no meaning in TensorRT.

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

core/lowering/lowering.cpp

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

core/lowering/passes/BUILD

@@ -21,7 +21,7 @@ cc_library(
         "remove_bn_dim_check.cpp",
         "remove_contiguous.cpp",
         "remove_dropout.cpp",
-        "remove_to.cpp",
+        "remove_nops.cpp",
         "unpack_addmm.cpp",
         "unpack_batch_norm.cpp",
         "unpack_log_softmax.cpp",

core/lowering/passes/passes.h

@@ -15,7 +15,7 @@ void EliminateExceptionOrPassPattern(std::shared_ptr<torch::jit::Graph> graph);
 void RemoveBNDimCheck(std::shared_ptr<torch::jit::Graph> graph);
 void RemoveContiguous(std::shared_ptr<torch::jit::Graph>& graph);
 void RemoveDropout(std::shared_ptr<torch::jit::Graph>& graph);
-void RemoveTo(std::shared_ptr<torch::jit::Graph> graph);
+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);

core/lowering/passes/remove_to.cpp renamed to core/lowering/passes/remove_nops.cpp

@@ -16,35 +16,35 @@ namespace lowering {
 namespace passes {
 namespace {
 using namespace torch::jit;
-struct ToRemoval {
-  ToRemoval(std::shared_ptr<Graph> graph) : graph_(std::move(graph)) {}
+struct NOPRemoval {
+  NOPRemoval(std::shared_ptr<Graph> graph) : graph_(std::move(graph)) {}
 
   void run() {
-    findTo(graph_->block());
-    findDetach(graph_->block());
+    removeTo(graph_->block());
+    removeDetach(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");
+        "RemoveNOPs - Note: Removing remaining aten::to operators (in addition to other ops that have no meaning in TRT), if type casts need to be preserved, add a pass before this pass is run");
     LOG_GRAPH("Post aten::to removal: " << *graph_);
   }
 
  private:
-  void findTo(Block* b) {
+  void removeTo(Block* b) {
     for (auto it = b->nodes().begin(); it != b->nodes().end(); it++) {
       auto n = *it;
       if (n->kind() == c10::Symbol::fromQualString("aten::to")) {
-        LOG_GRAPH("Found that node " << *n << "  is an to node (RemoveTo)" << std::endl);
+        LOG_GRAPH("Found that node " << *n << "  is an to node (RemoveNOPs)" << std::endl);
         n->outputs()[0]->replaceAllUsesWith(n->inputs()[0]);
         it.destroyCurrent();
       }
     }
   }
 
-  void findDetach(Block* b) {
+  void removeDetach(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);
+        LOG_GRAPH("Found that node " << *n << "  is an detach node (RemoveNOPs)" << std::endl);
         n->outputs()[0]->replaceAllUsesWith(n->inputs()[0]);
         it.destroyCurrent();
       }
@@ -55,8 +55,8 @@ struct ToRemoval {
 };
 } // namespace
 
-void RemoveTo(std::shared_ptr<Graph> graph) {
-  ToRemoval tr(std::move(graph));
+void RemoveNOPs(std::shared_ptr<Graph> graph) {
+  NOPRemoval tr(std::move(graph));
   tr.run();
 }
 

tests/core/lowering/test_remove_contiguous_pass.cpp

@@ -17,6 +17,7 @@ TEST(LoweringPasses, RemoveContiguousLowersCorrectly) {
       %3 = foo::bar(%input)
       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::RemoveContiguous(sg);

tests/core/lowering/test_remove_detach_pass.cpp

@@ -17,9 +17,10 @@ TEST(LoweringPasses, RemoveDetachCorrectly) {
       %3 = aten::sin(%input)
       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::RemoveTo(sg);
+  trtorch::core::lowering::passes::RemoveNOPs(sg);
 
   auto tg = std::make_shared<torch::jit::Graph>();
   torch::jit::parseIR(target_graph, &*tg);

tests/core/lowering/test_remove_to.cpp

@@ -1,28 +1,30 @@
 #include <string>
 #include "core/compiler.h"
 #include "core/lowering/passes/passes.h"
+#include "core/util/prelude.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) {
+TEST(LoweringPasses, RemoveToLowersCorrectly) {
   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";
+      return (%11))IR";
   std::string target_graph = R"IR(
     graph(%x.1):
       %11 : Tensor = aten::relu(%x.1)
       return (%11))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::RemoveContiguous(sg);
+  trtorch::core::lowering::passes::RemoveNOPs(sg);
 
   auto tg = std::make_shared<torch::jit::Graph>();
   torch::jit::parseIR(target_graph, &*tg);