Cortex-M: keep NHWC channel broadcasts quantized

backends/cortex_m/passes/explicit_layout_pass.py

@@ -4,9 +4,12 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import operator
+
 import executorch.backends.transforms.channels_last_ops  # noqa: F401
 
 import torch
+from executorch.backends.cortex_m.passes.passes_utils import is_flat_channel_broadcast
 
 from executorch.backends.transforms.replace_ops_with_channels_last_variants import (
     ChannelsLastOpSpec,
@@ -16,6 +19,7 @@
 )
 from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import PassResult
 
 
 def _is_rank4(node: torch.fx.Node) -> bool:
@@ -84,6 +88,68 @@ def _supports_max_pool2d(node: torch.fx.Node) -> bool:
     )
 
 
+_CORTEX_M_EXPLICIT_LAYOUT_OP_MAP = {
+    exir_ops.edge.aten.convolution.default: ChannelsLastOpSpec(
+        target=exir_ops.edge.channels_last.convolution.default,
+        input_indices=[0],
+        output_indices=[0],
+        filter_fn=lambda node: (
+            _is_rank4(node) and _has_input_and_output_qparams(node)
+        ),
+    ),
+    exir_ops.edge.aten.avg_pool2d.default: ChannelsLastOpSpec(
+        target=exir_ops.edge.channels_last.avg_pool2d.default,
+        input_indices=[0],
+        output_indices=[0],
+        filter_fn=_supports_avg_pool2d,
+    ),
+    exir_ops.edge.aten.max_pool2d.default: ChannelsLastOpSpec(
+        target=exir_ops.edge.channels_last.max_pool2d.default,
+        input_indices=[0],
+        output_indices=[0],
+        filter_fn=_supports_max_pool2d,
+    ),
+}
+
+_SOURCE_ANCHORS_BY_EDGE_TARGET = {
+    exir_ops.edge.aten.convolution.default: frozenset(
+        {
+            torch.ops.aten.conv1d.default,
+            torch.ops.aten.conv2d.default,
+            torch.ops.aten.conv_transpose2d.input,
+        }
+    ),
+    exir_ops.edge.aten.avg_pool2d.default: frozenset(
+        {torch.ops.aten.avg_pool2d.default}
+    ),
+    exir_ops.edge.aten.max_pool2d.default: frozenset(
+        {
+            torch.ops.aten.max_pool2d.default,
+            torch.ops.aten.max_pool2d_with_indices.default,
+        }
+    ),
+}
+assert _SOURCE_ANCHORS_BY_EDGE_TARGET.keys() == _CORTEX_M_EXPLICIT_LAYOUT_OP_MAP.keys()
+
+CORTEX_M_EXPLICIT_LAYOUT_SOURCE_ANCHORS = frozenset(
+    target for targets in _SOURCE_ANCHORS_BY_EDGE_TARGET.values() for target in targets
+)
+
+CORTEX_M_EXPLICIT_LAYOUT_TRANSPARENT_OPS = frozenset(
+    {
+        operator.getitem,
+        torch.ops.aten.relu.default,
+        torch.ops.aten.relu_.default,
+        torch.ops.aten.hardtanh.default,
+        torch.ops.aten.hardtanh_.default,
+        torch.ops.aten.clamp.default,
+        torch.ops.aten.clamp_.default,
+        torch.ops.aten.hardsigmoid.default,
+        torch.ops.aten.hardsigmoid_.default,
+    }
+)
+
+
 def _can_propagate(node: torch.fx.Node) -> bool:
     if node.target == exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default:
         return False
@@ -114,28 +180,29 @@ def __init__(
     ) -> None:
         super().__init__(
             exported_program,
-            op_map={
-                exir_ops.edge.aten.convolution.default: ChannelsLastOpSpec(
-                    target=exir_ops.edge.channels_last.convolution.default,
-                    input_indices=[0],
-                    output_indices=[0],
-                    filter_fn=lambda node: (
-                        _is_rank4(node) and _has_input_and_output_qparams(node)
-                    ),
-                ),
-                exir_ops.edge.aten.avg_pool2d.default: ChannelsLastOpSpec(
-                    target=exir_ops.edge.channels_last.avg_pool2d.default,
-                    input_indices=[0],
-                    output_indices=[0],
-                    filter_fn=_supports_avg_pool2d,
-                ),
-                exir_ops.edge.aten.max_pool2d.default: ChannelsLastOpSpec(
-                    target=exir_ops.edge.channels_last.max_pool2d.default,
-                    input_indices=[0],
-                    output_indices=[0],
-                    filter_fn=_supports_max_pool2d,
-                ),
-            },
+            op_map=dict(_CORTEX_M_EXPLICIT_LAYOUT_OP_MAP),
             can_propagate=_can_propagate,
             strict=strict,
         )
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        result = super().call(graph_module)
+        for node in result.graph_module.graph.nodes:
+            if node.target not in {
+                exir_ops.edge.aten.add.Tensor,
+                exir_ops.edge.aten.mul.Tensor,
+            } or not _has_input_and_output_qparams(node):
+                continue
+            input1, input2 = node.args[:2]
+            if not isinstance(input1, torch.fx.Node) or not isinstance(
+                input2, torch.fx.Node
+            ):
+                continue
+            if input1.meta["val"].shape == input2.meta["val"].shape:
+                continue
+            if not is_flat_channel_broadcast(input1.meta["val"], input2.meta["val"]):
+                raise RuntimeError(
+                    f"Quantized channel-broadcast node {node.name} did not join "
+                    "an explicit NHWC layout region."
+                )
+        return result

