From 561ec76048f563d32e12175f472458db1ab57dec Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 19 Aug 2026 21:51:12 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../cortex_m/passes/aten_to_cortex_m_pass.py | 22 ++- .../cortex_m/passes/explicit_layout_pass.py | 123 +++++++++--- backends/cortex_m/passes/passes_utils.py | 12 +- .../cortex_m/quantizer/pattern_checkers.py | 48 +++++ backends/cortex_m/quantizer/quantizer.py | 5 + .../cortex_m/test/test_explicit_layout.py | 176 +++++++++++++++++- 6 files changed, 352 insertions(+), 34 deletions(-) diff --git a/backends/cortex_m/passes/aten_to_cortex_m_pass.py b/backends/cortex_m/passes/aten_to_cortex_m_pass.py index 333e76e18b2..d8136131c54 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py @@ -61,9 +61,11 @@ def __init__( self, exported_program: ExportedProgram, target_config: CortexMTargetConfig, + use_explicit_layout: bool = False, ) -> None: super().__init__(exported_program=exported_program) self.target_config = target_config + self.use_explicit_layout = use_explicit_layout def call(self, graph_module: torch.fx.GraphModule) -> PassResult: result = super().call(graph_module) @@ -943,7 +945,6 @@ def _get_dequantize_per_tensor_replacement( def _get_add_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: - del dialect_pass if not _has_qparams(node): return None @@ -979,14 +980,20 @@ def _get_add_replacement( activation_min, activation_max, ) - return DialectNodeSpec(exir_ops.edge.cortex_m.quantized_add.default, args) + target = exir_ops.edge.cortex_m.quantized_add.default + if ( + cast(AtenToCortexMPass, dialect_pass).use_explicit_layout + and _get_input_tensor_data(node, 0).shape + != _get_input_tensor_data(node, 1).shape + ): + target = exir_ops.edge.cortex_m.quantized_add_nhwc.default + return DialectNodeSpec(target, args) @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.mul.Tensor) def _get_mul_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: - del dialect_pass if not _has_qparams(node): return None @@ -1009,7 +1016,14 @@ def _get_mul_replacement( output_mult, output_shift, ) - return DialectNodeSpec(exir_ops.edge.cortex_m.quantized_mul.default, args) + target = exir_ops.edge.cortex_m.quantized_mul.default + if ( + cast(AtenToCortexMPass, dialect_pass).use_explicit_layout + and _get_input_tensor_data(node, 0).shape + != _get_input_tensor_data(node, 1).shape + ): + target = exir_ops.edge.cortex_m.quantized_mul_nhwc.default + return DialectNodeSpec(target, args) @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.div.Tensor) diff --git a/backends/cortex_m/passes/explicit_layout_pass.py b/backends/cortex_m/passes/explicit_layout_pass.py index 4b908ebb3f7..08706e62ee0 100644 --- a/backends/cortex_m/passes/explicit_layout_pass.py +++ b/backends/cortex_m/passes/explicit_layout_pass.py @@ -4,6 +4,8 @@ # 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 @@ -16,6 +18,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 +87,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 @@ -104,6 +169,19 @@ def _can_propagate(node: torch.fx.Node) -> bool: return tensor1.shape == tensor2.shape or _has_input_and_output_qparams(node) +def _is_nhwc_channel_broadcast(node: torch.fx.Node) -> bool: + input1, input2 = node.args[:2] + if not isinstance(input1, torch.fx.Node) or not isinstance(input2, torch.fx.Node): + return False + tensor1 = input1.meta.get("val") + tensor2 = input2.meta.get("val") + if tensor1 is None or tensor2 is None or tensor1.dim() != 4 or tensor2.dim() != 4: + return False + return tensor1.size(3) == tensor2.size(3) and ( + tensor1.numel() == tensor1.size(3) or tensor2.numel() == tensor2.size(3) + ) + + class CortexMExplicitLayoutPass(ToContiguousChannelsLastPass): """Configure the common explicit-layout pipeline for Cortex-M kernels.""" @@ -114,29 +192,30 @@ 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, layout_pad_target=exir_ops.edge.channels_last.constant_pad_nd.default, 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_nhwc_channel_broadcast(node): + raise RuntimeError( + f"Quantized channel-broadcast node {node.name} did not join " + "an explicit NHWC layout region." + ) + return result diff --git a/backends/cortex_m/passes/passes_utils.py b/backends/cortex_m/passes/passes_utils.py index bcb828c5928..5bbab5c075e 100644 --- a/backends/cortex_m/passes/passes_utils.py +++ b/backends/cortex_m/passes/passes_utils.py @@ -336,16 +336,20 @@ def to_physical_order(logical_pad: list[int], tensor: torch.Tensor) -> list[int] return [logical_pad[_NHWC_DIM_ORDER[i]] for i in range(4)] -def is_channel_broadcast(tensor1: torch.Tensor, tensor2: torch.Tensor) -> bool: +def is_channel_broadcast( + tensor1: torch.Tensor, + tensor2: torch.Tensor, + require_channels_last: bool = True, +) -> bool: """ Check if tensor1 is broadcasted to tensor2 along channel dimension. Assumes tensor2 has shape [N, C, ...] and tensor1 has shape [N, 1, ...] or [1, C, ...]. """ if tensor1.dim() != tensor2.dim(): return False - if not is_channels_last(tensor1): - return False - if not is_channels_last(tensor2): + if require_channels_last and ( + not is_channels_last(tensor1) or not is_channels_last(tensor2) + ): return False channel_match = tensor1.size(1) == tensor2.size(1) diff --git a/backends/cortex_m/quantizer/pattern_checkers.py b/backends/cortex_m/quantizer/pattern_checkers.py index 34d1c0dcd29..b6654410119 100644 --- a/backends/cortex_m/quantizer/pattern_checkers.py +++ b/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, require_channels_last=False): + 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 diff --git a/backends/cortex_m/quantizer/quantizer.py b/backends/cortex_m/quantizer/quantizer.py index e1386511502..98f24123b95 100644 --- a/backends/cortex_m/quantizer/quantizer.py +++ b/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: diff --git a/backends/cortex_m/test/test_explicit_layout.py b/backends/cortex_m/test/test_explicit_layout.py index 50cade0f97f..764c42a3d9a 100644 --- a/backends/cortex_m/test/test_explicit_layout.py +++ b/backends/cortex_m/test/test_explicit_layout.py @@ -113,6 +113,51 @@ def forward(self, x): return self.conv2(torch.relu(self.conv1(x)) + self.bias) +class ConvBiasBranched(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 4, 3, padding=1) + self.conv2 = torch.nn.Conv2d(4, 5, 3, padding=1) + self.bias = torch.nn.Parameter(torch.randn(1, 4, 1, 1)) + + def forward(self, x): + biased = self.conv1(x) + self.bias + return self.conv2(biased), biased + + +class ConvBiasOutput(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 4, 3, padding=1) + self.bias = torch.nn.Parameter(torch.randn(1, 4, 1, 1)) + + def forward(self, x): + return self.conv(x) + self.bias + + +class ConvMulConv(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 4, 3, padding=1) + self.conv2 = torch.nn.Conv2d(4, 5, 3, padding=1) + self.scale = torch.nn.Parameter(torch.randn(1, 4, 1, 1)) + + def forward(self, x): + return self.conv2(self.conv1(x) * self.scale) + + +class ConvPoolBiasConv(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 4, 3, padding=1) + self.pool = torch.nn.AvgPool2d(2, 2) + self.conv2 = torch.nn.Conv2d(4, 5, 3, padding=1) + self.bias = torch.nn.Parameter(torch.randn(1, 4, 1, 1)) + + def forward(self, x): + return self.conv2(self.pool(self.conv1(x)) + self.bias) + + class ConvForkAdd(torch.nn.Module): def __init__(self): super().__init__() @@ -329,18 +374,62 @@ def test_softmax_stays_outside_explicit_layout_region(): assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 -def test_unquantized_channel_bias_stays_outside_explicit_layout_region(): +def test_channel_bias_is_quantized_inside_explicit_layout_region(): x = torch.randn(1, 3, 8, 8) program = _lower(ConvBiasConv(), (x,)) [add] = [ node for node in program.graph.nodes - if node.target == exir_ops.edge.aten.add.Tensor + if node.target == exir_ops.edge.cortex_m.quantized_add_nhwc.default ] - assert add.args[1].meta["val"].shape == torch.Size([1, 4, 1, 1]) + assert add.args[4].meta["val"].shape == torch.Size([1, 1, 1, 4]) + assert _count(program, exir_ops.edge.aten.add.Tensor) == 0 assert _count(program, exir_ops.edge.cortex_m.quantized_add.default) == 0 - assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 4 + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 + + +def test_branched_channel_bias_stays_quantized(): + x = torch.randn(1, 3, 8, 8) + program = _lower(ConvBiasBranched(), (x,)) + + assert _count(program, exir_ops.edge.cortex_m.quantized_add_nhwc.default) == 1 + assert _count(program, exir_ops.edge.aten.add.Tensor) == 0 + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 3 + + +def test_unanchored_channel_bias_falls_back_without_failing(): + x = torch.randn(1, 3, 8, 8) + program = _lower(ConvBiasOutput(), (x,)) + + assert _count(program, exir_ops.edge.cortex_m.quantized_add_nhwc.default) == 0 + assert _count(program, exir_ops.edge.aten.add.Tensor) == 1 + + +def test_channel_mul_is_quantized_inside_explicit_layout_region(): + x = torch.randn(1, 3, 8, 8) + program = _lower(ConvMulConv(), (x,)) + [mul] = [ + node + for node in program.graph.nodes + if node.target == exir_ops.edge.cortex_m.quantized_mul_nhwc.default + ] + + assert mul.args[2].meta["val"].shape == torch.Size([1, 1, 1, 4]) + assert _count(program, exir_ops.edge.aten.mul.Tensor) == 0 + assert _count(program, exir_ops.edge.cortex_m.quantized_mul.default) == 0 + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 + + +def test_pool_anchor_allows_quantized_channel_broadcast(): + x = torch.randn(1, 3, 8, 8) + program = _lower(ConvPoolBiasConv(), (x,)) + + assert _count(program, exir_ops.edge.cortex_m.quantized_add_nhwc.default) == 1 + assert ( + _count(program, exir_ops.edge.cortex_m.quantized_avg_pool2d_nhwc.default) == 1 + ) + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 def test_explicit_layout_does_not_increase_planned_memory_for_float_qdq(): @@ -409,6 +498,45 @@ def test_explicit_layout_does_not_increase_planned_memory_for_float_qdq(): {exir_ops.edge.cortex_m.quantized_conv2d.default: 3}, {exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: 3}, ), + ( + ConvBiasConv(), + (torch.randn(1, 3, 8, 8),), + m55, + { + exir_ops.edge.cortex_m.quantized_conv2d.default: 2, + exir_ops.edge.cortex_m.quantized_add.default: 1, + }, + { + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: 2, + exir_ops.edge.cortex_m.quantized_add_nhwc.default: 1, + }, + ), + ( + ConvBiasBranched(), + (torch.randn(1, 3, 8, 8),), + m55, + { + exir_ops.edge.cortex_m.quantized_conv2d.default: 2, + exir_ops.edge.cortex_m.quantized_add.default: 1, + }, + { + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: 2, + exir_ops.edge.cortex_m.quantized_add_nhwc.default: 1, + }, + ), + ( + ConvMulConv(), + (torch.randn(1, 3, 8, 8),), + m55, + { + exir_ops.edge.cortex_m.quantized_conv2d.default: 2, + exir_ops.edge.cortex_m.quantized_mul.default: 1, + }, + { + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: 2, + exir_ops.edge.cortex_m.quantized_mul_nhwc.default: 1, + }, + ), ( ConvPadConv(), (torch.randn(1, 3, 8, 8),), @@ -429,6 +557,19 @@ def test_explicit_layout_does_not_increase_planned_memory_for_float_qdq(): {exir_ops.edge.cortex_m.quantized_conv2d.default: 1}, {exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: 1}, ), + ( + ConvBiasConv(), + (torch.randn(1, 3, 8, 8),), + m33, + { + exir_ops.edge.cortex_m.quantized_conv2d.default: 2, + exir_ops.edge.cortex_m.quantized_add.default: 1, + }, + { + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: 2, + exir_ops.edge.cortex_m.quantized_add_nhwc.default: 1, + }, + ), ) for module, inputs, target_config, legacy_ops, explicit_ops in cases: @@ -520,6 +661,33 @@ def test_explicit_nhwc_pad_runs_on_fvp_with_singleton_height(): ) +def test_explicit_nhwc_channel_broadcast_add_runs_on_fvp(): + x = torch.linspace(-5, 5, steps=1 * 3 * 7 * 9).reshape(1, 3, 7, 9) + _run_explicit_layout_on_fvp( + ConvBiasConv(), + (x,), + exir_ops.edge.cortex_m.quantized_add_nhwc.default, + ) + + +def test_explicit_nhwc_branched_channel_broadcast_runs_on_fvp(): + x = torch.linspace(-5, 5, steps=1 * 3 * 7 * 9).reshape(1, 3, 7, 9) + _run_explicit_layout_on_fvp( + ConvBiasBranched(), + (x,), + exir_ops.edge.cortex_m.quantized_add_nhwc.default, + ) + + +def test_explicit_nhwc_channel_broadcast_mul_runs_on_fvp(): + x = torch.linspace(-5, 5, steps=1 * 3 * 7 * 9).reshape(1, 3, 7, 9) + _run_explicit_layout_on_fvp( + ConvMulConv(), + (x,), + exir_ops.edge.cortex_m.quantized_mul_nhwc.default, + ) + + def test_aot_explicit_layout_conv1d_runs_on_fvp(): from types import SimpleNamespace