feat(aten::to): Remove remaining typecast operators (should be a very

late pass to run)

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

Author: narendasan

Diff for: core/lowering/lowering.cpp

@@ -1,6 +1,8 @@
+#include "torch/csrc/jit/passes/common_subexpression_elimination.h"
 #include "torch/csrc/jit/passes/dead_code_elimination.h"
 #include "torch/csrc/jit/passes/fuse_linear.h"
 #include "torch/csrc/jit/passes/freeze_module.h"
+#include "torch/csrc/jit/passes/loop_unrolling.h"
 #include "torch/csrc/jit/passes/lower_graph.h"
 #include "torch/csrc/jit/passes/lower_tuples.h"
 #include "torch/csrc/jit/passes/quantization.h"
@@ -30,11 +32,13 @@ void LowerGraph(std::shared_ptr<torch::jit::Graph>& g) {
     passes::FuseFlattenLinear(g);
     passes::Conv2DToConvolution(g);
     passes::FuseAddMMBranches(g);
+    torch::jit::EliminateCommonSubexpression(g);
+    torch::jit::UnrollLoops(g);
+    torch::jit::EliminateCommonSubexpression(g);
     passes::UnpackAddMM(g);
     //passes::UnpackBatchNorm(g);
     passes::UnpackLogSoftmax(g);
-    //passes::RemoveDimExeception(g);
-    //irfusers::UnpackBatchNorm(g);
+    passes::RemoveTo(g);
     torch::jit::EliminateDeadCode(g);
     LOG_GRAPH(*g);
 }

Diff for: core/lowering/passes/BUILD

@@ -19,6 +19,7 @@ cc_library(
         "fuse_flatten_linear.cpp",
         "remove_contiguous.cpp",
         "remove_dropout.cpp",
+        "remove_to.cpp",
         "unpack_addmm.cpp",
         "unpack_batch_norm.cpp",
         "unpack_log_softmax.cpp",

Diff for: core/lowering/passes/passes.h

@@ -13,6 +13,7 @@ void FuseFlattenLinear(std::shared_ptr<torch::jit::Graph>& graph);
 void EliminateExceptionOrPassPattern(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 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);

Diff for: core/lowering/passes/remove_dropout.cpp

@@ -16,7 +16,6 @@ void RemoveDropout(std::shared_ptr<torch::jit::Graph>& graph) {
         graph(%input, %4, %5):
             return (%input))IR";
 
-    // replace matmul + add pattern to linear
     torch::jit::SubgraphRewriter remove_dropout;
     remove_dropout.RegisterRewritePattern(
         dropout_pattern, no_dropout_pattern);

Diff for: core/lowering/passes/remove_to.cpp

@@ -0,0 +1,54 @@
+#include "torch/csrc/jit/passes/guard_elimination.h"
+#include "torch/csrc/jit/ir/alias_analysis.h"
+#include "torch/csrc/jit/jit_log.h"
+#include "torch/csrc/jit/passes/constant_propagation.h"
+#include "torch/csrc/jit/passes/peephole.h"
+#include "torch/csrc/jit/runtime/graph_executor.h"
+#include "torch/csrc/jit/passes/dead_code_elimination.h"
+
+#include "core/util/prelude.h"
+
+#include <vector>
+
+namespace trtorch {
+namespace core {
+namespace lowering {
+namespace passes {
+namespace {
+using namespace torch::jit;
+struct ToRemoval {
+  ToRemoval(std::shared_ptr<Graph> graph)
+    : graph_(std::move(graph)) {}
+
+  void run() {
+    findTo(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");
+    LOG_GRAPH("Post aten::to removal: " << *graph_);
+  }
+
+private:
+  void findTo(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);
+        n->outputs()[0]->replaceAllUsesWith(n->inputs()[1]);
+        it.destroyCurrent();
+      }
+    }
+  }
+
+  std::shared_ptr<Graph> graph_;
+};
+} // namespace
+
+void RemoveTo(std::shared_ptr<Graph> graph) {
+  ToRemoval tr(std::move(graph));
+  tr.run();
+}
+
+} // namespace passes
+} // namespace lowering
+} // namespace core
+} // namespace trtorch