Add dead code elimination pass (#10101)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/10101

Simple DCE enabled by knowledge of the actual outputs (stacked beneath this diff)

Reviewed By: yinghai

Differential Revision: D9107853

fbshipit-source-id: 0c38fe5fe408be2b7fc9e1fe6a5b7160c06ce79b

2 files changed, 100 insertions, 0 deletions

diff --git a/caffe2/opt/dead_code_elim.cc b/caffe2/opt/dead_code_elim.cc
new file mode 100644
index 0000000000..5e7359329b
--- /dev/null
+++ b/caffe2/opt/dead_code_elim.cc
@@ -0,0 +1,43 @@
+#include "caffe2/core/logging.h"
+#include "caffe2/opt/converter.h"
+#include "caffe2/opt/passes.h"
+
+namespace caffe2 {
+namespace opt {
+
+using namespace nom;
+using namespace nom::repr;
+
+void deadCodeElim(NNModule* nn) {
+  // Iteratively remove unconsumed non-external outputs.
+  bool changed = false;
+  do {
+    changed = false;
+    for (const auto& node : nn->dataFlow.getMutableNodes()) {
+      NOM_REQUIRE_OR_CONT(nn::is<repr::NeuralNetOperator>(node));
+
+      bool isUsed = false;
+      for (const auto& output : nn::getOutputs(node)) {
+        if (nn::hasConsumer(output) || nn->outputs.count(output)) {
+          isUsed = true;
+          break;
+        }
+      }
+
+      NOM_REQUIRE_OR_CONT(!isUsed);
+
+      // No outputs are used, delete them and the node itself.
+      for (const auto& output : nn::getOutputs(node)) {
+        nn->dataFlow.deleteNode(output);
+      }
+      nn->dataFlow.deleteNode(node);
+      changed = true;
+      break;
+    }
+  } while (changed);
+}
+
+REGISTER_OPT_PASS_FROM_FUNC(DeadCodeElim, deadCodeElim);
+
+} // namespace opt
+} // namespace caffe2
diff --git a/caffe2/opt/dead_code_elim_test.cc b/caffe2/opt/dead_code_elim_test.cc
new file mode 100644
index 0000000000..86b969cad5
--- /dev/null
+++ b/caffe2/opt/dead_code_elim_test.cc
@@ -0,0 +1,57 @@
+#include "caffe2/core/common.h"
+#include "caffe2/opt/converter.h"
+#include "caffe2/opt/passes.h"
+
+#include <gtest/gtest.h>
+
+TEST(DeadCodeElim, BasicElim) {
+  caffe2::NetDef net;
+  {
+    caffe2::OperatorDef* def = net.add_op();
+    def->set_type("Fake");
+    def->add_input("X");
+    def->add_output("Y");
+  }
+
+  auto nn = caffe2::convertToNNModule(net);
+  auto pass = caffe2::OptimizationPassRegistry()->Create("DeadCodeElim", &nn);
+  pass->run();
+  auto optimized_net = caffe2::convertToCaffe2Proto(nn, net);
+  EXPECT_EQ(optimized_net.op().size(), 0);
+}
+
+TEST(DeadCodeElim, BasicNoElim) {
+  caffe2::NetDef net;
+  {
+    caffe2::OperatorDef* def = net.add_op();
+    def->set_type("Fake");
+    def->add_input("X");
+    def->add_output("Y");
+  }
+  net.add_external_output("Y");
+
+  auto nn = caffe2::convertToNNModule(net);
+  auto pass = caffe2::OptimizationPassRegistry()->Create("DeadCodeElim", &nn);
+  pass->run();
+  auto optimized_net = caffe2::convertToCaffe2Proto(nn, net);
+  EXPECT_EQ(optimized_net.op().size(), 1);
+}
+
+TEST(DeadCodeElim, PartiallyUsedNoElim) {
+  caffe2::NetDef net;
+  {
+    caffe2::OperatorDef* def = net.add_op();
+    def->set_type("Fake");
+    def->add_input("X");
+    def->add_output("Y");
+    def->add_output("Z");
+  }
+  net.add_external_output("Y");
+  // Z is unused, but we should keep Fake because Y is
+
+  auto nn = caffe2::convertToNNModule(net);
+  auto pass = caffe2::OptimizationPassRegistry()->Create("DeadCodeElim", &nn);
+  pass->run();
+  auto optimized_net = caffe2::convertToCaffe2Proto(nn, net);
+  EXPECT_EQ(optimized_net.op().size(), 1);
+}