From f8983a9df13ee3279a4d66fd089a5aeb10712c5e Mon Sep 17 00:00:00 2001 From: Oscar Andersson Date: Mon, 31 Aug 2026 08:57:27 +0200 Subject: [PATCH] Arm backend: Relax SymbolicShapeSupport-checks Allow TOSA support checks to distinguish between symbolic tensor metadata and symbolic shape materialization. Without the TOSA shape extension, symbolic tensor input/output shapes can still be accepted for ops whose lowering does not need materialized shape values. Reject symbolic SymInt arguments and the known shape- dependent edge cases instead. For targets that require fully resolved tensor shapes, add a separate partitioner-level check that rejects nodes with unresolved tensor input or output shapes. Signed-off-by: Oscar Andersson Change-Id: I1a65d72ceea18a9b5e08779c2987cdd44b70a589 --- backends/arm/ethosu/partitioner.py | 1 + .../tosa_supported_operators.py | 196 +++++++++++---- .../tosa_dialect/test_tosa_shape_support.py | 229 +++++++++++++++++- backends/arm/tosa/partitioner.py | 8 +- backends/arm/vgf/partitioner.py | 1 + 5 files changed, 381 insertions(+), 54 deletions(-) diff --git a/backends/arm/ethosu/partitioner.py b/backends/arm/ethosu/partitioner.py index eb04e69c8bf..0cf46cd59b0 100644 --- a/backends/arm/ethosu/partitioner.py +++ b/backends/arm/ethosu/partitioner.py @@ -37,6 +37,7 @@ def __init__( self.compile_spec = compile_spec self.additional_checks = additional_checks self.tosa_spec = compile_spec.tosa_spec + self._requires_resolved_tensor_shapes = True self._decomposable_resize_support = DecomposableResizeSupported(self.tosa_spec) self._custom_partition_ops: set[OpOverload] = set() self.intermediate_path = compile_spec._get_intermediate_path() diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index 4eae3f09dba..525cc565ea1 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -513,7 +513,10 @@ def tosa_support_factory( if additional_positive_checks: positive_checks.extend(additional_positive_checks) negative_checks = _negative_checks( - tosa_spec, exported_program, reporter, additional_checks + tosa_spec, + exported_program, + reporter, + additional_checks, ) # An op must be accepted by at least one postitive check, and not rejected by any @@ -530,8 +533,44 @@ def tosa_support_factory( return any_chain(*additional_positive_overrides or (), default_checks) +def _has_symbolic_shape(node: fx.Node) -> bool: + val = node.meta.get("val") + vals = val if isinstance(val, (list, tuple)) else (val,) + for node_val in vals: + if isinstance(node_val, torch.SymInt): + return True + + shape = getattr(node_val, "shape", None) + if shape is not None and any(isinstance(dim, torch.SymInt) for dim in shape): + return True + + return False + + class SymbolicShapeSupportCheck(OperatorSupportBase): - """Reject symbolic tensor shapes for specs without the shape extension.""" + """Reject symbolic shape constructs that require the TOSA shape + extension. + """ + + _SYMBOLIC_SPATIAL_DIM_TARGETS = ( + exir_ops.edge.aten.convolution.default, + exir_ops.edge.aten.avg_pool2d.default, + exir_ops.edge.aten.max_pool2d.default, + exir_ops.edge.aten.max_pool2d_with_indices.default, + exir_ops.edge.aten._adaptive_avg_pool2d.default, + torch.ops.aten.conv_transpose2d.input, + ) + _SYMBOLIC_MEAN_TARGETS = ( + exir_ops.edge.aten.mean.dim, + exir_ops.edge.aten.mean.default, + ) + _SYMBOLIC_VIEW_SHAPE_TARGETS = ( + exir_ops.edge.aten.squeeze_copy.dim, + exir_ops.edge.aten.squeeze_copy.dims, + exir_ops.edge.aten.unsqueeze_copy.default, + ) + _SYMBOLIC_PAD_TARGETS = (exir_ops.edge.aten.constant_pad_nd.default,) + _SYMBOLIC_SLICE_TARGETS = (exir_ops.edge.aten.slice_copy.Tensor,) def __init__(self, reporter: WhyNoPartitionReporter): """Initialize the check with a reporter. @@ -543,88 +582,147 @@ def __init__(self, reporter: WhyNoPartitionReporter): self.reporter = reporter @staticmethod - def _has_symbolic_shape(node: fx.Node) -> bool: - val = node.meta.get("val") - vals = val if isinstance(val, (list, tuple)) else (val,) - for node_val in vals: - if isinstance(node_val, torch.SymInt): + def _has_symbolic_shape_argument(arg: object) -> bool: + if isinstance(arg, torch.SymInt): + return True + + if isinstance(arg, fx.Node): + return SymbolicShapeSupportCheck._has_symbolic_shape_argument( + arg.meta.get("val") + ) + + if isinstance(arg, (list, tuple)): + return any( + SymbolicShapeSupportCheck._has_symbolic_shape_argument(item) + for item in arg + ) + + return False + + @staticmethod + def _get_mean_reduction_dims(node: fx.Node, input_rank: int) -> tuple[int, ...]: + if node.target == exir_ops.edge.aten.mean.default: + return tuple(range(input_rank)) + + dims = node.kwargs.get("dim", node.args[1] if len(node.args) > 1 else None) + if dims is None: + return tuple(range(input_rank)) + if isinstance(dims, int): + return (dims % input_rank,) + return tuple(dim % input_rank for dim in typing.cast(Sequence[int], dims)) + + def _has_unsupported_symbolic_tensor_shape(self, node: fx.Node) -> bool: + if node.target not in ( + *self._SYMBOLIC_SPATIAL_DIM_TARGETS, + *self._SYMBOLIC_MEAN_TARGETS, + *self._SYMBOLIC_VIEW_SHAPE_TARGETS, + *self._SYMBOLIC_PAD_TARGETS, + *self._SYMBOLIC_SLICE_TARGETS, + ): + return False + if not node.all_input_nodes: + return False + + input_node = node.all_input_nodes[0] + input_fake_tensor = get_first_fake_tensor(input_node) + if not any(isinstance(s, torch.SymInt) for s in input_fake_tensor.shape): + return False + + if node.target in self._SYMBOLIC_SPATIAL_DIM_TARGETS: + if any(isinstance(s, torch.SymInt) for s in input_fake_tensor.shape[2:]): + self.reporter.report_reject(node, "Symbolic spatial dims unsupported") return True - shape = getattr(node_val, "shape", None) - if shape is not None and any( - isinstance(dim, torch.SymInt) for dim in shape + if node.target in self._SYMBOLIC_MEAN_TARGETS: + if any( + isinstance(input_fake_tensor.shape[dim], torch.SymInt) + for dim in self._get_mean_reduction_dims( + node, len(input_fake_tensor.shape) + ) ): + self.reporter.report_reject(node, "Symbolic mean dims unsupported") return True + if node.target in self._SYMBOLIC_VIEW_SHAPE_TARGETS: + self.reporter.report_reject(node, "Symbolic view dims unsupported") + return True + + if node.target in self._SYMBOLIC_PAD_TARGETS: + self.reporter.report_reject(node, "Symbolic pad dims unsupported") + return True + + if node.target in self._SYMBOLIC_SLICE_TARGETS: + self.reporter.report_reject(node, "Symbolic slices unsupported") + return True + return False - def _partition_dynamic_upmsample_nearest2d(self, node: fx.Node) -> bool: - """Check if the node is an upsample_nearest2d with symbolic shapes. + def is_node_supported( + self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node + ) -> bool: + """Return False for symbolic shape uses needing shape extension. + + Without TOSA shape extension, symbolic input/output tensor dimensions + are generally allowed because they are tensor metadata. Symbolic shape + arguments and known shape-materialization edge cases are rejected. Args: + submodules (typing.Mapping[str, torch.nn.Module]): Exported modules. node (fx.Node): FX node to check. Returns: - bool: True if the node is an upsample_nearest2d with symbolic - shapes; otherwise, False. + bool: False if rejected by constraints; otherwise, True. """ - if node.target != exir_ops.edge.aten.upsample_nearest2d.vec: - return False + del submodules + if node.op in ("placeholder", "output"): + return True + if node.op == "call_function" and node.target in (*Q_OPS, *DQ_OPS): + return True - try: - input_tensor = get_first_fake_tensor(node.all_input_nodes[0]) - output_tensor = get_first_fake_tensor(node) - except Exception as exc: + if self._has_symbolic_shape_argument(node.args): self.reporter.report_reject( node, - f"upsample_nearest2d symbolic shapes need tensor metadata: {exc}", + "Node has symbolic shape arguments, has the TOSA spec shape extension support?", ) return False - input_size_xy = input_tensor.shape[2:4] - output_size_xy = output_tensor.shape[2:4] - if len(input_size_xy) != 2 or len(output_size_xy) != 2: - self.reporter.report_reject( - node, "upsample_nearest2d expects 2D spatial input/output." - ) + if self._has_unsupported_symbolic_tensor_shape(node): return False return True - def is_node_supported( - self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node - ) -> bool: - """Return False for nodes with symbolic tensor input or output shapes. - Dynamic shapes require the TOSA shape extension. Reject nodes with - symbolic tensor dimensions before partitioning when the active spec - does not enable that extension. +class CheckResolvedTensorShapes(OperatorSupportBase): + """Reject nodes with unresolved tensor input or output shapes.""" - Args: - submodules (typing.Mapping[str, torch.nn.Module]): Exported modules. - node (fx.Node): FX node to check. + def __init__(self, reporter: WhyNoPartitionReporter): + """Initialize the check with a reporter. - Returns: - bool: False if rejected by constraints; otherwise, True. + Args: + reporter (WhyNoPartitionReporter): Reporter for rejection reasons. """ + self.reporter = reporter + + def is_node_supported( + self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node + ) -> bool: + """Return False when the node depends on unresolved tensor shapes.""" + del submodules if node.op in ("placeholder", "output"): return True if node.op == "call_function" and node.target in (*Q_OPS, *DQ_OPS): return True - if self._has_symbolic_shape(node) or any( - self._has_symbolic_shape(input_node) for input_node in node.all_input_nodes + if _has_symbolic_shape(node) or any( + _has_symbolic_shape(input_node) for input_node in node.all_input_nodes ): - if node.target == exir_ops.edge.aten.upsample_nearest2d.vec: - return self._partition_dynamic_upmsample_nearest2d(node) - else: - self.reporter.report_reject( - node, - "Node has symbolic shape, has the TOSA spec shape extension support?", - ) - return False + self.reporter.report_reject( + node, + "Node has unresolved tensor shapes, which are not supported by this target.", + ) + return False return True diff --git a/backends/arm/test/misc/tosa_dialect/test_tosa_shape_support.py b/backends/arm/test/misc/tosa_dialect/test_tosa_shape_support.py index 1bc109b3dc1..e54451255e5 100644 --- a/backends/arm/test/misc/tosa_dialect/test_tosa_shape_support.py +++ b/backends/arm/test/misc/tosa_dialect/test_tosa_shape_support.py @@ -3,12 +3,18 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import executorch.backends.arm.operator_support.convolution_support # noqa: F401 +import executorch.backends.arm.operator_support.pool_2d_support # noqa: F401 +import executorch.backends.arm.operator_support.reduce_sum_support # noqa: F401 import executorch.backends.arm.operator_support.sym_size_int_support # noqa: F401 +import pytest import torch from executorch.backends.arm.operator_support.tosa_supported_operators import ( tosa_support_factory, ) +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.arm_tester import ArmTester from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec from executorch.backends.arm.tosa.partitioner import TOSAPartitioner from executorch.backends.arm.tosa.specification import TosaSpecification @@ -28,11 +34,60 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return torch.atan2(x, y) +class Conv2d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 4, 3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + +class AvgPool2d(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) + + +class MeanDim(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.mean(dim=2) + + +class MeanDefault(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.mean() + + +class Squeeze(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.squeeze(2) + + +class Unsqueeze(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.unsqueeze(2) + + +class Slice(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x[:, :, 1:, :] + + +class ScalarTensor(torch.nn.Module): + def forward(self) -> torch.Tensor: + return torch.scalar_tensor(1.0) + + class ReturnSymSize(torch.nn.Module): def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, int]: return x, x.shape[0] +class ReshapeWithSymSize(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.reshape(x.shape[0], 6) + + def _exported_program( module: torch.nn.Module, inputs: tuple[torch.Tensor, ...], @@ -100,7 +155,7 @@ def test_shape_extension_accepts_supported_symbolic_tensor_op(): assert support.is_node_supported(exported_program.graph_module, add_node) is True -def test_without_shape_extension_rejects_supported_symbolic_tensor_op(): +def test_without_shape_extension_accepts_supported_symbolic_tensor_op(): inputs = (torch.randn(2, 3), torch.randn(2, 3)) batch = Dim("batch", min=1, max=4) exported_program = _exported_program( @@ -108,11 +163,177 @@ def test_without_shape_extension_rejects_supported_symbolic_tensor_op(): inputs, dynamic_shapes=({0: batch}, {0: batch}), ) - support, reporter = _support("TOSA-1.0+FP", exported_program) + support, _ = _support("TOSA-1.0+FP", exported_program) add_node = _find_node(exported_program, exir_ops.edge.aten.add.Tensor) - assert support.is_node_supported(exported_program.graph_module, add_node) is False - assert "Node has symbolic shape" in reporter.get_table_report() + assert support.is_node_supported(exported_program.graph_module, add_node) is True + + +def _assert_rejected_with_reason(exported_program, target, reason: str) -> None: + support, reporter = _support("TOSA-1.0+FP", exported_program) + node = _find_node(exported_program, target) + + assert support.is_node_supported(exported_program.graph_module, node) is False + assert reason in reporter.get_table_report() + + +@pytest.mark.parametrize( + "compile_spec", + (common.get_u55_compile_spec(), common.get_u85_compile_spec()), +) +def test_ethos_rejects_unresolved_tensor_shapes(compile_spec): + inputs = (torch.randn(2, 3), torch.randn(2, 3)) + batch = Dim("batch", min=1, max=4) + tester = ArmTester( + Add(), + inputs, + compile_spec, + dynamic_shapes=({0: batch}, {0: batch}), + ) + + tester.quantize().export().to_edge().partition() + targets = { + node.target + for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes + } + + assert exir_ops.edge.aten.add.Tensor in targets + assert torch.ops.higher_order.executorch_call_delegate not in targets + + +def test_without_shape_extension_rejects_symbolic_spatial_conv2d(): + inputs = (torch.randn(2, 3, 8, 8),) + height = Dim("height", min=4, max=10) + exported_program = _exported_program( + Conv2d(), + inputs, + dynamic_shapes=({2: height},), + ) + + _assert_rejected_with_reason( + exported_program, + exir_ops.edge.aten.convolution.default, + "Symbolic spatial dims unsupported", + ) + + +def test_without_shape_extension_rejects_symbolic_spatial_pooling(): + inputs = (torch.randn(2, 3, 8, 8),) + height = Dim("height", min=2, max=5) * 2 + exported_program = _exported_program( + AvgPool2d(), + inputs, + dynamic_shapes=({2: height},), + ) + + _assert_rejected_with_reason( + exported_program, + exir_ops.edge.aten.avg_pool2d.default, + "Symbolic spatial dims unsupported", + ) + + +def test_without_shape_extension_rejects_symbolic_mean_reduction_dim(): + inputs = (torch.randn(2, 3, 8, 8),) + height = Dim("height", min=4, max=10) + exported_program = _exported_program( + MeanDim(), + inputs, + dynamic_shapes=({2: height},), + ) + + _assert_rejected_with_reason( + exported_program, + exir_ops.edge.aten.mean.dim, + "Symbolic mean dims unsupported", + ) + + +def test_without_shape_extension_rejects_symbolic_full_tensor_mean(): + inputs = (torch.randn(2, 3, 8, 8),) + height = Dim("height", min=4, max=10) + exported_program = _exported_program( + MeanDefault(), + inputs, + dynamic_shapes=({2: height},), + ) + + _assert_rejected_with_reason( + exported_program, + exir_ops.edge.aten.mean.default, + "Symbolic mean dims unsupported", + ) + + +def test_without_shape_extension_rejects_symbolic_squeeze(): + inputs = (torch.randn(2, 3, 1, 8),) + batch = Dim("batch", min=1, max=4) + exported_program = _exported_program( + Squeeze(), + inputs, + dynamic_shapes=({0: batch},), + ) + + _assert_rejected_with_reason( + exported_program, + exir_ops.edge.aten.squeeze_copy.dims, + "Symbolic view dims unsupported", + ) + + +def test_without_shape_extension_rejects_symbolic_unsqueeze(): + inputs = (torch.randn(2, 3, 8),) + batch = Dim("batch", min=1, max=4) + exported_program = _exported_program( + Unsqueeze(), + inputs, + dynamic_shapes=({0: batch},), + ) + + _assert_rejected_with_reason( + exported_program, + exir_ops.edge.aten.unsqueeze_copy.default, + "Symbolic view dims unsupported", + ) + + +def test_without_shape_extension_rejects_symbolic_slice(): + inputs = (torch.randn(2, 3, 8, 8),) + height = Dim("height", min=4, max=10) + exported_program = _exported_program( + Slice(), + inputs, + dynamic_shapes=({2: height},), + ) + + _assert_rejected_with_reason( + exported_program, + exir_ops.edge.aten.slice_copy.Tensor, + "Symbolic slices unsupported", + ) + + +def test_without_shape_extension_accepts_zero_input_supported_op(): + exported_program = _exported_program(ScalarTensor(), ()) + support, _ = _support("TOSA-1.0+FP", exported_program) + scalar_node = _find_node(exported_program, torch.ops.aten.scalar_tensor.default) + + assert support.is_node_supported(exported_program.graph_module, scalar_node) is True + + +def test_without_shape_extension_rejects_symbolic_shape_argument(): + inputs = (torch.randn(2, 2, 3),) + batch = Dim("batch", min=1, max=4) + exported_program = _exported_program( + ReshapeWithSymSize(), + inputs, + dynamic_shapes=({0: batch},), + ) + support, reporter = _support("TOSA-1.0+FP", exported_program) + view_node = _find_node(exported_program, exir_ops.edge.aten.view_copy.default) + + assert support.is_node_supported(exported_program.graph_module, view_node) is False + assert "Node has symbolic shape arguments" in reporter.get_table_report() def test_without_shape_extension_rejects_sym_size_int(): diff --git a/backends/arm/tosa/partitioner.py b/backends/arm/tosa/partitioner.py index e0fa9b5ee9f..91fe23dd7e4 100644 --- a/backends/arm/tosa/partitioner.py +++ b/backends/arm/tosa/partitioner.py @@ -38,6 +38,7 @@ from executorch.backends.arm.common.type import ensure_type from executorch.backends.arm.constants import DQ_OPS, MAX_RANK, Q_OPS from executorch.backends.arm.operator_support.tosa_supported_operators import ( + CheckResolvedTensorShapes, is_quantized, tosa_support_factory, ) @@ -453,6 +454,7 @@ def __init__( self.compile_spec = compile_spec self.tosa_spec = compile_spec.tosa_spec self.additional_checks = additional_checks + self._requires_resolved_tensor_shapes = False self._decomposable_resize_support = DecomposableResizeSupported(self.tosa_spec) self._custom_partition_ops: set[torch._ops.OpOverload] = set() self.intermediate_path = compile_spec._get_intermediate_path() @@ -746,11 +748,15 @@ def _create_operator_support( if self._custom_partition_ops else None ) + additional_checks = list(self.additional_checks or ()) + if self._requires_resolved_tensor_shapes: + additional_checks.append(CheckResolvedTensorShapes(reporter)) + return tosa_support_factory( self.tosa_spec, containing_program, reporter, - self.additional_checks, + additional_checks=additional_checks, additional_positive_checks=[ self._decomposable_resize_support, DecomposableLargeStrideMaxPool2dForU55Supported(self.tosa_spec), diff --git a/backends/arm/vgf/partitioner.py b/backends/arm/vgf/partitioner.py index be7de275af4..4d1b6c6f4b5 100644 --- a/backends/arm/vgf/partitioner.py +++ b/backends/arm/vgf/partitioner.py @@ -43,3 +43,4 @@ def __init__( self.intermediate_path = compile_spec._get_intermediate_path() # Preserve grid_sampler_2d for the VGF custom-lowering path only. self.register_custom_partition_op(exir_ops.edge.aten.grid_sampler_2d.default) + self._requires_resolved_tensor_shapes = False