Add constant prop pass (#1146)

Summary:
Pull Request resolved: #1146

Add a const prop pass for exported program

Reviewed By: angelayi

Differential Revision: D50961161

fbshipit-source-id: 5f4fc9e45c063d7697a81be4469b4c675ffac80d

Author: cccclai

exir/passes/TARGETS

@@ -77,6 +77,16 @@ python_library(
     ],
 )
 
+python_library(
+    name = "constant_prop_pass",
+    srcs = [
+        "constant_prop_pass.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+    ],
+)
+
 python_library(
     name = "remove_assert_async_pass",
     srcs = [

exir/passes/constant_prop_pass.py

@@ -0,0 +1,132 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from torch._export import ExportedProgram
+from torch._export.utils import get_buffer, get_param, is_buffer, is_param
+from torch._guards import detect_fake_mode
+from torch.export.exported_program import InputKind, InputSpec, TensorArgument
+
+
+def is_const(arg, exported_program, const_data_list) -> bool:
+    if isinstance(arg, (tuple, list)):
+        return all(is_const(x, exported_program, const_data_list) for x in arg)
+    elif isinstance(arg, dict):
+        return all(is_const(x, exported_program, const_data_list) for x in arg.values())
+    elif not isinstance(arg, torch.fx.Node) or arg.op != "placeholder":
+        return False
+    elif (
+        is_param(exported_program, arg)
+        or is_buffer(exported_program, arg)
+        or arg.name in const_data_list
+    ):
+        return True
+    return False
+
+
+def get_data(exported_program, arg):
+    if isinstance(arg, (tuple, list)):
+        return [get_data(exported_program, x) for x in arg]
+    elif is_param(exported_program, arg):
+        return get_param(exported_program, arg)
+    elif is_buffer(exported_program, arg):
+        return get_buffer(exported_program, arg)
+    return None
+
+
+def constant_prop_pass(exported_program: ExportedProgram) -> ExportedProgram:
+    """
+    This pass is for constant propagation for Exported Program with lifted parameters,
+    as the parameters will not be shown up as `get_attr` but as `placeholder` to the graph.
+    """
+    if (
+        len([node for node in exported_program.graph.nodes if node.op == "placeholder"])
+        == 0
+    ):
+        return exported_program
+
+    has_cond = [
+        node
+        for node in exported_program.graph.nodes
+        if node.target == torch.ops.higher_order.cond
+    ]
+    if len(has_cond) > 0:
+        raise RuntimeError("constant_prop_pass for control flow is not supported yet.")
+
+    first_user_input = None
+    for node in exported_program.graph.nodes:
+        if (
+            node.op == "placeholder"
+            and node.name in exported_program.graph_signature.user_inputs
+        ):
+            first_user_input = node
+            break
+
+    buffers = exported_program.graph_signature.buffers
+    prop_constant_data = []
+    const_data_to_be_removed = set()
+
+    fake_mode = detect_fake_mode(
+        tuple(
+            node.meta["val"]
+            for node in exported_program.graph.nodes
+            if node.op == "placeholder"
+        )
+    )
+    assert fake_mode is not None
+
+    for node in exported_program.graph.nodes:
+        if node.op == "call_function":
+            constant_data_name_list = [
+                input_spec.target for input_spec in prop_constant_data
+            ]
+            if is_const(node.args, exported_program, constant_data_name_list):
+                args_data = [get_data(exported_program, arg) for arg in node.args]
+                kwargs_data = node.kwargs
+                const_data_to_be_removed.update(node.args)
+                prop_constant_tensor = node.target(*args_data, **kwargs_data)
+                prop_constant_tensor_fqn = f"_prop_tensor_constant{len(buffers)}"
+
+                with exported_program.graph.inserting_before(first_user_input):
+                    const_placeholder_node = exported_program.graph.placeholder(
+                        prop_constant_tensor_fqn
+                    )
+                    # Update the meta data of the new placeholder (buffer) node
+                    for k, v in node.meta.items():
+                        const_placeholder_node.meta[k] = v
+                    const_placeholder_node.meta["val"] = fake_mode.from_tensor(
+                        prop_constant_tensor, static_shapes=True
+                    )
+                    const_placeholder_node.meta["val"].constant = prop_constant_tensor
+
+                    node.replace_all_uses_with(const_placeholder_node)
+                    exported_program.graph.erase_node(node)
+                    prop_constant_node_input_spec = InputSpec(
+                        kind=InputKind.BUFFER,
+                        arg=TensorArgument(name=const_placeholder_node.name),
+                        target=prop_constant_tensor_fqn,
+                    )
+                    prop_constant_data.append(prop_constant_node_input_spec)
+                    buffers.append(prop_constant_tensor_fqn)
+                    exported_program.state_dict[
+                        prop_constant_tensor_fqn
+                    ] = prop_constant_tensor
+                    exported_program.graph_signature.input_specs.append(
+                        prop_constant_node_input_spec
+                    )
+
+    # Remove the propogated buffer from the state dict
+    for node in exported_program.graph.nodes:
+        if (
+            node.op == "placeholder"
+            and node in const_data_to_be_removed
+            and len(node.users) == 0
+        ):
+            exported_program.state_dict.pop(node.name, None)
+            exported_program.graph.erase_node(node)
+
+    exported_program.graph_module.recompile()
+    return exported_program

exir/tests/TARGETS

@@ -212,6 +212,7 @@ python_unittest(
         "//executorch/exir/dialects/edge:lib",
         "//executorch/exir/emit:lib",
         "//executorch/exir/passes:const_prop_pass",
+        "//executorch/exir/passes:constant_prop_pass",
         "//executorch/exir/passes:debug_handle_generator_pass",
         "//executorch/exir/passes:lib",
         "//executorch/exir/passes:remove_assert_async_pass",

exir/tests/test_passes.py

@@ -32,6 +32,7 @@
     ToOutVarPass,
 )
 from executorch.exir.passes.const_prop_pass import ConstPropPass
+from executorch.exir.passes.constant_prop_pass import constant_prop_pass
 from executorch.exir.passes.debug_handle_generator_pass import DebugHandleGeneratorPass
 from executorch.exir.passes.remove_assert_async_pass import RemoveAssertAsyncPass
 from executorch.exir.passes.remove_mixed_type_operators import RemoveMixedTypeOperators
@@ -1057,3 +1058,84 @@ def forward(self, x):
         FileCheck().check(
             "executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor"
         ).run(gm.code)
+
+    def test_constant_prop_pass_for_add(self) -> None:
+        def add(x: torch.Tensor) -> torch.Tensor:
+            return x + 3
+
+        edge = exir.capture(add, (torch.ones(1),), exir.CaptureConfig(enable_aot=True))
+        edge = edge.transform(ScalarToTensorPass(), RemoveMixedTypeOperators())
+        edge.exported_program = lift_constant_tensor_pass(edge.exported_program)
+
+        # Check there is a lifted tensor followed by a to_copy node
+        FileCheck().check("_lifted_tensor_constant0").check(
+            "torch.ops.aten._to_copy.default"
+        ).run(edge.exported_program.graph_module.code)
+
+        new_ep = constant_prop_pass(edge.exported_program)
+
+        # Check (_lifted_tensor_constant + to_copy) node is replaced by prop tensor
+        FileCheck().check_not("_lifted_tensor_constant").check(
+            "_prop_tensor_constant1"
+        ).check_not("torch.ops.aten._to_copy.default").run(new_ep.graph_module.code)
+
+    def test_constant_prop_pass_for_parameter(self) -> None:
+        def count_additions(gm: torch.fx.GraphModule) -> int:
+            return sum(
+                (node.target == torch.ops.aten.add.Tensor) for node in gm.graph.nodes
+            )
+
+        class M(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.a = torch.nn.Parameter(torch.ones(1, 2, 3))
+
+            def forward(self, x):
+                b = self.a + self.a
+                c = torch.cat([self.a, b])
+                return (c + c) + x
+
+        edge = exir.capture(
+            M(),
+            (torch.zeros(2, 2, 3),),
+            exir.CaptureConfig(enable_aot=True),
+        )
+        self.assertEqual(count_additions(edge.exported_program.graph_module), 3)
+        edge.exported_program = constant_prop_pass(edge.exported_program)
+        self.assertEqual(count_additions(edge.exported_program.graph_module), 1)
+
+    def test_constant_prop_pass_for_control_flow(self) -> None:
+        class Module(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+
+            def t(self, val):
+                return val + 1
+
+            def f(self, val):
+                return val - 1
+
+            def true_fn(self, val):
+                return self.linear(val) + self.t(val)
+
+            def false_fn(self, val):
+                return self.linear(val) - self.f(val)
+
+            def forward(self, pred, x):
+                return torch.ops.higher_order.cond(
+                    pred, self.true_fn, self.false_fn, [x]
+                )
+
+        mod = Module()
+        x = torch.randn([3, 3])
+        pred = torch.tensor(x[0][0].item() < 0)
+        edge = exir.capture(mod, (pred, x), config=exir.CaptureConfig(enable_aot=True))
+        error_msg = r"constant_prop_pass for control flow is not supported yet."
+
+        # TODO(chenlai): enable constant prop pass for control flow
+        with self.assertRaisesRegex(
+            RuntimeError,
+            error_msg,
+        ):
+            _ = constant_prop_pass(edge.exported_program)