Arm backend: Add TOSA dialect conversion ops (#20786)

Add fake implementations for CAST and move RESCALE and
CAST_TO_BLOCK_SCALED into conversion_ops.py.



cc @digantdesai @freddan80 @per @zingo @mansnils @Sebastian-Larsson
@robell @rascani

Signed-off-by: Oscar Andersson <oscar.andersson@arm.com>

Author: oscarandersson8218

backends/arm/operator_support/TARGETS

@@ -7,6 +7,7 @@ runtime.python_library(
         "//executorch/backends/arm:ao_ext",
         "//executorch/backends/arm:constants",
         "//executorch/backends/arm/_passes:passes",
+        "//executorch/backends/arm/tosa:cast_support",
         "//executorch/backends/arm/tosa:resize_utils",
         "//executorch/backends/arm/tosa:tosa",
         "//executorch/backends/transforms:remove_getitem_op",

backends/arm/operator_support/to_dim_order_copy_support.py

@@ -10,7 +10,6 @@
 
 """
 
-import copy
 import logging
 
 import torch
@@ -22,123 +21,23 @@
     SupportedTOSAOperatorCheck,
 )
 from executorch.backends.arm.tosa import TosaSpecification
+from executorch.backends.arm.tosa.cast_support import (
+    supported_to_dim_order_casts,
+    SupportedCastMap,
+)
 from executorch.exir.dialects._ops import ops as exir_ops
 
 logger = logging.getLogger(__name__)
 
-SupportedTypeDict = dict[torch.dtype, list[torch.dtype]]
-
 
 @register_tosa_support_check
 class ToCopySupported(SupportedTOSAOperatorCheck):
-    """Provide TOSA support check for ``_to_dim_order_copy``.
-
-    Attributes:
-        SUPPORTED_INT_PROFILE_DTYPES (dict[torch.dtype, list[torch.dtype]]):
-            Allowed output dtypes for each integer input dtype.
-        SUPPORTED_FP_PROFILE_DTYPES (dict[torch.dtype, list[torch.dtype]]):
-            Allowed output dtypes for each floating input dtype.
-
-    """
+    """Provide TOSA support check for ``_to_dim_order_copy``."""
 
     targets = [
         exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
     ]
 
-    @staticmethod
-    def _merge_supported_types(
-        dtypes1: SupportedTypeDict,
-        dtypes2: SupportedTypeDict,
-    ) -> SupportedTypeDict:
-        """Return a merged mapping of supported dtype transitions.
-
-        Args:
-            dtypes1 (dict[torch.dtype, list[torch.dtype]]): Base mapping.
-            dtypes2 (dict[torch.dtype, list[torch.dtype]]): Mapping to merge in.
-
-        Returns:
-            dict[torch.dtype, list[torch.dtype]]: Combined mapping.
-
-        """
-        merged_dtypes = copy.deepcopy(
-            dtypes1
-        )  # Use deepcopy to avoid unintentionally modifying SUPPORTED_INT_PROFILE_DTYPES
-        for k, v in dtypes2.items():
-            merged_dtypes[k] = merged_dtypes.get(k, []) + v
-        return merged_dtypes
-
-    SUPPORTED_INT_PROFILE_DTYPES: SupportedTypeDict = {
-        torch.bool: [torch.bool, torch.int8, torch.int16, torch.int32],
-        torch.int8: [torch.bool, torch.int8, torch.int16, torch.int32],
-        torch.int16: [torch.bool, torch.int8, torch.int16, torch.int32],
-        torch.int32: [torch.bool, torch.int8, torch.int16, torch.int32],
-        torch.int64: [torch.bool, torch.int8, torch.int16, torch.int32],
-    }
-    SUPPORTED_FP_PROFILE_DTYPES: SupportedTypeDict = {
-        torch.int8: [torch.int8, torch.float16, torch.float32],
-        torch.int16: [torch.int16, torch.float16, torch.float32],
-        torch.int32: [torch.int32, torch.float16, torch.float32],
-        # INT64 inputs to casts *should* be ok, since they should be rejected by
-        # CheckInt64InputsAndOutputs if the cast can't be done AOT.
-        torch.int64: [
-            torch.int8,
-            torch.int16,
-            torch.int32,
-            torch.float16,
-            torch.float32,
-        ],
-        torch.float16: [
-            torch.int8,
-            torch.int16,
-            torch.int32,
-            torch.float16,
-            torch.float32,
-        ],
-        torch.float32: [
-            torch.float32,
-            torch.int8,
-            torch.int16,
-            torch.int32,
-            torch.float16,
-            torch.float32,
-        ],
-    }
-    SUPPORTED_INT_FP_PROFILE_DTYPES: SupportedTypeDict = {
-        torch.bool: [torch.float32],
-    }
-    SUPPORTED_BF16_EXTENSION_DTYPES: SupportedTypeDict = {
-        torch.int8: [torch.bfloat16],
-        torch.int16: [torch.bfloat16],
-        torch.int32: [torch.bfloat16],
-        torch.int64: [torch.bfloat16],
-        torch.float32: [torch.bfloat16],
-        torch.bfloat16: [
-            torch.int8,
-            torch.int16,
-            torch.int32,
-            torch.bfloat16,
-            torch.float32,
-        ],
-    }
-    SUPPORTED_FP8E4M3_EXTENSION_DTYPES: SupportedTypeDict = {
-        torch.float16: [torch.float8_e4m3fn],
-        torch.float32: [torch.float8_e4m3fn],
-        torch.float8_e4m3fn: [torch.float16, torch.float32],
-    }
-    SUPPORTED_FP8E5M2_EXTENSION_DTYPES: SupportedTypeDict = {
-        torch.float16: [torch.float8_e5m2],
-        torch.float32: [torch.float8_e5m2],
-        torch.float8_e5m2: [torch.float16, torch.float32],
-    }
-    SUPPORTED_BF16_FP8E4M3_EXTENSION_DTYPES: SupportedTypeDict = {
-        torch.bfloat16: [torch.float8_e4m3fn],
-        torch.float8_e4m3fn: [torch.bfloat16],
-    }
-    SUPPORTED_BF16_FP8E5M2_EXTENSION_DTYPES: SupportedTypeDict = {
-        torch.bfloat16: [torch.float8_e5m2],
-        torch.float8_e5m2: [torch.bfloat16],
-    }
-
     @staticmethod
     def _is_quantized_identity_cast(node: torch.fx.Node) -> bool:
         for user in node.users:
@@ -153,7 +52,7 @@ def _is_quantized_identity_cast(node: torch.fx.Node) -> bool:
                 return False
         return True
 
-    def is_node_tosa_supported(  # noqa: C901
+    def is_node_tosa_supported(
         self, node: fx.Node, tosa_spec: TosaSpecification
     ) -> bool:
         """Return True if the node is supported by TOSA.
@@ -163,43 +62,7 @@ def is_node_tosa_supported(  # noqa: C901
         input dtype.
 
         """
-        supported_dtypes: SupportedTypeDict = {}
-        if tosa_spec.support_integer():
-            supported_dtypes = self._merge_supported_types(
-                self.SUPPORTED_INT_PROFILE_DTYPES, supported_dtypes
-            )
-        if tosa_spec.support_float():
-            supported_dtypes = self._merge_supported_types(
-                self.SUPPORTED_FP_PROFILE_DTYPES, supported_dtypes
-            )
-        if tosa_spec.support_integer() and tosa_spec.support_float():
-            supported_dtypes = self._merge_supported_types(
-                self.SUPPORTED_INT_FP_PROFILE_DTYPES, supported_dtypes
-            )
-        if tosa_spec.support_extension("bf16"):
-            supported_dtypes = self._merge_supported_types(
-                self.SUPPORTED_BF16_EXTENSION_DTYPES, supported_dtypes
-            )
-        if tosa_spec.support_extension("fp8e4m3"):
-            supported_dtypes = self._merge_supported_types(
-                self.SUPPORTED_FP8E4M3_EXTENSION_DTYPES, supported_dtypes
-            )
-        if tosa_spec.support_extension("fp8e5m2"):
-            supported_dtypes = self._merge_supported_types(
-                self.SUPPORTED_FP8E5M2_EXTENSION_DTYPES, supported_dtypes
-            )
-        if tosa_spec.support_extension("bf16") and tosa_spec.support_extension(
-            "fp8e4m3"
-        ):
-            supported_dtypes = self._merge_supported_types(
-                self.SUPPORTED_BF16_FP8E4M3_EXTENSION_DTYPES, supported_dtypes
-            )
-        if tosa_spec.support_extension("bf16") and tosa_spec.support_extension(
-            "fp8e5m2"
-        ):
-            supported_dtypes = self._merge_supported_types(
-                self.SUPPORTED_BF16_FP8E5M2_EXTENSION_DTYPES, supported_dtypes
-            )
+        supported_dtypes: SupportedCastMap = supported_to_dim_order_casts(tosa_spec)
 
         if len(node.all_input_nodes) != 1:
             self.reporter.report_reject(
@@ -221,7 +84,6 @@ def is_node_tosa_supported(  # noqa: C901
             )
             return False
 
-        # Check input type
         input_dtype = input_val.dtype
         if input_dtype not in supported_dtypes:
             self.reporter.report_reject(
@@ -230,7 +92,6 @@ def is_node_tosa_supported(  # noqa: C901
             )
             return False
 
-        # Check output type
         output_val = node.meta["val"]
         if not isinstance(output_val, torch._subclasses.FakeTensor):
             self.reporter.report_reject(
@@ -242,10 +103,7 @@ def is_node_tosa_supported(  # noqa: C901
             )
             return False
         if output_val.dtype not in supported_dtypes[input_dtype]:
-            if (
-                tosa_spec.support_integer()
-                and ToCopySupported._is_quantized_identity_cast(node)
-            ):
+            if tosa_spec.support_integer() and self._is_quantized_identity_cast(node):
                 return True
             self.reporter.report_reject(
                 node,

backends/arm/test/misc/tosa_dialect/test_tosa_dialect_cast_to_block_scaled.py

@@ -7,7 +7,6 @@
 import torch
 from executorch.backends.arm.ao_ext.mxfp import mxfp_dtype_to_str
 from executorch.backends.arm.tosa.dialect.lib import TosaValueError
-from executorch.backends.arm.tosa.dialect.ops import cast_to_block_scaled  # noqa: F401
 from executorch.backends.arm.tosa.specification import (
     TosaLoweringContext,
     TosaSpecification,

backends/arm/test/misc/tosa_dialect/test_tosa_dialect_conversion_ops.py

@@ -0,0 +1,35 @@
+# 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.
+
+import executorch.backends.arm.tosa.dialect  # noqa: F401
+import pytest
+import torch
+from executorch.backends.arm.tosa.dialect.lib import TosaValueError
+from executorch.backends.arm.tosa.specification import (
+    TosaLoweringContext,
+    TosaSpecification,
+)
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+
+def test_cast_tosa_int() -> None:
+    tosa_spec = TosaSpecification.create_from_string("TOSA-1.0+INT")
+
+    with TosaLoweringContext(tosa_spec), FakeTensorMode() as mode:
+        input_tensor = mode.from_tensor(torch.ones((2, 3), dtype=torch.int8))
+        output = exir_ops.backend.tosa.CAST.default(input_tensor, torch.int32)
+
+    assert output.dtype == torch.int32
+    assert tuple(output.shape) == (2, 3)
+
+
+def test_cast_rejects_unsupported_profile_dtype_pair() -> None:
+    tosa_spec = TosaSpecification.create_from_string("TOSA-1.0+INT")
+
+    with TosaLoweringContext(tosa_spec), FakeTensorMode() as mode:
+        input_tensor = mode.from_tensor(torch.ones((2, 3), dtype=torch.int8))
+        with pytest.raises(TosaValueError, match="Unsupported CAST"):
+            exir_ops.backend.tosa.CAST.default(input_tensor, torch.float32)

backends/arm/tosa/BUCK

@@ -10,6 +10,17 @@ fbcode_target(_kind = runtime.python_library,
     visibility = ["PUBLIC"],
 )
 
+fbcode_target(_kind = runtime.python_library,
+    name = "cast_support",
+    srcs = [
+        "cast_support.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        ":specification",
+    ],
+)
+
 fbcode_target(_kind = runtime.python_library,
     name = "mapping",
     srcs = [