Arm backend: Support for dynamic shapes and fix resize bugs (#11310)

Burton2000 · web-flow · commit 50f73763aae1 · 2025-06-06T13:07:52.000+02:00
Dynamic shapes are detected when converting torch shapes to TOSA shapes, when they contain SymInts rather than regular ints. These are converted to -1s as this is how TOSA represents dynamic shapes. The resize op needs special handling for dynamic shapes, and will only work when all of the TOSA parameters (scale n/d, offset, border) work out to be constant values independent of the shape. Also fix bug where align_corners was always set to True, when it should be False Signed-off-by: Richard Burton <richard.burton@arm.com> Co-authored-by: Robert Hughes @Rob-Hughes-Arm
diff --git 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
@@ -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
@@ -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
@@ -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
