Arm backend: Preserve standalone noop-resizes

backends/arm/_passes/remove_noop_pass.py

@@ -9,10 +9,11 @@
 from typing import Any, Set, Type
 
 from executorch.backends.arm._passes import ArmOpTargetedPass
+from executorch.backends.arm.constants import DQ_OPS, Q_OPS
 
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult, ProxyValue
-from torch.fx import GraphModule
+from torch.fx import GraphModule, Node
 
 logger = logging.getLogger(__name__)
 
@@ -25,15 +26,18 @@ class RemoveNoopPass(ArmOpTargetedPass):
         exir_ops.edge.aten.cat.default,
         exir_ops.edge.aten.concatenate.default,
     )
+    _resize_ops = (
+        exir_ops.edge.aten.upsample_bilinear2d.vec,
+        exir_ops.edge.aten.upsample_nearest2d.vec,
+    )
     target_ops = (
         exir_ops.edge.dim_order_ops._clone_dim_order.default,
         exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
         exir_ops.edge.aten.alias_copy.default,
         exir_ops.edge.aten.copy.default,
         exir_ops.edge.aten.detach_copy.default,
         *_single_input_concat_ops,
-        exir_ops.edge.aten.upsample_bilinear2d.vec,
-        exir_ops.edge.aten.upsample_nearest2d.vec,
+        *_resize_ops,
         exir_ops.backend.tosa.PAD.default,
         exir_ops.backend.tosa.SLICE.default,
     )
@@ -54,7 +58,103 @@ def _get_static_shape(value: Any) -> tuple[int, ...] | None:
             return None
         return tuple(value)
 
+    @staticmethod
+    def _is_identity_scale_factor(scale_factors: Any) -> bool:
+        if type(scale_factors) in (int, float):
+            return scale_factors == 1.0
+        return (
+            isinstance(scale_factors, (list, tuple))
+            and len(scale_factors) == 2
+            and all(
+                type(scale) in (int, float) and scale == 1.0 for scale in scale_factors
+            )
+        )
+
+    def _resize_scale_factors(self, op, args: tuple[Any, ...]) -> Any:
+        return args[3] if op == exir_ops.edge.aten.upsample_bilinear2d.vec else args[2]
+
+    def _is_identity_resize(self, op, args: tuple[Any, ...]) -> bool:
+        scale_factors = self._resize_scale_factors(op, args)
+        if self._is_identity_scale_factor(scale_factors):
+            return True
+
+        output_spatial_size = self._get_static_shape(args[1])
+        input_spatial_size = self._get_static_shape(args[0].data.shape[2:])
+        return (
+            output_spatial_size is not None
+            and input_spatial_size is not None
+            and input_spatial_size == output_spatial_size
+            and scale_factors is None
+        )
+
+    def _is_identity_resize_node(self, node: Node) -> bool:
+        if node.target not in self._resize_ops:
+            return False
+        input_node = node.args[0]
+        if not isinstance(input_node, Node):
+            return False
+
+        scale_factors = self._resize_scale_factors(node.target, node.args)
+        if self._is_identity_scale_factor(scale_factors):
+            return True
+
+        output_spatial_size = self._get_static_shape(node.args[1])
+        input_spatial_size = self._get_static_shape(input_node.meta["val"].shape[2:])
+        return (
+            output_spatial_size is not None
+            and input_spatial_size is not None
+            and input_spatial_size == output_spatial_size
+            and scale_factors is None
+        )
+
+    def _is_removable_noop_node(self, node: Node) -> bool:
+        if node.target in self._single_input_concat_ops:
+            inputs = node.args[0]
+            return isinstance(inputs, (list, tuple)) and len(inputs) == 1
+
+        if node.target in (
+            exir_ops.edge.dim_order_ops._clone_dim_order.default,
+            exir_ops.edge.aten.alias_copy.default,
+            exir_ops.edge.aten.detach_copy.default,
+        ):
+            return True
+
+        if node.target == exir_ops.edge.aten.copy.default:
+            return True
+
+        if node.target == exir_ops.edge.dim_order_ops._to_dim_order_copy.default:
+            input_node = node.args[0]
+            if not isinstance(input_node, Node):
+                return False
+            input_dtype = input_node.meta["val"].dtype
+            output_dtype = node.kwargs.get("dtype", input_dtype)
+            return input_dtype == output_dtype
+
+        return False
+
+    def _partition_would_be_empty_without_identity_resize(
+        self, graph_module: GraphModule
+    ) -> bool:
+        has_identity_resize = False
+        for node in graph_module.graph.nodes:
+            if (
+                node.op != "call_function"
+                or node.target in Q_OPS
+                or node.target in DQ_OPS
+            ):
+                continue
+            if self._is_identity_resize_node(node):
+                has_identity_resize = True
+                continue
+            if self._is_removable_noop_node(node):
+                continue
+            return False
+        return has_identity_resize
+
     def call(self, graph_module: GraphModule) -> PassResult:
+        if self._partition_would_be_empty_without_identity_resize(graph_module):
+            return PassResult(graph_module, False)
+
         result = super().call(graph_module)
         # Removing a no-op can leave its shape operands without users.
         removed_dead_code = result.graph_module.graph.eliminate_dead_code()
@@ -63,7 +163,7 @@ def call(self, graph_module: GraphModule) -> PassResult:
             result.graph_module.recompile()
         return PassResult(result.graph_module, result.modified or removed_dead_code)
 
-    def call_operator(self, op, args, kwargs, meta, updated=False):
+    def call_operator(self, op, args, kwargs, meta, updated=False):  # noqa: C901
         if op not in self.target_ops:
             return super().call_operator(op, args, kwargs, meta, updated)
 
@@ -74,21 +174,8 @@ def call_operator(self, op, args, kwargs, meta, updated=False):
                 return inputs[0]
             return super().call_operator(op, args, kwargs, meta, updated)
 
-        if op in (
-            exir_ops.edge.aten.upsample_bilinear2d.vec,
-            exir_ops.edge.aten.upsample_nearest2d.vec,
-        ):
-            scale_factors = (
-                args[3] if op == exir_ops.edge.aten.upsample_bilinear2d.vec else args[2]
-            )
-            if (
-                isinstance(scale_factors, (list, tuple))
-                and len(scale_factors) == 2
-                and all(
-                    type(scale) in (int, float) and scale == 1.0
-                    for scale in scale_factors
-                )
-            ):
+        if op in self._resize_ops:
+            if self._is_identity_resize(op, args):
                 return args[0]
             return super().call_operator(op, args, kwargs, meta, updated)
 

backends/arm/test/ops/test_upsample_bilinear2d.py

@@ -131,6 +131,11 @@
     ),
 }
 
+test_data_suite_Uxx_same_size = {
+    "rand_same_size": lambda: (torch.rand(2, 3, 5, 5), (5, 5), None, False),
+    "rand_same_scale": lambda: (torch.rand(2, 3, 5, 5), None, 1.0, False),
+}
+
 test_data_u55 = {
     "rand_double_size": lambda: (torch.rand(2, 4, 8, 3), (16, 6), None, True),
 }
@@ -423,6 +428,24 @@ def test_upsample_bilinear2d_vec_u55_INT_UpsamplingBilinear2d_not_delegated(
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_suite_Uxx_same_size)
+def test_upsample_bilinear2d_vec_u85_INT_same_size(
+    test_data: torch.Tensor,
+):
+    test_data, size, scale_factor, compare_outputs = test_data()
+
+    pipeline = EthosU85PipelineINT[input_t1](
+        InterpolateAlignCornersFalse(size, scale_factor),
+        (test_data,),
+        aten_op,
+        qtol=1,
+        use_to_edge_transform_and_lower=True,
+    )
+    if not compare_outputs:
+        pipeline.pop_stage(-1)
+    pipeline.run()
+
+
 @common.parametrize("test_data", test_data_suite_Uxx)
 @common.XfailIfNoCorstone320
 def test_upsample_bilinear2d_vec_u85_INT_Upsample(test_data: input_t1):

backends/arm/test/ops/test_upsample_nearest2d.py

@@ -17,6 +17,7 @@
 
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineINT,
+    EthosU85PipelineINT,
     OpNotSupportedPipeline,
     TosaPipelineFP,
     TosaPipelineINT,
@@ -92,6 +93,11 @@
     "rand_octuple_size": lambda: (torch.rand(1, 4, 8, 3), (64, 24), None, True),
 }
 
+test_data_suite_u85_same_size = {
+    "rand_same_size": lambda: (torch.rand(2, 3, 5, 5), (5, 5), None, False),
+    "rand_same_scale": lambda: (torch.rand(2, 3, 5, 5), None, 1.0, 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),
@@ -513,6 +519,23 @@ def test_upsample_nearest2d_vec_u55_INT_UpsamplingNearest2d(
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_suite_u85_same_size)
+def test_upsample_nearest2d_vec_u85_INT_same_size(
+    test_data: torch.Tensor,
+):
+    test_data, size, scale_factor, compare_outputs = test_data()
+
+    pipeline = EthosU85PipelineINT[input_t1](
+        Interpolate(size, scale_factor),
+        (test_data,),
+        aten_op,
+        exir_op,
+    )
+    if not compare_outputs:
+        pipeline.pop_stage(-1)
+    pipeline.run()
+
+
 def test_upsample_nearest2d_vec_u55_INT_unsupported_scale_not_delegated():
     # 2.25 rounds a 2x2 input to 4x4, which looks like a supported 2x resize from shapes alone.
     test_data = torch.rand(1, 4, 2, 2)