Allow DCE to clean up some mutable ops (pytorch#14601)

Summary:
This PR makes DCE a little smarter in the presence of mutable ops. Previously mutable ops could never be cleaned up, now they can be cleaned up if we can prove there are no live uses of any alias sets that the op writes to.

This behavior is optional; if you pass DCE a block instead of a graph, it will do the same thing as before. Also changed `InlineAutographSubgraph` to use the common subgraph utils.

Tested on traced ResNet, and it gets rid of the dead code.
Pull Request resolved: pytorch#14601

Differential Revision: D13309118

Pulled By: suo

fbshipit-source-id: dac2791e7d2ecf219ae717a2759b83c1e927f254

Author: suo

test/expect/TestScript.test_mutable_dce.expect

@@ -0,0 +1,14 @@
+graph() {
+  %0 : int = prim::Constant[value=1]()
+  %1 : int[] = prim::Constant[value=[0, -1]]()
+  %2 : int = prim::Constant[value=0]()
+  %3 : int = prim::Constant[value=6]()
+  %4 : int = prim::Constant[value=2]()
+  %5 : int = prim::Constant[value=3]()
+  %6 : int[] = prim::ListConstruct(%4, %5)
+  %a.1 : Tensor = aten::rand(%6, %3, %2, %1)
+  %8 : int[] = prim::ListConstruct(%4, %5)
+  %9 : Tensor = aten::rand(%8, %3, %2, %1)
+  %a : Tensor = aten::add_(%a.1, %9, %0)
+  return (%a);
+}

test/expect/TestScript.test_mutable_dce_block.expect

@@ -0,0 +1,30 @@
+graph() {
+  %0 : int = prim::Constant[value=1]()
+  %1 : int[] = prim::Constant[value=[0, -1]]()
+  %2 : int = prim::Constant[value=0]()
+  %3 : int = prim::Constant[value=6]()
+  %4 : int = prim::Constant[value=2]()
+  %5 : int = prim::Constant[value=3]()
+  %6 : int[] = prim::ListConstruct(%4, %5)
+  %a.1 : Tensor = aten::rand(%6, %3, %2, %1)
+  %8 : int[] = prim::ListConstruct(%4, %5)
+  %9 : Tensor = aten::rand(%8, %3, %2, %1)
+  %a.2 : Tensor = aten::add_(%a.1, %9, %0)
+  %11 : int[] = prim::ListConstruct(%4, %5)
+  %b.1 : Tensor = aten::rand(%11, %3, %2, %1)
+  %13 : int[] = prim::ListConstruct(%4, %5)
+  %14 : Tensor = aten::zeros(%13, %3, %2, %1)
+  %15 : Tensor = aten::gt(%a.2, %14)
+  %16 : bool = prim::TensorToBool(%15)
+  %b : Tensor = prim::If(%16)
+    block0() {
+      %18 : int[] = prim::ListConstruct(%4, %5)
+      %19 : Tensor = aten::rand(%18, %3, %2, %1)
+      %b.2 : Tensor = aten::add_(%b.1, %19, %0)
+      -> (%b.2)
+    }
+    block1() {
+      -> (%b.1)
+    }
+  return (%b);
+}

test/expect/TestScript.test_mutable_dce_graph_input.expect

@@ -0,0 +1,12 @@
+graph(%a.1 : Tensor) {
+  %1 : int = prim::Constant[value=1]()
+  %2 : int[] = prim::Constant[value=[0, -1]]()
+  %3 : int = prim::Constant[value=0]()
+  %4 : int = prim::Constant[value=6]()
+  %5 : int = prim::Constant[value=2]()
+  %6 : int = prim::Constant[value=3]()
+  %7 : int[] = prim::ListConstruct(%5, %6)
+  %8 : Tensor = aten::rand(%7, %4, %3, %2)
+  %a : Tensor = aten::add_(%a.1, %8, %1)
+  return ();
+}

test/expect/TestScript.test_mutable_dce_list.expect

@@ -0,0 +1,17 @@
+graph(%a : Tensor) {
+  %1 : int = prim::Constant[value=1]()
+  %2 : int[] = prim::Constant[value=[0, -1]]()
+  %3 : int = prim::Constant[value=6]()
+  %4 : int = prim::Constant[value=0]()
+  %5 : int = prim::Constant[value=2]()
+  %6 : int = prim::Constant[value=3]()
+  %l : Tensor[] = prim::ListConstruct()
+  %8 : Tensor[] = aten::append(%l, %a)
+  %c.1 : Tensor = aten::select(%l, %4)
+  %10 : int[] = prim::ListConstruct(%5, %6)
+  %b : Tensor = aten::rand(%10, %3, %4, %2)
+  %12 : int[] = prim::ListConstruct(%5, %6)
+  %13 : Tensor = aten::rand(%12, %3, %4, %2)
+  %c : Tensor = aten::add_(%c.1, %13, %1)
+  return (%b);
+}

test/expect/TestScript.test_mutable_dce_loop.expect

@@ -0,0 +1,25 @@
+graph(%a : Tensor) {
+  %1 : int[] = prim::Constant[value=[0, -1]]()
+  %2 : int = prim::Constant[value=6]()
+  %i.1 : int = prim::Constant[value=0]()
+  %4 : int = prim::Constant[value=2]()
+  %5 : int = prim::Constant[value=3]()
+  %6 : int = prim::Constant[value=9223372036854775807]()
+  %7 : int = prim::Constant[value=1]()
+  %l : Tensor[] = prim::ListConstruct()
+  %9 : Tensor[] = aten::append(%l, %a)
+  %10 : int[] = prim::ListConstruct(%4, %5)
+  %b : Tensor = aten::rand(%10, %2, %i.1, %1)
+  %12 : bool = aten::lt(%i.1, %7)
+  %i : int = prim::Loop(%6, %12, %i.1)
+    block0(%14 : int, %15 : int) {
+      %c.1 : Tensor = aten::select(%l, %i.1)
+      %17 : int[] = prim::ListConstruct(%4, %5)
+      %18 : Tensor = aten::rand(%17, %2, %i.1, %1)
+      %c : Tensor = aten::add_(%c.1, %18, %7)
+      %i.2 : int = aten::add(%15, %7)
+      %21 : bool = aten::lt(%i.2, %7)
+      -> (%21, %i.2)
+    }
+  return (%b);
+}

test/test_jit.py

@@ -8881,6 +8881,68 @@ def fn(x, y):
 
         self.checkScript(fn, (torch.ones(2, 4, 2), torch.ones(2, 4, 2)))
 
+    def test_mutable_dce(self):
+        @torch.jit.script
+        def foo():
+            a = torch.rand(2, 3)
+            a += torch.rand(2, 3)
+            b = torch.rand(2, 3)
+            b += torch.rand(2, 3)
+            # b should be cleaned up but not a
+            return a
+
+        self.assertExpectedGraph(foo.graph)
+
+    def test_mutable_dce_block(self):
+        @torch.jit.script
+        def foo():
+            a = torch.rand(2, 3)
+            a += torch.rand(2, 3)
+            b = torch.rand(2, 3)
+            if bool(a > torch.zeros(2, 3)):
+                b += torch.rand(2, 3)
+                a += torch.rand(2, 3)
+            # a should be cleaned up but not b
+            return b
+
+        self.assertExpectedGraph(foo.graph)
+
+    def test_mutable_dce_graph_input(self):
+        @torch.jit.script
+        def foo(a):
+            a += torch.rand(2, 3)
+            # shouldn't clean up `a` even though it's not used in the output
+
+        self.assertExpectedGraph(foo.graph)
+
+    def test_mutable_dce_list(self):
+        @torch.jit.script
+        def foo(a):
+            l = []
+            l.append(a)
+            c = l[0]
+            b = torch.rand(2, 3)
+            c += torch.rand(2, 3)
+            return b
+
+        self.assertExpectedGraph(foo.graph)
+
+    def test_mutable_dce_loop(self):
+        @torch.jit.script
+        def foo(a):
+            l = []
+            l.append(a)
+            i = 0
+            b = torch.rand(2, 3)
+            while i < 1:
+                dead = torch.rand(2, 3)
+                c = l[0]
+                c += torch.rand(2, 3)
+                i += 1
+            return b
+
+        self.assertExpectedGraph(foo.graph)
+
 
 class MnistNet(nn.Module):
     def __init__(self):

torch/csrc/jit/export.cpp

@@ -96,7 +96,7 @@ void validateBlock(Block *b, onnx_torch::OperatorExportTypes operator_export_typ
 
 void validateGraph(const std::shared_ptr<Graph>& graph, onnx_torch::OperatorExportTypes operator_export_type) {
   validateBlock(graph->block(), operator_export_type);
-  EliminateDeadCode(graph);
+  EliminateDeadCode(graph->block());
 }
 
 class EncoderBase {

torch/csrc/jit/init.cpp

@@ -96,7 +96,7 @@ void initJITBindings(PyObject *module) {
    .def("_jit_pass_onnx_peephole", PeepholeOptimizeONNX)
    .def("_jit_pass_fuse", FuseGraph)
    .def("_jit_pass_dce", [](std::shared_ptr<Graph>& g) {
-     return EliminateDeadCode(g); // overload resolution
+     return EliminateDeadCode(g->block()); // overload resolution
    })
    .def("_jit_pass_cse", [](std::shared_ptr<Graph>& g) {
      return EliminateCommonSubexpression(g); // overload resolution

torch/csrc/jit/ir.cpp

@@ -862,18 +862,18 @@ Value* Node::insertOutput(size_t i) {
   return outputs_.at(i);
 }
 
-bool Node::isBefore(const Node * n) const {
-  if (this == n) {
-    return false;
-  }
-  return !isAfter(n);
-}
+bool Node::isBeforeOrAfter(const Node* n, MoveSide moveSide) const {
+  if (this->owningBlock() == n->owningBlock()) {
+    if (moveSide == MoveSide::BEFORE) {
+      return this->topo_position_ < n->topo_position_;
+    }
 
-bool Node::isAfter(const Node * n) const {
-  JIT_ASSERT(this->owningGraph() == n->owningGraph());
+    if (moveSide == MoveSide::AFTER) {
+      return this->topo_position_ > n->topo_position_;
+    }
 
-  if (this->owningBlock() == n->owningBlock()) {
-    return this->topo_position_ > n->topo_position_;
+    JIT_ASSERT(this == n);
+    return false;
   }
 
   // These nodes don't share a common block. Traverse the blockchains upward
@@ -887,7 +887,7 @@ bool Node::isAfter(const Node * n) const {
       JIT_ASSERT(rhs->owningBlock());
 
       if (lhs->owningBlock() == rhs->owningBlock()) {
-        return lhs->isAfter(rhs);
+        return lhs->isBeforeOrAfter(rhs, moveSide);
       }
       rhs = rhs->owningBlock()->owningNode();
     }
@@ -896,6 +896,15 @@ bool Node::isAfter(const Node * n) const {
   }
   // should never reach here, since both nodes are ultimately in the same graph
   JIT_ASSERT(false);
+
+}
+
+bool Node::isBefore(const Node * n) const {
+  return isBeforeOrAfter(n, MoveSide::BEFORE);
+}
+
+bool Node::isAfter(const Node * n) const {
+  return isBeforeOrAfter(n, MoveSide::AFTER);
 }
 
 Node* Node::insertBefore(Node * n) {

torch/csrc/jit/ir.h

@@ -599,6 +599,7 @@ struct Node : public Attributes<Node> {
   enum class MoveSide { BEFORE, AFTER };
   bool tryMove(Node* movePoint, MoveSide moveSide, const AliasDb& aliasDb, bool dryRun);
   void move(Node* movePoint, MoveSide moveSide);
+  bool isBeforeOrAfter(const Node* n, MoveSide moveSide) const;
 
   std::pair<Value*, const Argument&> findInput(Symbol name);
   void findSchema() const;
@@ -808,6 +809,12 @@ friend struct Block;
     const auto & block = *block_;
     return block.nodes();
   }
+  Node * param_node() {
+    return block_->param_node();
+  }
+  const Node * param_node() const {
+    return block_->param_node();
+  }
   Node * return_node() {
     return block_->return_node();
   }