diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index 2ae348a57e0..fbd03f3dcf7 100644 --- a/backends/cortex_m/passes/BUCK +++ b/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", diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index 1e007c32cc9..51cf9ddddf0 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/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 diff --git a/backends/cortex_m/test/targets.bzl b/backends/cortex_m/test/targets.bzl index e6a8a8ce252..4c800e818ec 100644 --- a/backends/cortex_m/test/targets.bzl +++ b/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", diff --git a/backends/cortex_m/test/test_explicit_layout_pipeline.py b/backends/cortex_m/test/test_explicit_layout_pipeline.py index ce01ceba8b9..969b7c33768 100644 --- a/backends/cortex_m/test/test_explicit_layout_pipeline.py +++ b/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))) diff --git a/backends/transforms/remove_unused_constants_pass.py b/backends/transforms/remove_unused_constants_pass.py new file mode 100644 index 00000000000..ef35dde09ab --- /dev/null +++ b/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) diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index 975b9ba600e..fefc9ceb583 100644 --- a/backends/transforms/targets.bzl +++ b/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"], diff --git a/backends/transforms/test/test_remove_unused_constants_pass.py b/backends/transforms/test/test_remove_unused_constants_pass.py new file mode 100644 index 00000000000..97cfb7fe307 --- /dev/null +++ b/backends/transforms/test/test_remove_unused_constants_pass.py @@ -0,0 +1,180 @@ +# 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 + +import torch +from executorch.backends.transforms.remove_unused_constants_pass import ( + RemoveUnusedConstantsPass, +) +from executorch.exir import to_edge +from torch.export.experimental import _export_forward_backward +from torch.export.graph_signature import InputKind, InputSpec, TensorArgument + + +class TensorStateModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(4, 4)) + self.register_buffer("persistent", torch.full((4, 4), 2.0)) + self.register_buffer("temporary", torch.full((4, 4), 3.0), persistent=False) + self.constant = torch.full((4, 4), 4.0) + self.register_buffer("updated", torch.zeros(4, 4)) + self.register_buffer("kept", torch.tensor(1.0)) + + def forward(self, x, unused_input): + self.updated.copy_(x) + return ( + x + + self.weight + + self.persistent + + self.temporary + + self.constant + + self.kept + ) + + +def test_removes_unused_state_and_preserves_mutation(): + inputs = (torch.randn(4, 4), torch.randn(1)) + ep = torch.export.export(TensorStateModel(), inputs).run_decompositions() + nodes = {node.name: node for node in ep.graph.nodes} + removed_targets = set() + for spec in ep.graph_signature.input_specs: + if spec.kind in ( + InputKind.PARAMETER, + InputKind.BUFFER, + InputKind.CONSTANT_TENSOR, + ) and spec.target not in ("updated", "kept"): + nodes[spec.arg.name].replace_all_uses_with(nodes["x"]) + removed_targets.add(spec.target) + ep.graph_module.recompile() + assert removed_targets == {"weight", "persistent", "temporary", "constant"} + assert not nodes["b_updated"].users + reference = copy.deepcopy(ep).module() + shared_state = ep.state_dict + shared_constants = ep.constants + + result = RemoveUnusedConstantsPass()(ep) + assert result.modified + ep.validate() + assert set(ep.state_dict) == {"updated", "kept"} + assert not ep.constants + assert ep.graph_signature.user_inputs == ("x", "unused_input") + assert "updated" in ep.graph_signature.buffers_to_mutate.values() + assert set(shared_state) == {"weight", "persistent", "updated", "kept"} + assert set(shared_constants) == {"temporary", "constant"} + assert ep.state_dict["updated"] is shared_state["updated"] + assert ep.state_dict["kept"] is shared_state["kept"] + actual = ep.module() + for _ in range(2): + inputs = (torch.randn(4, 4), torch.randn(1)) + torch.testing.assert_close(actual(*inputs), reference(*inputs)) + torch.testing.assert_close(actual.updated, inputs[0]) + assert not RemoveUnusedConstantsPass()(ep).modified + + +def test_edge_transform_preserves_original_program(): + inputs = (torch.randn(4, 4), torch.randn(1)) + edge = to_edge(torch.export.export(TensorStateModel(), inputs)) + original = edge.exported_program() + weight = next(node for node in original.graph.nodes if node.name == "p_weight") + x = next(node for node in original.graph.nodes if node.name == "x") + weight.replace_all_uses_with(x) + original.graph_module.recompile() + original.validate() + expected = original.module()(*inputs) + original_graph = original.graph + original_signature = original.graph_signature + original_state = original.state_dict + original_constants = original.constants + original.graph_module.meta["constant_pruning_test"] = "preserved" + + transformed = edge.transform([RemoveUnusedConstantsPass()]).exported_program() + + original.validate() + transformed.validate() + assert original.graph is original_graph + assert original.graph.owning_module is original.graph_module + assert original.graph_signature is original_signature + assert original.state_dict is original_state + assert original.constants is original_constants + assert original.graph_signature.parameters == ("weight",) + assert "weight" in original.state_dict + assert transformed.graph is not original.graph + assert transformed.graph_signature.parameters == () + assert "weight" not in transformed.state_dict + assert transformed.graph_module.meta["constant_pruning_test"] == "preserved" + for name, tensor in transformed.state_dict.items(): + assert tensor is original.state_dict[name] + for name, tensor in transformed.constants.items(): + assert tensor is original.constants[name] + torch.testing.assert_close(original.module()(*inputs), expected) + torch.testing.assert_close(transformed.module()(*inputs), expected) + + +def test_preserves_parameter_with_gradient_output(): + class Loss(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(2, 2)) + + def forward(self, x): + return (x * self.weight).sum() + + ep = _export_forward_backward(torch.export.export(Loss(), (torch.ones(2, 2),))) + nodes = {node.name: node for node in ep.graph.nodes} + nodes["p_weight"].replace_all_uses_with(nodes["x"]) + ep.graph_module.recompile() + assert not nodes["p_weight"].users + assert not RemoveUnusedConstantsPass()(ep).modified + ep.validate() + assert ep.graph_signature.parameters == ("weight",) + assert "weight" in ep.state_dict + + +def test_preserves_storage_referenced_by_another_placeholder(): + model = torch.nn.Linear(2, 2, bias=False) + inputs = (torch.randn(1, 2),) + ep = torch.export.export(model, inputs) + weight = next(node for node in ep.graph.nodes if node.name == "p_weight") + with ep.graph.inserting_before(weight): + unused = ep.graph.placeholder("unused_weight") + unused.meta = weight.meta.copy() + ep.graph_signature.input_specs.insert( + 0, InputSpec(InputKind.PARAMETER, TensorArgument(unused.name), "weight") + ) + ep.graph_module.recompile() + ep.validate() + original = ep.state_dict["weight"] + assert RemoveUnusedConstantsPass()(ep).modified + ep.validate() + assert ep.state_dict["weight"] is original + torch.testing.assert_close(ep.module()(*inputs), model(*inputs)) + + +def test_preserves_constant_in_module_call_signature(): + class Inner(torch.nn.Module): + def forward(self, x, weight): + return x + 1 + + class Outer(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(2, 2)) + self.inner = Inner() + + def forward(self, x): + return self.inner(x, self.weight) + + model = Outer() + inputs = (torch.randn(2, 2),) + ep = torch.export.export(model, inputs, preserve_module_call_signature=("inner",)) + ep.validate() + weight = next(node for node in ep.graph.nodes if node.name == "p_weight") + assert not weight.users + assert not RemoveUnusedConstantsPass()(ep).modified + ep.validate() + torch.testing.assert_close(ep.module()(*inputs), model(*inputs))