Remove unused exported tensor constants after Cortex-M lowering

backends/cortex_m/passes/BUCK

@@ -57,6 +57,7 @@ fbcode_target(_kind = runtime.python_library,
         "//executorch/backends/transforms:convert_conv1d_to_conv2d_pass",
         "//executorch/backends/transforms:remove_getitem_op",
         "//executorch/backends/transforms:remove_permutes_around_elementwise_ops",
+        "//executorch/backends/transforms:remove_unused_constants_pass",
         "//executorch/backends/transforms:replace_scalar_with_tensor",
         "//executorch/backends/transforms:replace_ops_with_channels_last_variants",
         "//executorch/backends/transforms:replace_squeeze_unsqueeze_with_view",

backends/cortex_m/passes/cortex_m_pass_manager.py

@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 
+import copy
 import inspect
 from typing import Any, Optional, Type
 
@@ -20,16 +21,24 @@
 from executorch.backends.transforms.remove_permutes_around_elementwise_ops import (
     RemovePermutesAroundElementwiseOps,
 )
+from executorch.backends.transforms.remove_unused_constants_pass import (
+    RemoveUnusedConstantsPass,
+)
 from executorch.backends.transforms.replace_scalar_with_tensor import (
     ReplaceScalarWithTensorArgPass,
 )
 from executorch.backends.transforms.replace_squeeze_unsqueeze_with_view import (
     ReplaceSqueezeAndUnsqueezeWithViewPass,
 )
-from executorch.exir.pass_base import ExportPass
+from executorch.exir.pass_base import (
+    ExportedProgramPassBase,
+    ExportedProgramPassResult,
+    ExportPass,
+)
 from executorch.exir.pass_manager import PassManager
 from executorch.exir.program._program import _transform, lift_constant_tensor_pass
 from torch.export import ExportedProgram
+from torch.fx import GraphModule
 
 from .activation_fusion_pass import ActivationFusionPass
 from .aten_to_cortex_m_pass import AtenToCortexMPass
@@ -47,7 +56,29 @@
 from .quantized_clamp_activation_pass import QuantizedClampActivationPass
 from .replace_quant_nodes_pass import ReplaceQuantNodesPass
 
-PassClass = Type[ExportPass]
+PassClass = Type[ExportPass | ExportedProgramPassBase]
+
+
+class LiftConstantTensorsPass(ExportedProgramPassBase):
+    def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
+        # The pass manager shallow-copies programs; lifting mutates shared structures.
+        graph = copy.deepcopy(exported_program.graph)
+        for original, cloned in zip(exported_program.graph.nodes, graph.nodes):
+            cloned.name = original.name
+        graph_module = GraphModule(exported_program.graph_module, graph)
+        graph_module.meta = exported_program.graph_module.meta.copy()
+        exported_program._graph_module = graph_module
+        exported_program._graph_signature = copy.deepcopy(
+            exported_program.graph_signature
+        )
+        exported_program._state_dict = exported_program.state_dict.copy()
+
+        buffer_count = len(exported_program.graph_signature.buffers)
+        exported_program = lift_constant_tensor_pass(exported_program)
+        return ExportedProgramPassResult(
+            exported_program,
+            len(exported_program.graph_signature.buffers) != buffer_count,
+        )
 
 
 class CortexMPassManager(PassManager):
@@ -63,6 +94,8 @@ class CortexMPassManager(PassManager):
         AtenToCortexMPass,
         FuseConvPaddingPass,
         InitializeScratchBuffersPass,
+        LiftConstantTensorsPass,
+        RemoveUnusedConstantsPass,
     ]
 
     explicit_layout_pass_list: list[PassClass] = [
@@ -83,11 +116,13 @@ class CortexMPassManager(PassManager):
         AtenToCortexMPass,
         FuseConvPaddingPass,
         InitializeScratchBuffersPass,
+        LiftConstantTensorsPass,
+        RemoveUnusedConstantsPass,
     ]
 
     pass_list = legacy_pass_list
 