backends/cortex_m/passes/passes_utils.py

@@ -345,10 +345,6 @@ def is_channel_broadcast(tensor1: torch.Tensor, tensor2: torch.Tensor) -> bool:
     """
     if tensor1.dim() != tensor2.dim():
         return False
-    if not is_channels_last(tensor1):
-        return False
-    if not is_channels_last(tensor2):
-        return False
 
     channel_match = tensor1.size(1) == tensor2.size(1)
     tensor1_channels_only = tensor1.numel() == tensor1.size(1)

backends/cortex_m/quantizer/pattern_checkers.py

@@ -7,6 +7,10 @@
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.arm.quantizer.arm_quantizer_utils import PatternCheck
 from executorch.backends.arm.quantizer.quantization_config import QuantizationConfig
+from executorch.backends.cortex_m.passes.explicit_layout_pass import (
+    CORTEX_M_EXPLICIT_LAYOUT_SOURCE_ANCHORS,
+    CORTEX_M_EXPLICIT_LAYOUT_TRANSPARENT_OPS,
+)
 from executorch.backends.cortex_m.passes.passes_utils import (
     coerce_int_pair,
     is_channel_broadcast,
@@ -56,6 +60,50 @@ def check_quantization_config(
         return is_per_tensor and is_int8
 
 
+class CortexMExplicitAddMulCheck(CortexMAddMulCheck):
+    @classmethod
+    def _reaches_layout_anchor(cls, starts: list[Node], traverse_inputs: bool) -> bool:
+        pending = list(starts)
+        visited: set[Node] = set()
+        while pending:
+            node = pending.pop()
+            if node in visited:
+                continue
+            visited.add(node)
+            if node.target in CORTEX_M_EXPLICIT_LAYOUT_SOURCE_ANCHORS:
+                return True
+            if node.target not in CORTEX_M_EXPLICIT_LAYOUT_TRANSPARENT_OPS:
+                continue
+            pending.extend(node.all_input_nodes if traverse_inputs else node.users)
+        return False
+
+    @classmethod
+    def check_pattern(cls, pattern):
+        pattern_nodes = set(pattern)
+        for node in pattern:
+            if len(node.all_input_nodes) != 2:
+                continue
+            tensor1 = get_first_fake_tensor(node.all_input_nodes[0])
+            tensor2 = get_first_fake_tensor(node.all_input_nodes[1])
+            if tensor1.shape == tensor2.shape:
+                continue
+            if not is_channel_broadcast(tensor1, tensor2):
+                return False
+
+            external_users = [
+                user
+                for pattern_node in pattern
+                for user in pattern_node.users
+                if user not in pattern_nodes
+            ]
+            if not cls._reaches_layout_anchor(
+                list(node.all_input_nodes), traverse_inputs=True
+            ) or not cls._reaches_layout_anchor(external_users, traverse_inputs=False):
+                return False
+
+        return True
+
+
 class CortexMDivCheck(PatternCheck):
 
     @classmethod

backends/cortex_m/quantizer/quantizer.py

@@ -19,6 +19,8 @@
     NodeTargetNodeFinder,
 )
 from executorch.backends.cortex_m.quantizer.pattern_checkers import (
+    CortexMAddMulCheck,
+    CortexMExplicitAddMulCheck,
     CortexMExplicitConv1DCheck,
     CortexMExplicitConv2DCheck,
     CortexMExplicitConvTranspose2DCheck,
@@ -85,6 +87,9 @@ def __init__(
         )
         support_dict = dict(CORTEX_M_QUANTIZER_SUPPORT_DICT)
         if use_explicit_layout:
+            for pattern, checker in support_dict.items():
+                if checker is CortexMAddMulCheck:
+                    support_dict[pattern] = CortexMExplicitAddMulCheck
             for pattern in CONV1D_OP_PATTERNS:
                 support_dict[pattern] = CortexMExplicitConv1DCheck
             for pattern in CONV_OP_PATTERNS: