Hoisting common expressions out of If blocks [retry] (pytorch#65645)

Summary:
Pull Request resolved: pytorch#65645

This is a retry of PR: pytorch#59492

Latest Changes: Added more tests, added the getOrCreateDB pattern, updated logic to remove unnecessary checks
addressed all comments.

Adding code to find common expressions from the two subblocks of an if
operation and hoist them before the if block.
This also allows Dead Code Elimination to
then eliminate some if blocks.

Test Plan: python test_jit.py TestIfHoisting

Reviewed By: eellison

Differential Revision: D33302065

Pulled By: Gamrix

fbshipit-source-id: a5a184a480cf07354359aaca344c6e27b687a3c2

Author: Gamrix

test/jit/test_if_hoisting.py

@@ -0,0 +1,213 @@
+
+import torch
+from torch.testing import FileCheck
+from torch.testing._internal.jit_utils import JitTestCase
+
+if __name__ == "__main__":
+    raise RuntimeError(
+        "This test file is not meant to be run directly, use:\n\n"
+        "\tpython test/test_jit.py TESTNAME\n\n"
+        "instead."
+    )
+
+
+class TestIfHoisting(JitTestCase):
+    def test_if_hoist_basic(self):
+        def fn(x: bool, y: int):
+            if x:
+                z = y + 3
+            else:
+                z = y + 3
+            return z
+
+
+        fn_script = torch.jit.script(fn)
+        op_graph = fn_script.graph
+        self.run_pass("common_expression_hoisting", op_graph)
+        self.run_pass("dce", op_graph)
+        FileCheck().check_count("prim::If", 0, exactly=True).run(op_graph)
+        FileCheck().check_count("aten::add", 1, exactly=True).run(op_graph)
+        self.assertEqual(fn(True, 1), fn_script(True, 1))
+
+    def test_if_hoist_transposed_expr(self):
+        """
+        Making sure that we can properly eliminate
+        an expression even if it is not at the start
+        of a block
+        """
+        def fn(x: bool, y: int):
+            if x:
+                a = y + 3
+                b = y * 2
+            else:
+                b = y * 2
+                a = y + 3
+            return a, b
+
+        fn_script = torch.jit.script(fn)
+        op_graph = fn_script.graph
+        self.run_pass("common_expression_hoisting", op_graph)
+        self.run_pass("dce", op_graph)
+
+        FileCheck().check_count("prim::If", 0, exactly=True).run(op_graph)
+        FileCheck().check_count("aten::add", 1, exactly=True).run(op_graph)
+
+        self.assertEqual(fn(True, 1), fn_script(True, 1))
+        self.assertEqual(fn(False, 5), fn_script(False, 5))
+
+    def test_if_hoist_swapped_expr(self):
+        """
+        Making sure that the if statement
+        doesn't get fully eliminated here
+        """
+        def fn(x: bool, y: int):
+            if x:
+                a = y + 3
+                b = y * 2
+            else:
+                a = y * 2
+                b = y + 3
+            return a, b
+
+        fn_script = torch.jit.script(fn)
+        op_graph = fn_script.graph
+        self.run_pass("common_expression_hoisting", op_graph)
+        self.run_pass("dce", op_graph)
+
+        FileCheck().check_count("prim::If", 1, exactly=True).run(op_graph)
+        FileCheck().check_count("aten::add", 1, exactly=True).run(op_graph)
+
+        self.assertEqual(fn(True, 1), fn_script(True, 1))
+        self.assertEqual(fn(False, 5), fn_script(False, 5))
+
+    def test_if_hoist_reused_var(self):
+        """
+        Making sure that cases where the python variable is reused
+        is handled correctly
+        """
+        def fn(x: bool, y: int):
+            b = 6
+            if x:
+                a = y + 3
+                a = y * 2
+            else:
+                a = y * 2
+                b = y + 3
+            return a, b
+
+        fn_script = torch.jit.script(fn)
+        op_graph = fn_script.graph
+        self.run_pass("common_expression_hoisting", op_graph)
+        self.run_pass("dce", op_graph)
+
+        FileCheck().check_count("prim::If", 1, exactly=True).run(op_graph)
+        FileCheck().check_count("aten::add", 1, exactly=True).run(op_graph)
+        FileCheck().check_count("aten::mul", 1, exactly=True).run(op_graph)
+
+        self.assertEqual(fn(True, 1), fn_script(True, 1))
+        self.assertEqual(fn(False, 5), fn_script(False, 5))
+
+    def test_no_hoist(self):
+        """
+        Nothing should happen here, expressions are different
+        """
+        def fn(x: bool, y: int, z: int):
+            if x:
+                a = y + 3
+            else:
+                a = z + 3
+            return a
+
+        fn_script = torch.jit.script(fn)
+        op_graph = fn_script.graph
+        self.run_pass("common_expression_hoisting", op_graph)
+        self.run_pass("dce", op_graph)
+
+        FileCheck().check_count("prim::If", 1, exactly=True).run(op_graph)
+        FileCheck().check_count("aten::add", 2, exactly=True).run(op_graph)
+
+        self.assertEqual(fn(True, 1, 3), fn_script(True, 1, 3))
+        self.assertEqual(fn(False, 5, 10), fn_script(False, 5, 10))
+
+    def test_mutate_before(self):
+        """
+        Make sure that if there is a mutation before the common
+        op, the hoist doesn't happen
+        """
+        def fn(x: bool, y: torch.Tensor):
+            if x:
+                y.add_(8)
+                a = y + 3
+            else:
+                a = y + 3
+            return a
+
+        fn_script = torch.jit.script(fn)
+        op_graph = fn_script.graph
+        self.run_pass("common_expression_hoisting", op_graph)
+        self.run_pass("dce", op_graph)
+
+        FileCheck().check_count("prim::If", 1, exactly=True).run(op_graph)
+        FileCheck().check_count("aten::add", 2, exactly=True).run(op_graph)
+        FileCheck().check_count("aten::add_", 1, exactly=True).run(op_graph)
+
+        t1 = torch.Tensor([1])
+        t2 = torch.Tensor([5, 6])
+        self.assertEqual(fn(True, t1), fn_script(True, t1))
+        self.assertEqual(fn(False, t2), fn_script(False, t2))
+
+    def test_mutate_after(self):
+        """
+        Check that the hoist can happen properly, and
+        that the output is still correct.
+        """
+        def fn(x: bool, y: torch.Tensor):
+            if x:
+                b = 1
+                a = y + 3
+                y.add_(8)
+            else:
+                b = 2
+                a = y + 3
+            c = b + a
+            return a
+
+        fn_script = torch.jit.script(fn)
+        op_graph = fn_script.graph
+        self.run_pass("common_expression_hoisting", op_graph)
+        self.run_pass("dce", op_graph)
+
+        FileCheck().check_count("prim::If", 1, exactly=True).run(op_graph)
+        FileCheck().check_count("aten::add", 2, exactly=True).run(op_graph)
+
+        t1 = torch.Tensor([1])
+        t2 = torch.Tensor([5, 6])
+        self.assertEqual(fn(True, t1.clone()), fn_script(True, t1.clone()))
+        self.assertEqual(fn(False, t2.clone()), fn_script(False, t2.clone()))
+
+    def test_multiple_hoists(self):
+        """
+        test that hoists that depend on other hoists are done correctly
+        """
+        def fn(x: bool, y: torch.Tensor):
+            if x:
+                a = y + 3
+                b = a + y
+            else:
+                a = y + 3
+                b = a + y
+            c = b * 2
+            return c
+
+        fn_script = torch.jit.script(fn)
+        op_graph = fn_script.graph
+        self.run_pass("common_expression_hoisting", op_graph)
+        self.run_pass("dce", op_graph)
+
+        FileCheck().check_count("prim::If", 0, exactly=True).run(op_graph)
+        FileCheck().check_count("aten::add", 2, exactly=True).run(op_graph)
+
+        t1 = torch.Tensor([1])
+        t2 = torch.Tensor([5, 6])
+        self.assertEqual(fn(True, t1), fn_script(True, t1))
+        self.assertEqual(fn(False, t2), fn_script(False, t2))

test/quantization/jit/test_quantize_jit.py

@@ -1256,6 +1256,7 @@ class Res(torch.nn.Module):
             def __init__(self):
                 super(Res, self).__init__()
                 self.conv = torch.nn.Conv2d(3, 3, 1).float()
+                self.conv2 = torch.nn.Conv2d(3, 3, 1).float()
                 self.use_skip = True
 
             def forward(self, x: torch.Tensor, cond: bool) -> torch.Tensor:
@@ -1264,7 +1265,7 @@ def forward(self, x: torch.Tensor, cond: bool) -> torch.Tensor:
                 if self.use_skip:
                     return self.conv(x)
                 else:
-                    return self.conv(x)
+                    return self.conv2(x)
 
         class M(torch.nn.Module):
             def __init__(self):

test/test_jit.py

@@ -25,6 +25,7 @@
 from jit.test_builtins import TestBuiltins, TestTensorBuiltins  # noqa: F401
 from jit.test_ignore_context_manager import TestIgnoreContextManager  # noqa: F401
 from jit.test_symbolic_shape_analysis import TestSymbolicShapeAnalysis  # noqa: F401
+from jit.test_if_hoisting import TestIfHoisting  # noqa: F401
 from jit.test_unsupported_ops import TestUnsupportedOps  # noqa: F401
 from jit.test_freezing import TestFreezing, TestFrozenOptimizations, TestMKLDNNReinplacing  # noqa: F401
 from jit.test_peephole import TestPeephole  # noqa: F401

tools/build_variables.bzl

@@ -219,6 +219,7 @@ core_sources_full_mobile_no_backend_interface = [
     "torch/csrc/jit/passes/clear_profiling.cpp",
     "torch/csrc/jit/passes/clear_undefinedness.cpp",
     "torch/csrc/jit/passes/common_subexpression_elimination.cpp",
+    "torch/csrc/jit/passes/common_expression_hoisting.cpp",
     "torch/csrc/jit/passes/concat_opt.cpp",
     "torch/csrc/jit/passes/constant_pooling.cpp",
     "torch/csrc/jit/passes/constant_propagation.cpp",

torch/_C/__init__.pyi.in

@@ -211,6 +211,7 @@ def _jit_pass_inline(Graph) -> None: ...
 def _jit_pass_constant_propagation(Graph) -> None: ...
 def _jit_pass_propagate_shapes_on_graph(Graph) -> None: ...
 def _jit_erase_non_input_shape_information(Graph) -> None: ...
+def _jit_pass_common_expression_hoisting(Graph) -> None: ...
 def _jit_get_schemas_for_operator(name :str) -> List[FunctionSchema]: ...
 def _jit_check_alias_annotation(g: Graph, args: Tuple[Any, ...], unqualified_op_name: str): ...
 def _jit_can_fuse_on_cpu() -> _bool: ...

torch/csrc/jit/ir/node_hashing.cpp

@@ -208,6 +208,8 @@ bool attributesEqualCSE(const Node* lhs, const Node* rhs) {
 
 } // anonymous namespace
 
+// Makes a hash that hashes the input Value, the output type
+// as well as the node attributes
 size_t HashNode::operator()(const Node* k) const {
   AT_ASSERT(k != nullptr);
   size_t constant_hash = 0;
@@ -235,6 +237,8 @@ size_t HashNode::operator()(const Node* k) const {
       constant_hash);
 };
 
+// Checks that two nodes have the same inputs, output types
+// and node attributes.
 bool EqualNode::operator()(const Node* lhs, const Node* rhs) const {
   if (lhs == nullptr && rhs == nullptr)
     return true;
@@ -267,6 +271,16 @@ bool EqualNode::operator()(const Node* lhs, const Node* rhs) const {
   if (!attributesEqualCSE(lhs, rhs))
     return false;
 
+  // Check if the blocks contained in a op are the same
+  if (lhs->blocks().size() != rhs->blocks().size()) {
+    return false;
+  }
+  for (size_t i = 0; i < lhs->blocks().size(); ++i) {
+    if (lhs->blocks()[i] != rhs->blocks()[i]) {
+      return false;
+    }
+  }
+
   return true;
 };