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
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

@@ -54,6 +54,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

@@ -59,6 +59,7 @@
     DecomposeGroupedConvPass,
     DecomposeGroupNormPass,
     DecomposeIndexSelectToGatherPass,
+    DecomposeIndexTensorToGatherPass,
     DecomposeIntPowPass,
     DecomposeLayerNormPass,
     DecomposeLeakyReLUPass,
@@ -303,6 +304,9 @@ def _tosa_pipeline(
                 DecomposeEmbeddingPass(),
                 DecomposeIndexSelectToGatherPass(),
                 DecomposeStridedSliceCopyPass(),
+                DecomposeSliceScatterPass(),
+                AccumulateIndexPutPass(),
+                DecomposeIndexTensorToGatherPass(),
                 Conv1dUnsqueezePass(),
             ]
         )
@@ -325,8 +329,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

@@ -17,6 +17,7 @@
 from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
+from executorch.exir.pass_base import NodeMetadata
 
 from torch._export.utils import (
     get_buffer,
@@ -197,6 +198,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'.