diff --git a/backends/arm/quantizer/arm_quantizer.py b/backends/arm/quantizer/arm_quantizer.py index 99fcbf3745c..1d5af210247 100644 --- a/backends/arm/quantizer/arm_quantizer.py +++ b/backends/arm/quantizer/arm_quantizer.py @@ -592,6 +592,39 @@ def _for_each_filtered_node( yield node +def _strip_exported_program_guards(model: GraphModule) -> None: + """Discard torch.export's side-effect-only input-guard call before PT2E. + + Since torch 2.15, ExportedProgram.module() may insert a call_module node + targeting ``_guards_fn``. The Arm annotation passes cannot lower that node. + This is equivalent to using ``module(check_guards=False)`` for PT2E; + unrelated call_module nodes must not be removed. + + """ + guard_nodes = [ + node + for node in model.graph.nodes + if node.op == "call_module" and node.target == "_guards_fn" + ] + if not guard_nodes: + return + + # The generated guard returns None and has no consumers. Never erase a + # call whose result is used, even if it happens to have this name. + if any(node.users for node in guard_nodes): + raise RuntimeError("Cannot remove _guards_fn: its result is in use") + + for node in guard_nodes: + model.graph.erase_node(node) + + # PyTorch's exported-module pre-hook tests for this attribute. Removing it + # also restores the standard input constraint checks provided by the hook, + # as happens with module(check_guards=False). + if hasattr(model, "_guards_fn"): + delattr(model, "_guards_fn") + model.recompile() + + class TOSAQuantizer(Quantizer): """Manage quantization annotations for TOSA-compatible backends. @@ -898,6 +931,7 @@ def transform_for_annotation(self, model: GraphModule) -> GraphModule: GraphModule: Transformed model prepared for annotation. """ + _strip_exported_program_guards(model) return self.quantizer.transform_for_annotation(model) def annotate(self, model: GraphModule) -> GraphModule: diff --git a/backends/arm/test/quantizer/test_arm_export_guards_pt2e.py b/backends/arm/test/quantizer/test_arm_export_guards_pt2e.py new file mode 100644 index 00000000000..ba0303b3f04 --- /dev/null +++ b/backends/arm/test/quantizer/test_arm_export_guards_pt2e.py @@ -0,0 +1,114 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +"""Tests for fix of Arm PT2E quantizer with export guards.""" + +import pytest +import torch +from executorch.backends.arm.ethosu import EthosUCompileSpec +from executorch.backends.arm.quantizer import ( + EthosUQuantizer, + get_symmetric_quantization_config, +) +from executorch.backends.arm.quantizer.arm_quantizer import ( + _strip_exported_program_guards, +) +from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e + + +class _ConvMean(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, 3) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x).mean(dim=(2, 3)) + + +class _NoOpGuard(torch.nn.Module): + def forward(self, *args) -> None: + return None + + +def _guard_calls(gm: torch.fx.GraphModule) -> list[torch.fx.Node]: + return [ + n for n in gm.graph.nodes if n.op == "call_module" and n.target == "_guards_fn" + ] + + +def _ensure_guard_call(gm: torch.fx.GraphModule) -> None: + """Force the real exported-graph shape when torch suppresses guards. + + Some PyTorch versions suppress _guards_fn if the caller's path contains + 'executorch'. Thus a test *inside this repo* needs an explicit fallback. + + """ + if _guard_calls(gm): + return + gm.add_module("_guards_fn", _NoOpGuard()) + placeholders = [n for n in gm.graph.nodes if n.op == "placeholder"] + assert placeholders + with gm.graph.inserting_after(placeholders[-1]): + gm.graph.call_module("_guards_fn", args=tuple(placeholders)) + gm.recompile() + + +@pytest.mark.parametrize("use_composable_quantizer", [True, False]) +@pytest.mark.parametrize("guarded", [True, False]) +def test_ethosu_prepare_pt2e_accepts_exported_module_guards( + use_composable_quantizer: bool, guarded: bool +) -> None: + inputs = (torch.randn(1, 3, 16, 16),) + ep = torch.export.export(_ConvMean().eval(), inputs) + gm = ep.module() if guarded else ep.module(check_guards=False) + if guarded: + _ensure_guard_call(gm) + assert _guard_calls(gm) + else: + assert not _guard_calls(gm) + + reference = gm(*inputs) + quantizer = EthosUQuantizer( + EthosUCompileSpec("ethos-u85-256"), + use_composable_quantizer=use_composable_quantizer, + ) + quantizer.set_global(get_symmetric_quantization_config(is_per_channel=True)) + + prepared = prepare_pt2e(gm, quantizer) + assert not _guard_calls(prepared) + assert not hasattr(prepared, "_guards_fn") + torch.testing.assert_close(prepared(*inputs), reference, atol=1e-4, rtol=1e-4) + + +def test_only_export_guard_call_is_removed() -> None: + gm = torch.export.export(_ConvMean().eval(), (torch.randn(1, 3, 16, 16),)).module( + check_guards=False + ) + _ensure_guard_call(gm) + gm.add_module("other_module", _NoOpGuard()) + guard = _guard_calls(gm)[0] + with gm.graph.inserting_after(guard): + gm.graph.call_module("other_module", args=()) + gm.recompile() + + _strip_exported_program_guards(gm) + assert not _guard_calls(gm) + assert not hasattr(gm, "_guards_fn") + assert any( + n.op == "call_module" and n.target == "other_module" for n in gm.graph.nodes + ) + + +def test_guard_with_used_output_is_not_discarded() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + guarded_result = graph.call_module("_guards_fn", args=(x,)) + graph.output(guarded_result) + + root = torch.nn.Module() + root.add_module("_guards_fn", torch.nn.Identity()) + gm = torch.fx.GraphModule(root, graph) + with pytest.raises(RuntimeError, match="its result is in use"): + _strip_exported_program_guards(gm) + assert len(_guard_calls(gm)) == 1