Arm backend: Support high rank index tensor (#17951)

  - Drop the rank>=4 rejection from index.Tensor TOSA support checks
  - Add rank-4 and rank-5 index tensor test cases
  - Note: rank>=4 support has been covered by the index.Tensor
    refactor, removing the need for special handling in
    ToTosaMemoryFormatPass

Change-Id: Ief40942a94040c02e54c7f276eecd660d571e46d

Arm backend: Refactor aten.index.tensor suppport
  - Decompose edge.aten.index.Tensor via backend tosa.GATHER + shape ops
  - Remove the index.Tensor node visitor
  - Reorder AccumulateIndexPutPass (generates index.Tensor) and
    DecomposeSliceScatterPass (may generate index_put) before
    DecomposeIndexTensorToGatherPass
  - Correct index_put tests to not test int inputs under FP-only profile

Change-Id: I5cfc7c110f0074463043ef1cb61165cc784a4db2


cc @digantdesai @SS-JIA @freddan80 @per @zingo @oscarandersson8218
@mansnils @Sebastian-Larsson @robell

---------

Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>
Co-authored-by: Erik Lundell <erik.lundell@arm.com>

Author: YufengShi-dudu

backends/arm/_passes/__init__.py

@@ -55,6 +55,9 @@
 from .decompose_index_select_to_gather_pass import (  # noqa
     DecomposeIndexSelectToGatherPass,
 )
+from .decompose_index_tensor_to_gather_pass import (  # noqa
+    DecomposeIndexTensorToGatherPass,
+)
 from .decompose_int16_activation_conv_pass import (  # noqa
     DecomposeConvWithInt16ActivationPass,
 )

backends/arm/_passes/accumulate_index_put_pass.py

@@ -7,6 +7,10 @@
 import torch
 
 from executorch.backends.arm._passes import ArmPass
+from executorch.backends.arm._passes.decompose_index_tensor_to_gather_pass import (
+    DecomposeIndexTensorToGatherPass,
+)
+from executorch.backends.arm._passes.rewrite_index_put_pass import RewriteIndexPutPass
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
@@ -33,7 +37,10 @@ class AccumulateIndexPutPass(ArmPass):
     for the index_put op.
     """
 
-    _passes_required_after: Set[Type[ExportPass]] = set()
+    _passes_required_after: Set[Type[ExportPass]] = {
+        DecomposeIndexTensorToGatherPass,
+        RewriteIndexPutPass,
+    }
 
     def call_operator(self, op, args, kwargs, meta):
         if op not in (aten_ops + edge_ops) or not self.allowed_to_transform(meta):

backends/arm/_passes/arm_pass_manager.py

@@ -60,6 +60,7 @@
     DecomposeGroupedConvPass,
     DecomposeGroupNormPass,
     DecomposeIndexSelectToGatherPass,
+    DecomposeIndexTensorToGatherPass,
     DecomposeIntPowPass,
     DecomposeLayerNormPass,
     DecomposeLeakyReLUPass,
@@ -307,6 +308,9 @@ def _tosa_pipeline(
                 DecomposeEmbeddingPass(),
                 DecomposeIndexSelectToGatherPass(),
                 DecomposeStridedSliceCopyPass(),
+                DecomposeSliceScatterPass(),
+                AccumulateIndexPutPass(),
+                DecomposeIndexTensorToGatherPass(),
                 Conv1dUnsqueezePass(),
             ]
         )
@@ -329,8 +333,6 @@ def _tosa_pipeline(
         # Node transformation passes (post scalar-removal)
         self.add_passes(
             [
-                DecomposeSliceScatterPass(),
-                AccumulateIndexPutPass(),
                 RewriteIndexPutPass(),
                 RewriteBoolBitwiseToLogicalPass(),
                 DecomposeRemainderPass(),

backends/arm/_passes/arm_pass_utils.py

@@ -21,6 +21,7 @@
     _get_control_flow_submodules,
     get_control_flow_submodules,
 )
+from executorch.exir.pass_base import NodeMetadata
 
 from torch._export.utils import (
     get_buffer,
@@ -202,6 +203,14 @@ def insert_q_dq_pair(
     return dq
 
 
+def meta_without_qparams(meta: NodeMetadata) -> NodeMetadata:
+    """Return a copy of NodeMetadata with input/output qparams cleared."""
+    plain_meta_dict = dict(meta.data)
+    plain_meta_dict["input_qparams"] = {}
+    plain_meta_dict["output_qparams"] = {}
+    return NodeMetadata(plain_meta_dict)
+
+
 def get_first_fake_tensor(node: torch.fx.Node) -> FakeTensor:
     """Returns a FakeTensor from the meta field of 'node'.