Arm backend: Add bf16 support for aten.index_select and aten.unfold_copy (#19751)

Follow-up to #17097, which added BF16 support to the TOSA GATHER op.
`aten.index_select` and `aten.unfold_copy` both lower via TOSA GATHER
but their support checks were not updated at the time.

In both decompositions(`DecomposeIndexSelectToGatherPass()` and
`DecomposeUnfoldToGatherPass()`),
the bf16 values tensor flows through dtype-agnostic reshape ops and
`tosa.GATHER`, which accepts `BF16`.
The support check was the only blocker.

| Op                  | bf16 before | bf16 after |
|---------------------|:-----------:|:----------:|
| `aten.gather`       | ✅          | ✅         |
| `aten.index.Tensor` | ✅          | ✅         |
| `aten.slice_copy`   | ✅          | ✅         |
| `aten.index_select` | ❌          | ✅         |
| `aten.unfold_copy`  | ❌          | ✅         |

Changes:
- `index_select_support.py`, `unfold_copy_support.py`: extend float
branch
to include `bfloat16`; add bf16 extension guard; update rejection
message.
- `test_index_select.py`, `test_unfold_copy.py`: add isolated
  `_tosa_FP_bf16` test functions using
  `TosaPipelineFP(..., tosa_extensions=["bf16"])`.

### Test plan

`test_index_select_tosa_FP_bf16` and `test_unfold_copy_tosa_FP_bf16`
exercise the bf16 path end-to-end through `TosaPipelineFP` with the bf16
extension enabled, following the same pattern of the existing
`test_slice_tensor_tosa_FP_bf16` from #17492

Author: vacu9708

backends/arm/operator_support/index_select_support.py

@@ -77,8 +77,16 @@ def is_node_tosa_supported(
                     f"{node.target}: dtype {values_dtype} requires INT profile.",
                 )
                 return False
-        # fp16/fp32: either FP profile, or INT profile (via quantization)
-        elif values_dtype in (torch.float16, torch.float32):
+        # fp16/fp32/bf16: either FP profile, or INT profile (via quantization)
+        elif values_dtype in (torch.float16, torch.float32, torch.bfloat16):
+            if values_dtype == torch.bfloat16 and not tosa_spec.support_extension(
+                "bf16"
+            ):
+                self.reporter.report_reject(
+                    node,
+                    f"{node.target}: dtype {values_dtype} requires bf16 extension.",
+                )
+                return False
             if not (tosa_spec.support_float() or tosa_spec.support_integer()):
                 self.reporter.report_reject(
                     node,
@@ -90,7 +98,7 @@ def is_node_tosa_supported(
             self.reporter.report_reject(
                 node,
                 f"{node.target}: unsupported values dtype {values_dtype}; "
-                "expected bool/int8/int16/int32/float16/float32.",
+                "expected bool/int8/int16/int32/float16/bfloat16/float32.",
             )
             return False
 

backends/arm/operator_support/unfold_copy_support.py

@@ -84,8 +84,16 @@ def is_node_tosa_supported(
                     f"{node.target}: dtype {values_dtype} requires INT profile.",
                 )
                 return False
-        # fp16/fp32: either FP profile, or INT profile (via quantization)
-        elif values_dtype in (torch.float16, torch.float32):
+        # fp16/fp32/bf16: either FP profile, or INT profile (via quantization)
+        elif values_dtype in (torch.float16, torch.float32, torch.bfloat16):
+            if values_dtype == torch.bfloat16 and not tosa_spec.support_extension(
+                "bf16"
+            ):
+                self.reporter.report_reject(
+                    node,
+                    f"{node.target}: dtype {values_dtype} requires bf16 extension.",
+                )
+                return False
             if not (tosa_spec.support_float() or tosa_spec.support_integer()):
                 self.reporter.report_reject(
                     node,
@@ -97,7 +105,7 @@ def is_node_tosa_supported(
             self.reporter.report_reject(
                 node,
                 f"{node.target}: unsupported values dtype {values_dtype}; "
-                "expected bool/int8/int16/int32/float16/float32.",
+                "expected bool/int8/int16/int32/float16/bfloat16/float32.",
             )
             return False
 

backends/arm/test/ops/test_index_select.py

@@ -61,6 +61,26 @@ def forward(self, input_: torch.Tensor, dim: int, index_: torch.Tensor):
         torch.tensor([3, 1], dtype=torch.int32),  # [W=2]
     ),
 }
+test_data_fp_bf16: dict[str, input_params] = {
+    # Rank-2: [K, C] -> index_select dim=0 => [W, C]
+    "test_bf16_rank2_dim0": (
+        torch.tensor(
+            [[0.5, 1.25, 2.5], [3.5, 4.25, 5.75], [6.5, 7.25, 8.75]],
+            dtype=torch.bfloat16,
+        ),  # [K=3, C=3]
+        0,
+        torch.tensor([2, 0], dtype=torch.int32),  # [W=2]
+    ),
+    # Rank-3: [N, K, C] -> index_select dim=-1 => [N, K, W]
+    "test_bf16_rank3_dim_neg1": (
+        torch.tensor(
+            [[[0.5, 1.5], [2.5, 3.5]], [[4.5, 5.5], [6.5, 7.5]]],
+            dtype=torch.bfloat16,
+        ),  # [N=2, K=2, C=2]
+        -1,
+        torch.tensor([1, 0], dtype=torch.int32),  # [W=2]
+    ),
+}
 
 # ---- INT profile: integer inputs + bool ----
 test_data_int: dict[str, input_params] = {
@@ -104,6 +124,18 @@ def test_index_select_tosa_FP(test_data: input_params):
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_fp_bf16)
+def test_index_select_tosa_FP_bf16(test_data: input_params):
+    pipeline = TosaPipelineFP[input_params](
+        IndexSelect(),
+        test_data,
+        aten_op=IndexSelect.aten_op,
+        exir_op=IndexSelect.exir_op,
+        tosa_extensions=["bf16"],
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", test_data_int | test_data_fp)
 def test_index_select_tosa_INT(test_data: input_params):
     # INT profile runs quantized, so we test both int inputs and float inputs here.

backends/arm/test/ops/test_unfold_copy.py

@@ -120,6 +120,18 @@ def forward(self, input_: torch.Tensor, dim_: int, size_: int, step_: int):
     ),
 }
 
+test_data_bf16: dict[str, input_params] = {
+    "test_bf16_2d_dim1": (
+        torch.tensor(
+            [[0.1, 0.2, 0.3, 0.4, 0.5], [1.1, 1.2, 1.3, 1.4, 1.5]],
+            dtype=torch.bfloat16,
+        ),  # [B=2, T=5]
+        1,
+        3,
+        2,  # U=(5-3)//2+1=2 -> [B=2, U=2, C=3]
+    ),
+}
+
 
 @common.parametrize("test_data", test_data_fp)
 def test_unfold_copy_tosa_FP(test_data: input_params):
@@ -132,6 +144,18 @@ def test_unfold_copy_tosa_FP(test_data: input_params):
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_bf16)
+def test_unfold_copy_tosa_FP_bf16(test_data: input_params):
+    pipeline = TosaPipelineFP[input_params](
+        UnfoldCopy(),
+        test_data,
+        aten_op=UnfoldCopy.aten_op,
+        exir_op=UnfoldCopy.exir_op,
+        tosa_extensions=["bf16"],
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", test_data_int | test_data_fp)
 def test_unfold_copy_tosa_INT(test_data: input_params):
     pipeline = TosaPipelineINT[input_params](