Arm backend: Share operator dtype support checks (#24495)

Signed-off-by: Sebastian Larsson <sebastian.larsson@arm.com>

Author: Sebastian-Larsson

backends/arm/BUCK

@@ -18,6 +18,7 @@ fbcode_target(
         "constants.py",
     ],
     deps = [
+        "//caffe2:torch",
         "//executorch/exir/dialects:lib",
     ],
 )

backends/arm/constants.py

@@ -5,6 +5,7 @@
 
 from typing import Any, cast, Final
 
+import torch
 from executorch.exir.dialects._ops import ops as exir_ops
 
 exir_ops = cast(Any, exir_ops)
@@ -30,6 +31,15 @@
 )
 PER_CHANNEL_QDQ_OPS: Final = (QUANT_PER_CHANNEL_OP, DEQUANT_PER_CHANNEL_OP)
 
+INT_DTYPES: Final = (torch.bool, torch.int8, torch.int16, torch.int32)
+FP_DTYPES: Final = (
+    torch.float16,
+    torch.float32,
+    torch.bfloat16,
+    torch.float8_e4m3fn,
+    torch.float8_e5m2,
+)
+
 NHWC_ORDER: Final = (0, 2, 3, 1)
 NHWC_INVERSE_ORDER: Final = (0, 3, 1, 2)
 NNHWC_ORDER: Final = (0, 1, 3, 4, 2)

backends/arm/operator_support/dtype_support.py

@@ -0,0 +1,73 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+"""Provide shared dtype checks for TOSA operator support."""
+
+import torch
+from executorch.backends.arm.constants import FP_DTYPES, INT_DTYPES
+from executorch.backends.arm.tosa import TosaSpecification
+from executorch.exir.backend.utils import WhyNoPartitionReporter
+
+
+def is_indexing_dtype_supported(
+    node: torch.fx.Node,
+    values_dtype: torch.dtype,
+    tosa_spec: TosaSpecification,
+    reporter: WhyNoPartitionReporter,
+) -> bool:
+    """Check value dtypes for indexing against TOSA profiles and extensions.
+
+    Supported integer dtypes and bool require the INT profile. Supported
+    floating-point dtypes allow the FP profile or the INT profile via
+    quantization. BF16 and FP8 also require their corresponding extensions.
+
+    Args:
+        node (torch.fx.Node): Operator node used in rejection messages.
+        values_dtype (torch.dtype): Dtype of the values being indexed.
+        tosa_spec (TosaSpecification): Active TOSA profiles and extensions.
+        reporter (WhyNoPartitionReporter): Reporter for rejection reasons.
+
+    Returns:
+        bool: True if the dtype is supported; otherwise, False.
+
+    """
+    if values_dtype in INT_DTYPES:
+        if not tosa_spec.support_integer():
+            reporter.report_reject(
+                node,
+                f"{node.target}: dtype {values_dtype} requires INT profile.",
+            )
+            return False
+    elif values_dtype in FP_DTYPES:
+        if values_dtype == torch.bfloat16 and not tosa_spec.support_extension("bf16"):
+            reporter.report_reject(
+                node,
+                f"{node.target}: dtype {values_dtype} requires bf16 extension.",
+            )
+            return False
+        if values_dtype == torch.float8_e4m3fn and not tosa_spec.support_extension(
+            "fp8e4m3"
+        ):
+            reporter.report_reject(
+                node,
+                f"{node.target}: dtype {values_dtype} requires fp8e4m3 extension.",
+            )
+            return False
+        if values_dtype == torch.float8_e5m2 and not tosa_spec.support_extension(
+            "fp8e5m2"
+        ):
+            reporter.report_reject(
+                node,
+                f"{node.target}: dtype {values_dtype} requires fp8e5m2 extension.",
+            )
+            return False
+    else:
+        reporter.report_reject(
+            node,
+            f"{node.target}: unsupported values dtype {values_dtype}; "
+            f"expected one of {INT_DTYPES + FP_DTYPES}.",
+        )
+        return False
+
+    return True

backends/arm/operator_support/gather_support.py

