From aeacfa5112b7f1d9b1fcb13ecb826e36e5cf21d0 Mon Sep 17 00:00:00 2001
From: Sebastian Larsson <sebastian.larsson@arm.com>
Date: Tue, 24 Jun 2025 16:05:19 +0200
Subject: [PATCH] Arm backend: Add U55 operator check for View

In order for the view operator to be compatible with the channels-last
format of TosaBackend, transposes may need to be inserted before and
after the view op. The transose operator comes with some constraints on
U55. One such constraint is that the product of the axes should not
exceed 65536. With this patch a check if implemented for the view
operator to prevent such operators from being delegated.

Change-Id: Ic6259899fecc157452bed48ea82096089e125789
Signed-off-by: Sebastian Larsson <sebastian.larsson@arm.com>
---
 .../arm/operator_support/ethos_u55_support.py | 63 +++++++++++++++++++
 .../tosa_supported_operators.py               |  2 +
 backends/arm/test/ops/test_view.py            | 20 ++++++
 3 files changed, 85 insertions(+)

diff --git a/backends/arm/operator_support/ethos_u55_support.py b/backends/arm/operator_support/ethos_u55_support.py
index 372dab4c363..5a705acd877 100644
--- a/backends/arm/operator_support/ethos_u55_support.py
+++ b/backends/arm/operator_support/ethos_u55_support.py
@@ -174,6 +174,69 @@ def is_node_supported(
 shape_t = list[int]
 
 
+class EthosU55ViewCheck(OperatorSupportBase):
+
+    def __init__(self, reporter: WhyNoPartitionReporter):
+        super().__init__()
+        self.reporter = reporter
+
+    def axes_product(self, nhwc_shape: shape_t) -> int:
+        product = 1
+        for axes in nhwc_shape:
+            product *= axes
+        return product
+
+    # TODO: Extend this check to comply with u55 restrictions
+    def is_node_supported(
+        self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
+    ) -> bool:
+        """
+        Check whether a given view node is supported on U55.
+
+        Currently only checks dtypes and product of axes.
+
+        It is not the view operator itself that is not supported on U55. In order for the
+        view operator to be compatible with the channels-last format of TosaBackend,
+        transposes may need to be inserted before and after the view op. If that happens
+        and that transpose operator does not adhere to the limitations then it will
+        result in the following error:
+
+            CPU performance estimation for "Transpose" not implemented.
+            ...
+            CPU operations are not supported for GraphAPI input
+
+        Args:
+            node: The FX node representing the view_copy operator.
+
+        Returns:
+            False if the operator is not support and True if it is supported.
+        """
+        if not node.target == exir_ops.edge.aten.view_copy.default:
+            return True
+
+        shape = list(get_first_fake_tensor(node).shape)
+        dtype = _try_determine_dtype(node)
+        permutation = list(typing.cast(list[int], node.args[1]))
+
+        rank = len(shape)
+        if rank > 4:
+            if dtype == torch.int32:
+                self.reporter.report_reject(
+                    node, f"No support for {permutation=} in int32."
+                )
+                return False
+
+        if dtype in (torch.int8, torch.int16):
+            if self.axes_product(shape) > 65536:
+                self.reporter.report_reject(
+                    node,
+                    f"No support for {shape=}, {dtype=}. Product of axes must be <65536",
+                )
+                return False
+
+        return True
+
+
 class EthosU55TransposeCheck(OperatorSupportBase):
 
     def __init__(self, reporter: WhyNoPartitionReporter):
diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py
index cdb27b7c31e..14336ca83ce 100644
--- a/backends/arm/operator_support/tosa_supported_operators.py
+++ b/backends/arm/operator_support/tosa_supported_operators.py
@@ -23,6 +23,7 @@
     EthosU55DtypeSupport,
     EthosU55NotSupported,
     EthosU55TransposeCheck,
+    EthosU55ViewCheck,
 )
 from executorch.backends.arm.tosa_quant_utils import dq_ops, q_ops
 from executorch.backends.arm.tosa_specification import TosaSpecification
@@ -133,6 +134,7 @@ def tosa_support_factory(
         negative_checks.append(EthosU55NotSupported(reporter))
         negative_checks.append(EthosU55DtypeSupport(reporter))
         negative_checks.append(EthosU55TransposeCheck(reporter))
+        negative_checks.append(EthosU55ViewCheck(reporter))
 
     return chain(
         reporter.wrap_check(
diff --git a/backends/arm/test/ops/test_view.py b/backends/arm/test/ops/test_view.py
index 47e9627645d..fc780b1d32c 100644
--- a/backends/arm/test/ops/test_view.py
+++ b/backends/arm/test/ops/test_view.py
@@ -15,6 +15,7 @@
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineBI,
     EthosU85PipelineBI,
+    OpNotSupportedPipeline,
     TosaPipelineBI,
     TosaPipelineMI,
 )
@@ -44,6 +45,10 @@ class View(torch.nn.Module):
         "rand_4d_2_4_same": lambda: (torch.rand(2, 3, 2, 3), (2, 3, 3, 2)),
     }
 
+    rank_product_too_large = {
+        "rand_4d_large": lambda: (torch.rand(1, 49, 16, 128), (1, 16, 49, 128)),
+    }
+
     def __init__(self, new_shape):
         super().__init__()
         self.new_shape = new_shape
@@ -104,6 +109,21 @@ def test_view_u55_BI(test_data: Tuple):
     pipeline.run()
 
 
+@common.parametrize("test_data", View.rank_product_too_large, xfails=xfails)
+@common.XfailIfNoCorstone300
+def test_view_u55_BI_not_delegated(test_data: Tuple):
+    test_tensor, new_shape = test_data()
+    pipeline = OpNotSupportedPipeline[input_t1](
+        View(new_shape),
+        (test_tensor,),
+        {"executorch_exir_dialects_edge__ops_aten_view_copy": 1},
+        n_expected_delegates=0,
+        quantize=True,
+        u55_subset=True,
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", View.needs_transpose_tests, xfails=xfails)
 @common.XfailIfNoCorstone320
 def test_view_u85_BI(test_data: Tuple):