Arm backend: Remove RESIZE that don't change shape

Sometimes, we get a TOSA RESIZE op stemming from
upsample_bilinear2d or upsample_nearest2d that doesn't
change the input & output resolution, it just passes from
int8 to int32 and back to int8. Removing these
resizes provides good uplift at runtime, especially
if they operate on large tensors.

Signed-off-by: George Gekov <george.gekov@arm.com>
Change-Id: If0215efe423228f05c092795cbb28d13b61bcfc6

Author: gggekov

backends/arm/_passes/remove_noop_pass.py

@@ -32,6 +32,8 @@ class RemoveNoopPass(ArmOpTargetedPass):
         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,
         exir_ops.backend.tosa.PAD.default,
         exir_ops.backend.tosa.SLICE.default,
     )
@@ -72,6 +74,24 @@ 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
+                )
+            ):
+                return args[0]
+            return super().call_operator(op, args, kwargs, meta, updated)
+
         if op == exir_ops.backend.tosa.PAD.default:
             padding = self._get_static_shape(args[1])
             # PAD is an identity when every before/after padding value is zero.

backends/arm/test/ops/test_upsample_bilinear2d.py

@@ -34,8 +34,6 @@
     "rand_double_size": lambda: (torch.rand(2, 4, 8, 3), (16, 6), None, True),
     "rand_one_double_scale": lambda: (torch.rand(2, 4, 1, 1), None, 2.0, True),
     "rand_one_double_size": lambda: (torch.rand(2, 4, 1, 1), (2, 2), None, True),
-    "rand_one_same_scale": lambda: (torch.rand(2, 4, 1, 1), None, 1.0, True),
-    "rand_one_same_size": lambda: (torch.rand(2, 4, 1, 1), (1, 1), None, True),
     # Can't compare outputs as the rounding when selecting the nearest pixel is
     # different between PyTorch and TOSA. Just check the legalization went well.
     # TODO Improve the test infrastructure to support more in depth verification
@@ -86,18 +84,6 @@
         None,
         True,
     ),
-    "randn_one_same_scale_negative": lambda: (
-        torch.randn(2, 4, 1, 1),
-        None,
-        1.0,
-        True,
-    ),
-    "randn_one_same_size_negative": lambda: (
-        torch.randn(2, 4, 1, 1),
-        (1, 1),
-        None,
-        True,
-    ),
 }
 test_data_suite_tosa_bf16 = {
     "randn_double_scale_bf16": lambda: (

backends/arm/test/ops/test_upsample_nearest2d.py

@@ -41,7 +41,6 @@
     "rand_double_size": lambda: (torch.rand(2, 4, 8, 3), (16, 6), None, True),
     "rand_one_double_scale": lambda: (torch.rand(2, 4, 1, 1), None, 2.0, True),
     "rand_one_double_size": lambda: (torch.rand(2, 4, 1, 1), (2, 2), None, True),
-    "rand_one_same_scale": lambda: (torch.rand(2, 4, 1, 1), None, 1.0, True),
     "rand_one_same_size": lambda: (torch.rand(2, 4, 1, 1), (1, 1), None, True),
     # Can't compare outputs as the rounding when selecting the nearest pixel is
     # different between PyTorch and TOSA. Just check the legalization went well.

backends/arm/test/passes/test_remove_data_layout_noops.py

@@ -118,6 +118,89 @@ def test_keep_multi_input_concat():
     assert _count_target(result, exir_ops.edge.aten.cat.default) == 1
 
 
+def test_remove_identity_bilinear_upsample():
+    graph = Graph()
+    x = graph.placeholder("x")
+    x.meta["val"] = torch.ones(1, 4, 768, 384)
+    upsample = _call(
+        graph,
+        exir_ops.edge.aten.upsample_bilinear2d.vec,
+        (x, None, False, [1.0, 1.0]),
+        torch.ones(1, 4, 768, 384),
+    )
+    graph.output((upsample,))
+
+    result = _run_remove_noop(GraphModule(torch.nn.Module(), graph))
+
+    assert _count_target(result, exir_ops.edge.aten.upsample_bilinear2d.vec) == 0
+    assert result.graph.output_node().args[0][0].op == "placeholder"
+
+
+def test_remove_identity_nearest_upsample():
+    graph = Graph()
+    x = graph.placeholder("x")
+    x.meta["val"] = torch.ones(1, 4, 768, 384)
+    upsample = _call(
+        graph,
+        exir_ops.edge.aten.upsample_nearest2d.vec,
+        (x, None, [1.0, 1.0]),
+        torch.ones(1, 4, 768, 384),
+    )
+    graph.output((upsample,))
+
+    result = _run_remove_noop(GraphModule(torch.nn.Module(), graph))
+
+    assert _count_target(result, exir_ops.edge.aten.upsample_nearest2d.vec) == 0
+    assert result.graph.output_node().args[0][0].op == "placeholder"
+
+
+def test_keep_bilinear_upsample_with_rounded_identity_shape():
+    graph = Graph()
+    x = graph.placeholder("x")
+    x.meta["val"] = torch.ones(1, 4, 3, 3)
+    upsample = _call(
+        graph,
+        exir_ops.edge.aten.upsample_bilinear2d.vec,
+        (x, None, False, [1.1, 1.1]),
+        torch.ones(1, 4, 3, 3),
+    )
+    graph.output((upsample,))
+
+    result = _run_remove_noop(GraphModule(torch.nn.Module(), graph))
+
+    assert _count_target(result, exir_ops.edge.aten.upsample_bilinear2d.vec) == 1
+
+
+class _IdentityUpsampleModule(torch.nn.Module):
+    def forward(self, x):
+        return torch.nn.functional.interpolate(
+            x, scale_factor=1.0, mode="bilinear", align_corners=False
+        )
+
+
+def test_remove_identity_bilinear_upsample_backend_pipeline():
+    exported_program = export(
+        _IdentityUpsampleModule(), (torch.ones(1, 4, 768, 384),), strict=True
+    )
+    edge_program = to_edge(
+        exported_program,
+        compile_config=EdgeCompileConfig(_check_ir_validity=False),
+    ).exported_program()
+
+    assert (
+        _count_target(
+            edge_program.graph_module, exir_ops.edge.aten.upsample_bilinear2d.vec
+        )
+        == 1
+    )
+
+    graph_module = ArmPassManager(
+        TosaCompileSpec("TOSA-1.0+FP")
+    ).transform_to_backend_pipeline(edge_program, edge_program.graph_module)
+
+    assert _count_target(graph_module, exir_ops.backend.tosa.RESIZE.default) == 0
+
+
 def test_remove_full_slice_and_unused_shape_constants():
     graph_module = _tosa_data_layout_graph(
         exir_ops.backend.tosa.SLICE.default,