Arm backend: Add TOSA dialect op for RESIZE (#14513)

Add TOSA backend dialect op for TOSA RESIZE. The dialect op replaces
upsample_nearest2d and upsample_bilinear_2d in RewriteUpsamplePass. Also
the Nodevisitors of upsample_nearest2d and upsample_bilinear2d are
replaced by one NodeVisitor for the resize backend dialect op.


Signed-off-by: Oscar Andersson <[email protected]>

Author: oscarandersson8218

backends/arm/_passes/__init__.py

@@ -91,6 +91,7 @@
     ReplaceScalarWithTensorArgPassTOSABI,
     ReplaceScalarWithTensorArgPassTOSAMI,
 )
+from .rewrite_upsample import RewriteUpsamplePass  # noqa
 from .scalars_to_attribute_pass import ScalarsToAttributePass  # noqa
 from .size_adjust_input_pass import SizeAdjustInputPass  # noqa
 from .to_tosa_memory_format_pass import ToTosaMemoryFormatPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -91,6 +91,7 @@
     ReplaceScalarWithTensorArgPassTOSABI,
     ReplaceScalarWithTensorArgPassTOSAMI,
     RetraceFoldedDtypesPass,
+    RewriteUpsamplePass,
     ScalarsToAttributePass,
     SizeAdjustInputPass,
     ToTosaMemoryFormatPass,
@@ -206,6 +207,7 @@ def _tosa_INT_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         # needs to happen before AddBiasPass, but after the table ops are inserted
         # to be able to validate that conv2d has right dtype arguments.
         self.add_pass(DecomposeConv2dWithInt16ActivationPass())
+        self.add_pass(RewriteUpsamplePass(exported_program))
         self.add_pass(AddBiasPass(exported_program))
 
         self.add_pass(FuseEqualPlaceholdersPass(exported_program))
@@ -290,6 +292,7 @@ def _tosa_FP_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(FuseViewCopyTransform())
         self.add_pass(FuseConstantArgsPass(exported_program))
         self.add_pass(CastInt64BuffersToInt32Pass(exported_program))
+        self.add_pass(RewriteUpsamplePass(exported_program))
         self.add_pass(AddBiasPass(exported_program))
         self.add_pass(InsertTableOpsPass(exported_program))
         self.add_pass(FuseEqualPlaceholdersPass(exported_program))

backends/arm/_passes/fuse_constant_ops_pass.py

@@ -114,8 +114,9 @@ def call(self, graph_module):
             if node.op != "call_function":
                 continue
             if node.target in [
-                exir_ops.backend.tosa.TABLE.default,
                 exir_ops.backend.tosa.RESCALE.default,
+                exir_ops.backend.tosa.RESIZE.default,
+                exir_ops.backend.tosa.TABLE.default,
                 exir_ops.backend.tosa.TRANSPOSE.default,
             ]:
                 continue

backends/arm/_passes/rewrite_upsample.py

@@ -0,0 +1,84 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Set, Type
+
+import torch
+from executorch.backends.arm._passes import ArmPass
+from executorch.backends.arm._passes.arm_pass_utils import (
+    create_node,
+    get_first_fake_tensor,
+)
+from executorch.backends.arm.tosa.utils import get_resize_parameters
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass, PassResult
+
+
+class RewriteUpsamplePass(ArmPass):
+    """Rewrite upsample2d nodes to TOSA.RESIZE nodes."""
+
+    targeted_ops = (
+        exir_ops.edge.aten.upsample_nearest2d.vec,
+        exir_ops.edge.aten.upsample_bilinear2d.vec,
+    )
+
+    _passes_required_after: Set[Type[ExportPass]] = set()
+
+    def call(self, graph_module):
+        modified = False
+        for node in graph_module.graph.nodes:
+            if node.op != "call_function" or node.target not in self.targeted_ops:
+                continue
+            modified = True
+
+            if node.target == exir_ops.edge.aten.upsample_bilinear2d.vec:
+                x, output_size, align_corners, scale_factors = node.args
+                resize_mode = "bilinear"
+            else:
+                x, output_size, scale_factors = node.args
+                align_corners = False
+                resize_mode = "nearest"
+
+            with graph_module.graph.inserting_before(node):
+                tosa_resize_node = create_node(
+                    graph_module.graph,
+                    op_target=exir_ops.backend.tosa.RESIZE.default,
+                    args=(x, output_size, align_corners, scale_factors),
+                    kwargs={"resize_mode": resize_mode},
+                    from_node=node,
+                )
+                node.replace_all_uses_with(tosa_resize_node)
+                graph_module.graph.erase_node(node)
+            input_dtype = get_first_fake_tensor(x).dtype
+            if input_dtype == torch.int8 and resize_mode == "bilinear":
+                input_size = get_first_fake_tensor(x).shape
+                input_size_xy = input_size[2:]
+                output_size = get_first_fake_tensor(node).shape
+                output_size_xy = output_size[2:]
+                scale_n_yx, _, _, _ = get_resize_parameters(
+                    input_size_xy=input_size_xy,
+                    output_size_xy=output_size_xy,
+                    resize_mode=1,
+                    align_corners=align_corners,
+                )
+                output_dtype = get_first_fake_tensor(node).dtype
+                output_scale = float(1 / (scale_n_yx[0] * scale_n_yx[1]))
+                with graph_module.graph.inserting_after(tosa_resize_node):
+                    rescale_node = create_node(
+                        graph_module.graph,
+                        exir_ops.backend.tosa.RESCALE.default,
+                    )
+                    tosa_resize_node.replace_all_uses_with(rescale_node)
+                    rescale_node.args = (
+                        tosa_resize_node,
+                        output_dtype,
+                        output_scale,
+                        0,  # zero point
+                        0,  # zero point
+                    )
+
+        if modified:
+            graph_module = super().call(graph_module).graph_module
+        return PassResult(graph_module, modified)

backends/arm/operators/__init__.py

@@ -43,6 +43,7 @@
     op_reciprocal,
     op_repeat,
     op_rescale,
+    op_resize,
     op_rshift_tensor,
     op_rsqrt,
     op_sigmoid,
@@ -54,8 +55,6 @@
     op_tanh,
     op_to_dim_order_copy,
     op_transpose,
-    op_upsample_bilinear2d,
-    op_upsample_nearest2d,
     op_view,
     op_where,
     ops_binary,

backends/arm/operators/op_upsample_nearest2d.py renamed to backends/arm/operators/op_resize.py

@@ -24,8 +24,8 @@
 
 
 @register_node_visitor
-class UpsampleNearest2dVisitor(NodeVisitor):
-    target = "aten.upsample_nearest2d.vec"
+class ResizeVisitor(NodeVisitor):
+    target = "tosa.RESIZE.default"
 
     tosa_specs = NodeVisitor.tosa_specs
 
@@ -41,12 +41,18 @@ def define_node(
     ) -> None:
         import serializer.tosa_serializer as ts
 
-        validate_num_inputs(self.target, inputs, 3)
-        validate_same_dtype(self.target, [inputs[0], output], ts)
+        validate_num_inputs(self.target, inputs, [3, 4])
+        if node.kwargs.get("resize_mode") == "bilinear":
+            resize_mode = ResizeMode.BILINEAR
+            align_corners = bool(node.args[2])
+        else:
+            resize_mode = ResizeMode.NEAREST
+            align_corners = False
+            validate_same_dtype(self.target, [inputs[0], output], ts)
         validate_valid_dtype(
             self.target,
             [inputs[0], output],
-            [ts.DType.INT8, ts.DType.INT32, ts.DType.FP32],
+            [ts.DType.INT8, ts.DType.INT32, ts.DType.FP16, ts.DType.FP32],
             output.tosa_spec,
         )
 
@@ -59,7 +65,7 @@ def define_node(
         # 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=False
+            input_size_yx, output_size_yx, resize_mode, align_corners=align_corners
         )
 
         def in_int16_range(x):
@@ -86,7 +92,7 @@ def in_int16_range(x):
         )
         attr = ts.TosaSerializerAttribute()
         attr.ResizeAttribute(
-            mode=ResizeMode.NEAREST,
+            mode=resize_mode,
         )
 
         self._serialize_operator(

backends/arm/operators/op_upsample_bilinear2d.py

@@ -5,6 +5,7 @@
 
 from executorch.backends.arm.tosa.dialect.ops import (  # noqa F401
     rescale,
+    resize,
     table,
     transpose,
 )