-    pass_list_transform_for_annotation: list[PassClass] = [
+    pass_list_transform_for_annotation: list[Type[ExportPass]] = [
         ScalarsToAttributePass,
         ReplaceScalarWithTensorArgPass,
         ClampHardswishPass,
@@ -165,7 +200,4 @@ def transform(self) -> ExportedProgram:
             transform_pass = pass_cls(**kwargs)
             exported_program = _transform(exported_program, transform_pass)
 
-        # All constant tensors should be lifted to buffers at this point, re-run
-        # lift_constant_tensor_pass in case new ones have been introduced.
-        exported_program = lift_constant_tensor_pass(exported_program)
         return exported_program

backends/cortex_m/test/targets.bzl

@@ -81,6 +81,8 @@ def define_common_targets(is_fbcode = False):
                 "//executorch/backends/cortex_m/passes:cortex_passes",
                 "//executorch/backends/cortex_m/quantizer:quantizer",
                 "//executorch/backends/test/harness:tester",
+                "//executorch/backends/transforms:remove_unused_constants_pass",
+                "//executorch/exir:lib",
                 "//executorch/exir/dialects:lib",
                 ":tester",
                 "fbsource//third-party/pypi/pytest:pytest",

backends/cortex_m/test/test_explicit_layout_pipeline.py

@@ -4,15 +4,23 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import copy
 from functools import partial
 
 import pytest
 import torch
-from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager
+from executorch.backends.cortex_m.passes.cortex_m_pass_manager import (
+    CortexMPassManager,
+    LiftConstantTensorsPass,
+)
 from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer
 from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig
 from executorch.backends.cortex_m.test.tester import CortexMTester
 from executorch.backends.test.harness.stages import Quantize, RunPasses, StageType
+from executorch.backends.transforms.remove_unused_constants_pass import (
+    RemoveUnusedConstantsPass,
+)
+from executorch.exir import to_edge
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.fx import Node
 
@@ -107,6 +115,11 @@ def test_layout_pipelines_select_distinct_spatial_operators():
     )
     explicit_program = explicit.get_artifact(StageType.RUN_PASSES).exported_program()
 
+    for program in (legacy_program, explicit_program):
+        assert all(
+            node.users for node in program.graph.nodes if node.op == "placeholder"
+        )
+
     assert _count(legacy_program, exir_ops.edge.cortex_m.quantized_conv2d.default) == 1
     assert (
         _count(
@@ -128,6 +141,99 @@ def test_layout_pipelines_select_distinct_spatial_operators():
     assert _count(explicit_program, exir_ops.edge.cortex_m.transpose.default) == 2
 
 
+@pytest.mark.parametrize("use_edge_transform", [False, True], ids=["cortex_m", "edge"])
+@pytest.mark.parametrize(
+    "passes,lifted,pruned",
+    [
+        pytest.param([], False, False, id="empty"),
+        pytest.param([LiftConstantTensorsPass], True, False, id="lift"),
+        pytest.param(
+            [LiftConstantTensorsPass, RemoveUnusedConstantsPass],
+            True,
+            True,
+            id="lift_and_prune",
+        ),
+    ],
+)
+def test_constant_cleanup_respects_pass_list_and_lift_order(
+    passes, lifted, pruned, use_edge_transform
+):
+    class Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.register_buffer("_lifted_tensor_constant0", torch.tensor([1.0]))
+            self.register_buffer("_lifted_tensor_constant1", torch.tensor([2.0]))
+
+        def forward(self, input):
+            return input + self._lifted_tensor_constant1
+
+    class OtherEntryPoint(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.register_buffer("_lifted_tensor_constant2", torch.tensor([7.0]))
+
+        def forward(self, x):
+            return x + self._lifted_tensor_constant2
+
+    inputs = (torch.zeros(1),)
+    edge = to_edge(torch.export.export(Model(), inputs))
+    program = edge.exported_program()
+    assert "_lifted_tensor_constant0" in program.graph_signature.buffers
+    program.graph_module.register_buffer("new_tensor", torch.tensor([3.0]))
+    output = next(node for node in program.graph.nodes if node.op == "output")
+    original = output.args[0][0]
+    with program.graph.inserting_before(original):
+        constant = program.graph.get_attr("new_tensor")
+        constant.meta = original.meta.copy()
+        result = program.graph.call_function(
+            exir_ops.edge.aten.add.Tensor, (original.args[1], constant)
+        )
+        result.meta = original.meta.copy()
+    original.replace_input_with(original.args[1], result)
+    program.graph_module.recompile()
+    program.validate()
+
+    other = to_edge(torch.export.export(OtherEntryPoint(), inputs)).exported_program()
+    program.state_dict.update(other.state_dict)
+    other._state_dict = program.state_dict
+    other.validate()
+    original_code = program.graph_module.code
+    original_signature = copy.deepcopy(program.graph_signature)
+    original_state = program.state_dict.copy()
+    program.graph_module.meta["constant_lifting_test"] = "preserved"
+
+    transformed = (
+        edge.transform([pass_cls() for pass_cls in passes]).exported_program()
+        if use_edge_transform
+        else CortexMPassManager(program, passes=passes).transform()
+    )
+
+    program.validate()
+    other.validate()
+    torch.testing.assert_close(program.module()(*inputs), torch.tensor([5.0]))
+    torch.testing.assert_close(other.module()(*inputs), torch.tensor([7.0]))
+    assert program.graph_module.code == original_code
+    assert program.graph_signature == original_signature
+    assert program.state_dict.keys() == original_state.keys()
+    for name, tensor in original_state.items():
+        assert program.state_dict[name] is tensor
+        assert other.state_dict[name] is tensor
+    transformed.validate()
+    assert transformed.graph_signature.user_inputs == original_signature.user_inputs
+    assert transformed.graph_module.meta["constant_lifting_test"] == "preserved"
+    assert ("_lifted_tensor_constant0" in transformed.graph_signature.buffers) == (
+        not pruned
+    )
+    assert sum(node.op == "get_attr" for node in transformed.graph.nodes) == (
+        not lifted
+    )
+    assert (
+        transformed.state_dict["_lifted_tensor_constant1"]
+        is original_state["_lifted_tensor_constant1"]
+    )
+    torch.testing.assert_close(transformed.module()(*inputs), torch.tensor([5.0]))
+
+
 def test_conv1d_is_quantized_before_layout_conversion():
     tester = CortexMTester(Conv1d().eval(), (torch.randn(1, 2, 8),))
     tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=True)))

backends/transforms/remove_unused_constants_pass.py

@@ -0,0 +1,79 @@
+# 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 copy
+
+from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult
+from torch.export import ExportedProgram
+from torch.export.graph_signature import ExportGraphSignature, InputKind, TensorArgument
+from torch.fx import GraphModule
+
+
+class RemoveUnusedConstantsPass(ExportedProgramPassBase):
+    """Retire unused tensor constants after backend graph rewrites."""
+
+    def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
+        signature = exported_program.graph_signature
+        placeholders = {
+            node.name: node
+            for node in exported_program.graph.nodes
+            if node.op == "placeholder"
+        }
+        protected_targets = {
+            spec.target for spec in signature.output_specs if spec.target is not None
+        }
+        preserved_names = {
+            argument.name
+            for entry in exported_program.module_call_graph
+            if entry.signature is not None
+            for argument in (*entry.signature.inputs, *entry.signature.outputs)
+            if isinstance(argument, TensorArgument)
+        }
+        unused = [
+            spec
+            for spec in signature.input_specs
+            if spec.kind
+            in (InputKind.PARAMETER, InputKind.BUFFER, InputKind.CONSTANT_TENSOR)
+            and spec.target not in protected_targets
+            and spec.arg.name not in preserved_names
+            and not placeholders[spec.arg.name].users
+        ]
+        if not unused:
+            return ExportedProgramPassResult(exported_program, False)
+
+        unused_names = {spec.arg.name for spec in unused}
+        signature = ExportGraphSignature(
+            input_specs=[
+                spec
+                for spec in signature.input_specs
+                if spec.arg.name not in unused_names
+            ],
+            output_specs=list(signature.output_specs),
+        )
+        # The pass manager shallow-copies programs, so their graph is still shared.
+        graph = copy.deepcopy(exported_program.graph)
+        for original_node, node in zip(exported_program.graph.nodes, list(graph.nodes)):
+            # FX copying can rename built-ins such as "input", used by the signature.
+            node.name = original_node.name
+            if node.op == "placeholder" and node.name in unused_names:
+                graph.erase_node(node)
+        graph_module = GraphModule(exported_program.graph_module, graph)
+        graph_module.meta = exported_program.graph_module.meta.copy()
+
+        remaining_targets = {spec.target for spec in signature.input_specs}
+        # Entry points can share state dictionaries; retain their tensor identities.
+        state_dict = exported_program.state_dict.copy()
+        constants = exported_program.constants.copy()
+        for spec in unused:
+            if spec.target not in remaining_targets:
+                state_dict.pop(spec.target, None)
+                constants.pop(spec.target, None)
+
+        exported_program._state_dict = state_dict
+        exported_program._constants = constants
+        exported_program._graph_signature = signature
+        exported_program._graph_module = graph_module
+        return ExportedProgramPassResult(exported_program, True)

backends/transforms/targets.bzl

@@ -214,6 +214,16 @@ def define_common_targets():
         ],
     )
 
+    runtime.python_library(
+        name = "remove_unused_constants_pass",
+        srcs = ["remove_unused_constants_pass.py"],
+        visibility = ["PUBLIC"],
+        deps = [
+            "//caffe2:torch",
+            "//executorch/exir:pass_base",
+        ],
+    )
+
     runtime.python_library(
         name = "mean_to_sum_div",
         srcs = ["mean_to_sum_div.py"],