Remove unused exported tensor constants after Cortex-M lowering

backends/cortex_m/passes/BUCK

@@ -58,6 +58,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,6 +21,9 @@
 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,
 )
@@ -34,6 +38,7 @@
 from executorch.exir.pass_manager import ExportedProgramPassManager, PassType
 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
@@ -52,7 +57,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 _CortexMLoweringPass(ExportedProgramPassBase):
@@ -77,10 +104,6 @@ def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
             modified |= transformed is not exported_program
             exported_program = transformed
 
-        # Passes can introduce tensor attributes that must become program inputs.
-        buffer_count = len(exported_program.graph_signature.buffers)
-        exported_program = lift_constant_tensor_pass(exported_program)
-        modified |= len(exported_program.graph_signature.buffers) != buffer_count
         return ExportedProgramPassResult(exported_program, modified)
 
 
@@ -97,6 +120,8 @@ class CortexMPassManager(ExportedProgramPassManager):
         AtenToCortexMPass,
         FuseConvPaddingPass,
         InitializeScratchBuffersPass,
+        LiftConstantTensorsPass,
+        RemoveUnusedConstantsPass,
     ]
 
     explicit_layout_pass_list: list[PassClass] = [
@@ -117,11 +142,13 @@ class CortexMPassManager(ExportedProgramPassManager):
         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]] = [
         DecomposeSDPAPass,
         ScalarsToAttributePass,
         ReplaceScalarWithTensorArgPass,

backends/cortex_m/test/targets.bzl

@@ -82,6 +82,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

@@ -5,8 +5,14 @@
 # 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
+
 import pytest
 import torch
+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 (
@@ -15,6 +21,10 @@
     CortexMTester,
 )
 from executorch.backends.test.harness.stages import Quantize, 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
 
@@ -156,6 +166,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(
@@ -177,6 +192,104 @@ def test_layout_pipelines_select_distinct_spatial_operators():
     assert _count(explicit_program, exir_ops.edge.cortex_m.transpose.default) == 2
 
 
+@pytest.mark.parametrize("entry_point", ["legacy", "edge_passes", "edge_manager"])
+@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, entry_point
+):
+    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"
+
+    if entry_point == "edge_passes":
+        transformed = edge.transform(
+            [pass_cls() for pass_cls in passes]
+        ).exported_program()
+    elif entry_point == "edge_manager":
+        transformed = edge.transform(
+            CortexMPassManager(passes=passes)
+        ).exported_program()
+    else:
+        transformed = 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/cortex_m/test/test_pass_manager.py

@@ -142,7 +142,7 @@ def test_empty_pass_list_reports_no_change():
         manager.transform()
 
 
-def test_empty_pass_list_still_lifts_constants():
+def test_empty_pass_list_preserves_unlifted_constants():
     class AddConstant(ExportPass):
         def call_operator(self, op, args, kwargs, meta):
             result = super().call_operator(op, args, kwargs, meta)
@@ -159,9 +159,10 @@ def call_operator(self, op, args, kwargs, meta):
     assert any(node.op == "get_attr" for node in program.graph.nodes)
 
     result = CortexMPassManager(passes=[])(program)
-    assert result.modified
-    lowered = result.exported_program
-    lowered.validate()
-    assert all(node.op != "get_attr" for node in lowered.graph.nodes)
-    assert len(lowered.graph_signature.buffers) == 1
-    torch.testing.assert_close(lowered.module()(*inputs), torch.tensor([2.0, 4.0]))
+    assert not result.modified
+    assert result.exported_program.graph is program.graph
+    assert result.exported_program.graph_signature is program.graph_signature
+    program.validate()
+    assert any(node.op == "get_attr" for node in program.graph.nodes)
+    assert not program.graph_signature.buffers
+    torch.testing.assert_close(program.module()(*inputs), torch.tensor([2.0, 4.0]))

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"],