@@ -38,6 +38,9 @@
 from executorch.backends.arm._passes.prepare_gather_indices_pass import (
     is_safe_int32_to_int64_gather_boundary,
 )
+from executorch.backends.arm.operator_support.dtype_support import (
+    is_indexing_dtype_supported,
+)
 from executorch.backends.arm.operator_support.tosa_supported_operators import (
     register_tosa_support_check,
     SupportedTOSAOperatorCheck,
@@ -52,9 +55,10 @@ class GatherSupported(SupportedTOSAOperatorCheck):
 
     targets = [exir_ops.edge.aten.gather.default]
 
-    def is_node_tosa_supported(  # noqa: C901
+    def is_node_tosa_supported(
         self, node: fx.Node, tosa_spec: TosaSpecification
     ) -> bool:  # type: ignore[override, misc]
+        """Return True if the gather node passes the TOSA support checks."""
         if len(node.args) != 3:
             self.reporter.report_reject(
                 node,
@@ -111,62 +115,4 @@ def is_node_tosa_supported(  # noqa: C901
             )
             return False
 
-        # ---- values dtype ----
-        values_dtype = x_val.dtype
-        # ints (and bool via casts) require INT profile
-        if values_dtype in (torch.bool, torch.int8, torch.int16, torch.int32):
-            if not tosa_spec.support_integer():
-                self.reporter.report_reject(
-                    node,
-                    f"{node.target}: dtype {values_dtype} requires INT profile.",
-                )
-                return False
-        # fp16/fp32/bf16/fp8: either FP profile, or INT profile (via quantization)
-        elif values_dtype in (
-            torch.float16,
-            torch.float32,
-            torch.bfloat16,
-            torch.float8_e4m3fn,
-            torch.float8_e5m2,
-        ):
-            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 values_dtype == torch.float8_e4m3fn and not tosa_spec.support_extension(
-                "fp8e4m3"
-            ):
-                self.reporter.report_reject(
-                    node,
-                    f"{node.target}: dtype {values_dtype} requires fp8e4m3 extension.",
-                )
-                return False
-            if values_dtype == torch.float8_e5m2 and not tosa_spec.support_extension(
-                "fp8e5m2"
-            ):
-                self.reporter.report_reject(
-                    node,
-                    f"{node.target}: dtype {values_dtype} requires fp8e5m2 extension.",
-                )
-                return False
-            if not (tosa_spec.support_float() or tosa_spec.support_integer()):
-                self.reporter.report_reject(
-                    node,
-                    f"{node.target}: dtype {values_dtype} requires FP profile or "
-                    "INT profile (with quantization).",
-                )
-                return False
-        else:
-            self.reporter.report_reject(
-                node,
-                f"{node.target}: unsupported values dtype {values_dtype}; "
-                "expected bool/int8/int16/int32/float16/bfloat16/float32/"
-                "float8_e4m3fn/float8_e5m2.",
-            )
-            return False
-
-        return True
+        return is_indexing_dtype_supported(node, x_val.dtype, tosa_spec, self.reporter)

backends/arm/operator_support/index_select_support.py

@@ -18,6 +18,9 @@
 import torch
 import torch.fx as fx
 
+from executorch.backends.arm.operator_support.dtype_support import (
+    is_indexing_dtype_supported,
+)
 from executorch.backends.arm.operator_support.tosa_supported_operators import (
     register_tosa_support_check,
     SupportedTOSAOperatorCheck,
@@ -35,6 +38,7 @@ class IndexSelectSupported(SupportedTOSAOperatorCheck):
     def is_node_tosa_supported(
         self, node: fx.Node, tosa_spec: TosaSpecification
     ) -> bool:  # type: ignore[override, misc]
+        """Return True if the index_select node passes the support checks."""
         if len(node.args) != 3:
             self.reporter.report_reject(
                 node,
@@ -67,62 +71,4 @@ def is_node_tosa_supported(
             )
             return False
 
-        # Dtype constraints
-        values_dtype = x_val.dtype
-        # ints (and bool via casts) require INT profile
-        if values_dtype in (torch.bool, torch.int8, torch.int16, torch.int32):
-            if not tosa_spec.support_integer():
-                self.reporter.report_reject(
-                    node,
-                    f"{node.target}: dtype {values_dtype} requires INT profile.",
-                )
-                return False
-        # fp16/fp32/bf16/fp8: either FP profile, or INT profile (via quantization)
-        elif values_dtype in (
-            torch.float16,
-            torch.float32,
-            torch.bfloat16,
-            torch.float8_e4m3fn,
-            torch.float8_e5m2,
-        ):
-            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 values_dtype == torch.float8_e4m3fn and not tosa_spec.support_extension(
-                "fp8e4m3"
-            ):
-                self.reporter.report_reject(
-                    node,
-                    f"{node.target}: dtype {values_dtype} requires fp8e4m3 extension.",
-                )
-                return False
-            if values_dtype == torch.float8_e5m2 and not tosa_spec.support_extension(
-                "fp8e5m2"
-            ):
-                self.reporter.report_reject(
-                    node,
-                    f"{node.target}: dtype {values_dtype} requires fp8e5m2 extension.",
-                )
-                return False
-            if not (tosa_spec.support_float() or tosa_spec.support_integer()):
-                self.reporter.report_reject(
-                    node,
-                    f"{node.target}: dtype {values_dtype} requires FP profile or "
-                    "INT profile (with quantization).",
-                )
-                return False
-        else:
-            self.reporter.report_reject(
-                node,
-                f"{node.target}: unsupported values dtype {values_dtype}; "
-                "expected bool/int8/int16/int32/float16/bfloat16/float32/"
-                "float8_e4m3fn/float8_e5m2.",
-            )
-            return False
-
-        return True
+        return is_indexing_dtype_supported(node, x_val.dtype, tosa_spec, self.reporter)

backends/arm/operator_support/index_tensor_support.py

@@ -16,6 +16,9 @@
 import torch.fx as fx
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.arm.common.type import ensure_type
+from executorch.backends.arm.operator_support.dtype_support import (
+    is_indexing_dtype_supported,
+)
 from executorch.backends.arm.operator_support.tosa_supported_operators import (
     register_tosa_support_check,
     SupportedTOSAOperatorCheck,
@@ -178,58 +181,6 @@ def is_node_tosa_supported(
         if not self._are_shapes_supported(node, input_val, index_vals):
             return False
 
-        values_dtype = input_val.dtype
-        if values_dtype in (torch.bool, torch.int8, torch.int16, torch.int32):
-            if not tosa_spec.support_integer():
-                self.reporter.report_reject(
-                    node,
-                    f"{node.target}: dtype {values_dtype} requires INT profile.",
-                )
-                return False
-        elif values_dtype in (
-            torch.float16,
-            torch.float32,
-            torch.bfloat16,
-            torch.float8_e4m3fn,
-            torch.float8_e5m2,
-        ):
-            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 values_dtype == torch.float8_e4m3fn and not tosa_spec.support_extension(
-                "fp8e4m3"
-            ):
-                self.reporter.report_reject(
-                    node,
-                    f"{node.target}: dtype {values_dtype} requires fp8e4m3 extension.",
-                )
-                return False
-            if values_dtype == torch.float8_e5m2 and not tosa_spec.support_extension(
-                "fp8e5m2"
-            ):
-                self.reporter.report_reject(
-                    node,
-                    f"{node.target}: dtype {values_dtype} requires fp8e5m2 extension.",
-                )
-                return False
-            if not (tosa_spec.support_float() or tosa_spec.support_integer()):
-                self.reporter.report_reject(
-                    node,
-                    f"{node.target}: dtype {values_dtype} requires FP profile or "
-                    "INT profile (with quantization).",
-                )
-                return False
-        else:
-            self.reporter.report_reject(
-                node,
-                f"{node.target}: unsupported values dtype {values_dtype}; "
-                "expected bool/int8/int16/int32/float16/bfloat16/float32/float8_e4m3fn/float8_e5m2.",
-            )
-            return False
-
-        return True
+        return is_indexing_dtype_supported(
+            node, input_val.dtype, tosa_spec, self.reporter
+        )