diff --git a/backends/arm/operators/op_upsample_bilinear2d.py b/backends/arm/operators/op_upsample_bilinear2d.py index 4d8c0ff9320..8705930d3b0 100644 --- a/backends/arm/operators/op_upsample_bilinear2d.py +++ b/backends/arm/operators/op_upsample_bilinear2d.py @@ -49,15 +49,18 @@ def define_node( input_dtype = inputs[0].dtype # tosa_shape output is NHWC, take HW - input_size_yx = torch.tensor( - tosa_shape(inputs[0].shape, inputs[0].dim_order)[1:3] - ) - # Ignore scale and size parameters, directly use the output size as - # we only support static shapes currently - output_size_yx = torch.tensor(tosa_shape(output.shape, output.dim_order)[1:3]) + input_size_yx = tuple([inputs[0].shape[dim] for dim in inputs[0].dim_order])[ + 1:3 + ] + output_size_yx = tuple([output.shape[dim] for dim in output.dim_order])[1:3] + # Get align_corners value from the node arguments. + align_corners = bool(node.args[2]) scale_n_yx, scale_d_yx, offset_yx, border_yx = get_resize_parameters( - input_size_yx, output_size_yx, ResizeMode.NEAREST, align_corners=True + input_size_yx, + output_size_yx, + ResizeMode.NEAREST, + align_corners=align_corners, ) def in_int16_range(x): @@ -139,15 +142,18 @@ def define_node( input_dtype = inputs[0].dtype # tosa_shape output is NHWC, take HW - input_size_yx = torch.tensor( - tosa_shape(inputs[0].shape, inputs[0].dim_order)[1:3] - ) - # Ignore scale and size parameters, directly use the output size as - # we only support static shapes currently - output_size_yx = torch.tensor(tosa_shape(output.shape, output.dim_order)[1:3]) + input_size_yx = tuple([inputs[0].shape[dim] for dim in inputs[0].dim_order])[ + 1:3 + ] + output_size_yx = tuple([output.shape[dim] for dim in output.dim_order])[1:3] + # Get align_corners value from the node arguments. + align_corners = bool(node.args[2]) scale_n_yx, scale_d_yx, offset_yx, border_yx = get_resize_parameters( - input_size_yx, output_size_yx, ResizeMode.NEAREST, align_corners=True + input_size_yx, + output_size_yx, + ResizeMode.NEAREST, + align_corners=align_corners, ) def in_int16_range(x): diff --git a/backends/arm/operators/op_upsample_nearest2d.py b/backends/arm/operators/op_upsample_nearest2d.py index e9c72145555..0a3cb356257 100644 --- a/backends/arm/operators/op_upsample_nearest2d.py +++ b/backends/arm/operators/op_upsample_nearest2d.py @@ -17,7 +17,7 @@ validate_same_dtype, ) from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_utils import get_resize_parameters, tosa_shape +from executorch.backends.arm.tosa_utils import get_resize_parameters from tosa_tools.v0_80.tosa.ResizeMode import ResizeMode # type: ignore @@ -43,19 +43,16 @@ def define_node( validate_num_inputs(self.target, inputs, 3) validate_same_dtype(self.target, [inputs[0], output]) - if inputs[0].shape is None or output.shape is None: - raise ValueError("Only static shapes are supported") - # tosa_shape output is NHWC, take HW - input_size_yx = torch.tensor( - tosa_shape(inputs[0].shape, inputs[0].dim_order)[1:3] - ) - # Ignore scale and size parameters, directly use the output size as - # we only support static shapes currently - output_size_yx = torch.tensor(tosa_shape(output.shape, output.dim_order)[1:3]) + input_size_yx = tuple([inputs[0].shape[dim] for dim in inputs[0].dim_order])[ + 1:3 + ] + output_size_yx = tuple([output.shape[dim] for dim in output.dim_order])[1:3] + # Align corners shouldn't make a difference for nearest upsampling. We set to False so + # half pixel centers are used for resize parameter logic. scale_n_yx, scale_d_yx, offset_yx, border_yx = get_resize_parameters( - input_size_yx, output_size_yx, ResizeMode.NEAREST, align_corners=True + input_size_yx, output_size_yx, ResizeMode.NEAREST, align_corners=False ) def in_int16_range(x): @@ -102,19 +99,16 @@ def define_node( validate_num_inputs(self.target, inputs, 3) validate_same_dtype(self.target, [inputs[0], output]) - if inputs[0].shape is None or output.shape is None: - raise ValueError("Only static shapes are supported") - # tosa_shape output is NHWC, take HW - input_size_yx = torch.tensor( - tosa_shape(inputs[0].shape, inputs[0].dim_order)[1:3] - ) - # Ignore scale and size parameters, directly use the output size as - # we only support static shapes currently - output_size_yx = torch.tensor(tosa_shape(output.shape, output.dim_order)[1:3]) + input_size_yx = tuple([inputs[0].shape[dim] for dim in inputs[0].dim_order])[ + 1:3 + ] + output_size_yx = tuple([output.shape[dim] for dim in output.dim_order])[1:3] + # Align corners shouldn't make a difference for nearest upsampling. We set to False so + # half pixel centers are used for resize parameter logic. scale_n_yx, scale_d_yx, offset_yx, border_yx = get_resize_parameters( - input_size_yx, output_size_yx, ResizeMode.NEAREST, align_corners=True + input_size_yx, output_size_yx, ResizeMode.NEAREST, align_corners=False ) def in_int16_range(x): diff --git a/backends/arm/test/ops/test_upsample_nearest2d.py b/backends/arm/test/ops/test_upsample_nearest2d.py index 7809d5fdee2..d38d9fbe380 100644 --- a/backends/arm/test/ops/test_upsample_nearest2d.py +++ b/backends/arm/test/ops/test_upsample_nearest2d.py @@ -40,6 +40,17 @@ "rand_one_and_half_size": lambda: (torch.rand(2, 4, 8, 3), (12, 4), None, False), } +test_data_suite_dynamic = { + # (test_name, test_data, size, scale_factor, compare_outputs) + "rand_double_scale": lambda: (torch.rand(2, 4, 8, 3), None, 2.0, False), + "rand_double_scale_one_dim": lambda: ( + torch.rand(2, 4, 8, 3), + None, + (1.0, 2.0), + False, + ), +} + class UpsamplingNearest2d(torch.nn.Module): def __init__( @@ -161,3 +172,159 @@ def test_upsample_nearest2d_vec_tosa_BI_nearest(test_data: torch.Tensor): pipeline.pop_stage(-1) pipeline.run() + + +@common.parametrize("test_data", test_data_suite_dynamic) +def test_upsample_nearest2d_dynamic_MI_nearest(test_data: torch.Tensor): + test_data, size, scale_factor, compare_outputs = test_data() + + batch_size = torch.export.Dim("batch", min=0, max=1000) + input_height = torch.export.Dim("input_height", min=0, max=1000) + input_width = torch.export.Dim("input_width", min=0, max=1000) + + dynamic_shapes = {"x": {0: batch_size, 2: input_height, 3: input_width}} + + pipeline = TosaPipelineMI[input_t1]( + UpsamplingNearest2d(size, scale_factor), + (test_data,), + aten_op, + exir_op=[], + dynamic_shapes=dynamic_shapes, + ) + if not compare_outputs: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite_dynamic) +def test_upsample_nearest2d_dynamic_BI_nearest(test_data: torch.Tensor): + test_data, size, scale_factor, compare_outputs = test_data() + + batch_size = torch.export.Dim("batch", min=0, max=2) + input_height = torch.export.Dim("input_height", min=0, max=8) + input_width = torch.export.Dim("input_width", min=0, max=8) + + dynamic_shapes = {"x": {0: batch_size, 2: input_height, 3: input_width}} + + pipeline = TosaPipelineBI[input_t1]( + UpsamplingNearest2d(size, scale_factor), + (test_data,), + aten_op, + exir_op=[], + dynamic_shapes=dynamic_shapes, + ) + if not compare_outputs: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite_dynamic) +def test_upsample_nearest2d_dynamic_MI_interpolate(test_data: torch.Tensor): + test_data, size, scale_factor, compare_outputs = test_data() + + batch_size = torch.export.Dim("batch", min=0, max=2) + input_height = torch.export.Dim("input_height", min=4, max=8) + input_width = torch.export.Dim("input_width", min=3, max=8) + + dynamic_shapes = { + "x": { + 0: batch_size, + 2: input_height, + 3: input_width, + } + } + + pipeline = TosaPipelineMI[input_t1]( + Interpolate(size, scale_factor), + (test_data,), + aten_op, + exir_op=[], + dynamic_shapes=dynamic_shapes, + ) + if not compare_outputs: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite_dynamic) +def test_upsample_nearest2d_dynamic_BI_interpolate(test_data: torch.Tensor): + test_data, size, scale_factor, compare_outputs = test_data() + + batch_size = torch.export.Dim("batch", min=0, max=2) + input_height = torch.export.Dim("input_height", min=4, max=8) + input_width = torch.export.Dim("input_width", min=3, max=8) + + dynamic_shapes = { + "x": { + 0: batch_size, + 2: input_height, + 3: input_width, + } + } + + pipeline = TosaPipelineBI[input_t1]( + Interpolate(size, scale_factor), + (test_data,), + aten_op, + exir_op=[], + dynamic_shapes=dynamic_shapes, + ) + if not compare_outputs: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite_dynamic) +def test_upsample_nearest2d_dynamic_MI_upsample(test_data: torch.Tensor): + test_data, size, scale_factor, compare_outputs = test_data() + + batch_size = torch.export.Dim("batch", min=0, max=1000) + input_height = torch.export.Dim("input_height", min=0, max=1000) + input_width = torch.export.Dim("input_width", min=0, max=1000) + + dynamic_shapes = { + "x": { + 0: batch_size, + 2: input_height, + 3: input_width, + } + } + + pipeline = TosaPipelineMI[input_t1]( + Upsample(size, scale_factor), + (test_data,), + aten_op, + exir_op=[], + dynamic_shapes=dynamic_shapes, + ) + if not compare_outputs: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite_dynamic) +def test_upsample_nearest2d_dynamic_BI_upsample(test_data: torch.Tensor): + test_data, size, scale_factor, compare_outputs = test_data() + + batch_size = torch.export.Dim("batch", min=0, max=2) + input_height = torch.export.Dim("input_height", min=0, max=8) + input_width = torch.export.Dim("input_width", min=0, max=8) + + dynamic_shapes = { + "x": { + 0: batch_size, + 2: input_height, + 3: input_width, + } + } + + pipeline = TosaPipelineBI[input_t1]( + Upsample(size, scale_factor), + (test_data,), + aten_op, + exir_op=[], + dynamic_shapes=dynamic_shapes, + ) + if not compare_outputs: + pipeline.pop_stage(-1) + pipeline.run() diff --git a/backends/arm/test/tester/test_pipeline.py b/backends/arm/test/tester/test_pipeline.py index c066f301c70..7642a1b5cb0 100644 --- a/backends/arm/test/tester/test_pipeline.py +++ b/backends/arm/test/tester/test_pipeline.py @@ -4,7 +4,7 @@ # LICENSE file in the root directory of this source tree. import logging -from typing import Callable, Dict, Generic, List, Optional, Type, TypeVar +from typing import Any, Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar import torch @@ -88,10 +88,14 @@ def __init__( compile_spec: List[CompileSpec], exir_ops: Optional[str | List[str]] = None, use_to_edge_transform_and_lower: bool = True, + dynamic_shapes: Optional[Tuple[Any]] = None, ): self.tester = ArmTester( - module, example_inputs=test_data, compile_spec=compile_spec + module, + example_inputs=test_data, + compile_spec=compile_spec, + dynamic_shapes=dynamic_shapes, ) self.aten_ops = aten_ops if isinstance(aten_ops, list) else [aten_ops] @@ -283,6 +287,7 @@ def __init__( atol: float = 1e-03, rtol: float = 1e-03, qtol: int = 1, + dynamic_shapes: Optional[Tuple[Any]] = None, ): tosa_profiles = { "0.80": TosaSpecification.create_from_string("TOSA-0.80+BI"), @@ -310,6 +315,7 @@ def __init__( compile_spec, exir_op, use_to_edge_transform_and_lower, + dynamic_shapes, ) self.add_stage(self.tester.quantize, quant_stage, pos=0) @@ -381,6 +387,7 @@ def __init__( atol: float = 1e-03, rtol: float = 1e-03, qtol: int = 0, + dynamic_shapes: Optional[Tuple[Any]] = None, ): tosa_profiles = { "0.80": TosaSpecification.create_from_string("TOSA-0.80+MI"), @@ -398,6 +405,7 @@ def __init__( compile_spec, exir_op, use_to_edge_transform_and_lower, + dynamic_shapes, ) self.add_stage_after( "export", diff --git a/backends/arm/tosa_utils.py b/backends/arm/tosa_utils.py index c5546647b10..e4bdccb93fd 100644 --- a/backends/arm/tosa_utils.py +++ b/backends/arm/tosa_utils.py @@ -7,8 +7,9 @@ import logging import os -from typing import Any, Optional, Tuple +from typing import Any, Optional +import sympy # type: ignore import torch import tosa_tools.v0_80.serializer.tosa_serializer as ts # type: ignore from executorch.backends.arm.tosa_mapping import TosaArg @@ -168,7 +169,13 @@ def is_consumer_node_depthwise_conv2d(node: Node): def tosa_shape(shape, dim_order): - return tuple([shape[dim] for dim in dim_order]) + reordered = tuple([shape[dim] for dim in dim_order]) + # Dynamic shapes in executorch are represented with torch.SymInt objects in the shapes, + # in TOSA we do not have this concept and instead use -1. + removed_symints = tuple( + [-1 if isinstance(d, torch.SymInt) else d for d in reordered] + ) + return removed_symints def expand_dims( @@ -200,46 +207,90 @@ def expand_dims( return intermediate +def get_resize_parameters_1d( + input_size: int | torch.SymInt, + output_size: int | torch.SymInt, + resize_mode: int, + align_corners: bool, +): + # We don't support align_corners for symbolic shapes, because handling the edge case where size == 1 is tricky. + if align_corners: + if (not isinstance(input_size, int)) or (not isinstance(output_size, int)): + raise RuntimeError( + "We do not support align_corners=True for symbolic shapes." + ) + + # SymInt seems to not actually work for symbolic expressions, so use the underlying sympy objects instead + input_size = ( + input_size.node._expr if isinstance(input_size, torch.SymInt) else input_size + ) + output_size = ( + output_size.node._expr if isinstance(output_size, torch.SymInt) else output_size + ) + if align_corners and input_size > 1 and output_size > 1: + scale_n = output_size - 1 + else: + scale_n = output_size + if align_corners and input_size > 1 and output_size > 1: + scale_d = input_size - 1 + else: + scale_d = input_size + ratio = scale_n / scale_d + if not sympy.sympify(ratio).is_constant(): + raise RuntimeError( + "Resize requires a constant ratio: " + str(ratio) + " is not constant!" + ) + gcd = sympy.gcd(scale_n, scale_d) + scale_n = 2 * scale_n // gcd + scale_d = 2 * scale_d // gcd + # These should always be whole integers, based on the above calculations + scale_n = int(scale_n.evalf()) + scale_d = int(scale_d.evalf()) + + if align_corners: + offset = 0 + else: + # Half pixel centers so input and output sampling positions are offset by 1/2 pixel. + offset = scale_d // 2 - scale_n // 2 + + # Calculate border to maintain the correct the output size. + # Note that this should always result in a constant value, as the ratio is constant. + border = scale_d * (output_size - 1) - scale_n * (input_size - 1) + offset + + if not sympy.sympify(border).is_constant(): + raise RuntimeError( + "Resize requires a constant border: " + str(border) + " is not constant!" + ) + + border = int(sympy.sympify(border).evalf()) + return scale_n, scale_d, offset, border + + def get_resize_parameters( - input_size: torch.Tensor, - output_size: torch.Tensor, + input_size_xy: tuple[int | torch.SymInt, int | torch.SymInt], + output_size_xy: tuple[int | torch.SymInt, int | torch.SymInt], resize_mode: int, align_corners: bool, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.IntTensor, ...]: """Get the tosa.resize parameters based on the input and output size. Args: - input_size (torch.Tensor): Size of the input - output_size (torch.Tensor): Size of the output + input_size_xy (tuple[int | torch.SymInt]): Size of the input + output_size_xy (tuple[int | torch.SymInt]): Size of the output resize_mode (tosa.ResizeMode): The TOSA resize mode align_corners (bool): Align the corners pixels of the input and output Returns: - scale_n (torch.Tensor), scale_d (torch.Tensor), - offset (torch.Tensor), border (torch.Tensor) + scale_n (torch.IntTensor), scale_d (torch.IntTensor), + offset (torch.IntTensor), border (torch.IntTensor) """ - assert torch.all(input_size > 0) - assert torch.all(output_size > 0) - - scale_n = torch.tensor( - [ - so - 1 if align_corners and si > 1 and so > 1 else so - for si, so in zip(input_size, output_size) - ] + + # Get the parameters for each dimension independently + y_params = get_resize_parameters_1d( + input_size_xy[0], output_size_xy[0], resize_mode, align_corners ) - scale_d = torch.tensor( - [ - si - 1 if align_corners and si > 1 and so > 1 else si - for si, so in zip(input_size, output_size) - ] + x_params = get_resize_parameters_1d( + input_size_xy[1], output_size_xy[1], resize_mode, align_corners ) - - gcd = torch.gcd(scale_n, scale_d) - scale_n = scale_n // gcd - scale_d = scale_d // gcd - - # No half-pixel centre support in PyTorch, no offset needed - offset = torch.zeros_like(input_size) - border = scale_d * (output_size - 1) - scale_n * (input_size - 1) + offset - - return scale_n, scale_d, offset, border + # Combine them together, so we return four 2-element tensors (scale_n, scale_d, offset, border) + return tuple(map(torch.IntTensor, zip(y_params, x_params)))