From 627aef730e449446520e601ff4e00ec92c34e3ae Mon Sep 17 00:00:00 2001
From: RJ Ascani <rja@meta.com>
Date: Thu, 27 Aug 2026 11:33:16 -0700
Subject: [PATCH] Cortex-M: keep the activation on a max pool that will not
 lower

ActivationFusionPass folds a following relu, hardtanh or clamp into its producer
by narrowing that producer's quantized output range and erasing the activation
node. CortexMMaxPool2DCheck annotates a pool it cannot lower before tagging it,
so such a pool still carries qparams and reads as a fusion target. It then falls
back to the portable kernel, which reads no range, and the activation is gone:
relu(max_pool2d(dilation=2)) over randn(1, 4, 8, 8) * 20 came out unclamped,
19.8 away from the reference at its worst element.

Nothing caught it because the op counts look right either way -- the activation
is supposed to disappear -- and no test paired a pool with one.

Authored with Claude Code.
---
 .../cortex_m/passes/activation_fusion_pass.py |  9 +++++
 backends/cortex_m/passes/passes_utils.py      | 13 +++++++
 backends/cortex_m/test/ops/test_max_pool2d.py | 38 +++++++++++++++++++
 3 files changed, 60 insertions(+)

diff --git a/backends/cortex_m/passes/activation_fusion_pass.py b/backends/cortex_m/passes/activation_fusion_pass.py
index eb11e4ebc4c..36372fb98dd 100644
--- a/backends/cortex_m/passes/activation_fusion_pass.py
+++ b/backends/cortex_m/passes/activation_fusion_pass.py
@@ -11,6 +11,7 @@
 from executorch.backends.cortex_m.passes.passes_utils import (
     get_activation_bounds,
     quantize_val,
+    skips_quantized_max_pool2d,
 )
 
 from executorch.exir.dialects._ops import ops as exir_ops
@@ -179,6 +180,14 @@ def call(self, graph_module: GraphModule) -> PassResult:
                     f"Cannot fuse activation {node.name} as input node {input_node.name} has multiple users."
                 )
                 continue
+            if skips_quantized_max_pool2d(input_node):
+                # Fusing means narrowing the producer's output range and erasing
+                # this node. A pool that will not lower runs on the portable
+                # kernel, which reads no range, so the activation would be lost.
+                logger.warning(
+                    f"Cannot fuse activation {node.name} as input node {input_node.name} will not lower."
+                )
+                continue
 
             if (qparams_dict := self._get_validated_qparams(node, input_node)) is None:
                 continue
diff --git a/backends/cortex_m/passes/passes_utils.py b/backends/cortex_m/passes/passes_utils.py
index fe18cc1b141..b8d9134c0a4 100644
--- a/backends/cortex_m/passes/passes_utils.py
+++ b/backends/cortex_m/passes/passes_utils.py
@@ -185,6 +185,19 @@ def is_foldable_alpha(alpha: Any) -> TypeGuard[int]:
     return isinstance(alpha, int)
 
 
+def skips_quantized_max_pool2d(node: Node) -> bool:
+    """Whether CortexMMaxPool2DCheck tagged this pool as one that cannot lower.
+
+    The checker annotates the node before tagging it, so a tagged pool still
+    carries qparams and looks lowerable to anything that only reads those.
+    """
+    return (
+        node.meta.get("custom", {})
+        .get("cortex_m", {})
+        .get("skip_quantized_max_pool2d", False)
+    )
+
+
 def is_qualified_int8_node(args) -> bool:
     try:
         if len(args) < 6:
diff --git a/backends/cortex_m/test/ops/test_max_pool2d.py b/backends/cortex_m/test/ops/test_max_pool2d.py
index d394747dfb5..2b75b55ecf9 100644
--- a/backends/cortex_m/test/ops/test_max_pool2d.py
+++ b/backends/cortex_m/test/ops/test_max_pool2d.py
@@ -77,6 +77,33 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.pool(x)[1]
 
 
+class CortexMMaxPool2dRelu(torch.nn.Module):
+    """A pool that declines to lower, with an activation after it.
+
+    ActivationFusionPass folds the activation into the producer's output range
+    and erases its node. The portable fallback reads no range, so a pool that
+    declines has to keep the activation as a node of its own.
+    """
+
+    ops_before_transforms = {
+        "executorch_exir_dialects_edge__ops_aten_max_pool2d_with_indices_default": 1,
+        "executorch_exir_dialects_edge__ops_aten_relu_default": 1,
+    }
+    ops_after_transforms = {
+        "executorch_exir_dialects_edge__ops_aten_max_pool2d_with_indices_default": 1,
+        "executorch_exir_dialects_edge__ops_aten_clamp_default": 1,
+        "executorch_exir_dialects_edge__ops_cortex_m_quantized_max_pool2d_default": 0,
+    }
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.pool = torch.nn.MaxPool2d(*args, **kwargs)
+        self.relu = torch.nn.ReLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.relu(self.pool(x))
+
+
 test_cases = {
     "maxpool_2x2": McuTestCase(
         CortexMMaxPool2d(kernel_size=2, stride=2),
@@ -164,6 +191,17 @@ def test_dialect_max_pool2d_fallback(test_case, cortex_m_target):
     )
 
 
+def test_activation_after_a_declining_pool_survives(cortex_m_target):
+    """Erasing it would be silent: the graph still runs, just unclamped."""
+    model = CortexMMaxPool2dRelu(kernel_size=2, stride=1, dilation=2)
+    tester = CortexMTester(
+        model,
+        ((torch.randn(1, 4, 8, 8) * 20),),
+        target_config=cortex_m_target,
+    )
+    tester.test_dialect(model.ops_before_transforms, model.ops_after_transforms, qtol=1)
+
+
 @parametrize("test_case", fallback_test_cases)
 def test_executorch_max_pool2d_fallback(test_case, cortex_m_target):
     tester = CortexMTester(