Arm backend: Support high-rank index.Tensor indices

  - 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
Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>

Author: YufengShi-dudu

backends/arm/operator_support/index_tensor_support.py

@@ -4,8 +4,8 @@
 # LICENSE file in the root directory of this source tree.
 """Provide TOSA support checks for ``aten.index.Tensor``.
 
-Reject unsupported patterns such as high-rank index tensors, front-positioned
-slice/ellipsis/None markers, and cases that exceed ``int32`` element limits.
+Reject unsupported patterns such as front-positioned slice/ellipsis/None
+markers and cases that exceed ``int32`` element limits.
 
 """
 
@@ -30,18 +30,12 @@ class IndexTensorSupported(SupportedTOSAOperatorCheck):
     This support check is intended to prevent the partitioning of
     currently unsupported usages of the index.Tensor operator.
 
-    1. Usages where indexing tensors are of rank 4 or higher.
-        This is due to the AnnotateChannelsLastDimOrder pass and
-        the rarity of such operation.
-        Support is possible but would require further changes to the above
-        pass which can be added at such a time as is necessary.
-
-    2. Usages where slice, ellipsis or None are present before an indexing tensor:
+    1. Usages where slice, ellipsis or None are present before an indexing tensor:
         t[{start}:{end}, indexTensor] - slicing
         t[None, indexTensor] - unsqueeze
         t[..., indexTensor] - ellipsis
 
-    3. Usages where the value tensor contains more than int32.max elements
+    2. Usages where the value tensor contains more than int32.max elements
         This is due to int32 TOSA limitation and the fact that we flatten out
         and accumulate all index tensors.
         As such to avoid overflow we reject lowering of this operator if it is
@@ -115,13 +109,12 @@ def is_node_tosa_supported(
 
         Enforces the following constraints:
         - No ``None`` (unsqueeze), slice, or ellipsis before an indexing tensor.
-        - Indexing tensors have rank <= 3.
         - The value tensor element count fits in ``int32``.
 
         """
         indices = node.args[1]
         for index in indices:  # type: ignore[union-attr]
-            # Usage 2 guard
+            # Usage 1 guard
             if index is None:
                 self.reporter.report_reject(
                     node,
@@ -132,17 +125,7 @@ def is_node_tosa_supported(
                 )
                 return False
 
-            # Usage 1 guard
-            index = ensure_type(torch.fx.Node, index)
-            fake_tensor = get_first_fake_tensor(index)
-            if len(fake_tensor.size()) > 3:
-                self.reporter.report_reject(
-                    node,
-                    ("Indexing tensors of rank >= 4 is not supported."),
-                )
-                return False
-
-        # Usage 3 guard
+        # Usage 2 guard
         input_node = ensure_type(torch.fx.Node, node.args[0])
         total_vals = math.prod(get_first_fake_tensor(input_node).shape)
         if total_vals > torch.iinfo(torch.int32).max:

backends/arm/test/ops/test_index_tensor.py

@@ -318,6 +318,28 @@ class IndexTensor(torch.nn.Module):
                 torch.randint(2, size=(15,), dtype=torch.int32),
             ),
         ),
+        "test_1d_rank4_index": (
+            torch.rand(12),
+            (torch.randint(12, size=(1, 2, 1, 3), dtype=torch.int32),),
+        ),
+        "test_2d_rank4_broadcastable_indices": (
+            torch.rand(4, 6),
+            (
+                torch.randint(4, size=(1, 2, 1, 1), dtype=torch.int32),
+                torch.randint(6, size=(1, 1, 3, 1), dtype=torch.int32),
+            ),
+        ),
+        "test_1d_high_rank_index": (
+            torch.rand(24),
+            (torch.randint(24, size=(1, 1, 2, 1, 3), dtype=torch.int32),),
+        ),
+        "test_2d_high_rank_broadcastable_indices": (
+            torch.rand(4, 5),
+            (
+                torch.randint(4, size=(1, 2, 1, 1, 1), dtype=torch.int32),
+                torch.randint(5, size=(1, 1, 3, 1, 1), dtype=torch.int32),
+            ),
+        ),
     }
     test_data_bf16: dict[input_params] = {
         "test_2d_1_idx_bf16": (