Arm backend: Fix Arm PT2E quantization with export guards

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:

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