Reland unused exported constant cleanup after Cortex-M lowering

### Summary

Cortex-M layout conversion leaves original tensor buffers in the exported program. Reland pytorch#22862 to lift new constants and remove unused placeholders and storage after lowering. This restores the exact eight-file patch from the original landing, applied cleanly to current main, including preservation of the original program, shared state, and custom pass-list behavior.

The original landing was reverted because a downstream custom adapter invoked ExportedProgram-aware passes on a GraphModule. The adapter migration is tracked in [D123765406](https://www.internalfb.com/diff/D123765406) and uses the native CortexMPassManager(...).passes API. Keep this reland in draft until that prerequisite lands.

Authored with AI assistance from OpenAI Codex.

### Test plan

35 Python tests passed:

```bash
pytest -q backends/transforms/test/test_remove_unused_constants_pass.py \
  backends/cortex_m/test/test_pass_manager.py \
  backends/cortex_m/test/test_explicit_layout_pipeline.py \
  -k 'not test_implementation_explicit_layout_fuses_same_padding and not test_implementation_transpose_conv2d_strided_pointwise'
```

This covers constant pruning, original-program validity, shared state, mutation and gradient preservation, both layouts, native and legacy export entry points, PTE serialization, and custom pass-list ordering. Six FVP implementation cases were deselected.

Lintrunner passed on all eight changed files. The reviewed downstream recipe passed twelve compatibility cases with this cleanup. Four additional ModAI integration tests passed, covering ten primary exports on the internal baseline where the cleanup remains reverted.

Differential Revision: D124196546

Pull Request resolved: pytorch#23599

Author: rascani

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