Arm backend: Add TOSA elementwise ternerary node visitors

backends/arm/_passes/aten_to_tosa_ternary.py

@@ -0,0 +1,25 @@
+# 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.
+
+from executorch.backends.transforms.aten_to_dialect_pass import (
+    AtenToDialectPass,
+    DialectNodeSpec,
+)
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch.fx import Node
+
+
+def rewrite_ternary_operator(
+    node: Node, pass_: AtenToDialectPass
+) -> DialectNodeSpec | None:
+    match node.target:
+        case exir_ops.edge.aten.where.self:
+            return DialectNodeSpec(
+                exir_ops.backend.tosa.SELECT.default,
+                node.args,
+                dict(node.kwargs),
+            )
+        case _:
+            return None

backends/arm/_passes/exir_to_tosa_pass.py

@@ -18,6 +18,9 @@
     rewrite_rfft2,
     rewrite_unary_operator,
 )
+from executorch.backends.arm._passes.aten_to_tosa_ternary import (
+    rewrite_ternary_operator,
+)
 from executorch.backends.transforms.aten_to_dialect_pass import (
     AtenToDialectPass,
     DialectNodeSpec,
@@ -123,6 +126,15 @@ def _get_activation_replacement(
     return get_activation_replacement(node, pass_)
 
 
+@register_dialect_substitutions(
+    exir_ops.edge.aten.where.self,
+)
+def _get_ternary_replacement(
+    node: Node, pass_: AtenToDialectPass
+) -> DialectNodeSpec | None:
+    return rewrite_ternary_operator(node, pass_)
+
+
 @register_dialect_substitutions(
     exir_ops.edge.aten.cat.default,
     exir_ops.edge.aten.flip.default,

backends/arm/operators/op_where.py

@@ -23,7 +23,7 @@
 
 @register_node_visitor
 class WhereVisitor(NodeVisitor):
-    target = "aten.where.self"
+    target = "tosa.SELECT.default"
 
     def __init__(self, *args):
         super().__init__(*args)

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

@@ -0,0 +1,149 @@
+# 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.dialect.ops_registration import (
+    get_registered_tosa_ops,
+)
+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 _to_fake(mode: FakeTensorMode, *values):
+    return [
+        mode.from_tensor(value) if isinstance(value, torch.Tensor) else value
+        for value in values
+    ]
+
+
+@pytest.mark.parametrize(
+    (
+        "spec",
+        "condition",
+        "input1",
+        "input2",
+        "expected_shape",
+        "expected_dtype",
+    ),
+    [
+        pytest.param(
+            "TOSA-1.1+INT",
+            torch.randint(0, 2, (2, 1, 3), dtype=torch.bool),
+            torch.randint(-8, 8, (1, 4, 3), dtype=torch.int32),
+            torch.randint(-8, 8, (1, 4, 3), dtype=torch.int32),
+            (2, 4, 3),
+            torch.int32,
+            id="int32",
+        ),
+        pytest.param(
+            "TOSA-1.1+FP",
+            torch.randint(0, 2, (2, 1, 3), dtype=torch.bool),
+            torch.randn((1, 4, 3), dtype=torch.float32),
+            torch.randn((1, 4, 3), dtype=torch.float32),
+            (2, 4, 3),
+            torch.float32,
+            id="fp32",
+        ),
+        pytest.param(
+            "TOSA-1.1+FP",
+            torch.randint(0, 2, (2, 1, 3), dtype=torch.bool),
+            torch.randint(0, 2, (1, 4, 3), dtype=torch.bool),
+            torch.randint(0, 2, (1, 4, 3), dtype=torch.bool),
+            (2, 4, 3),
+            torch.bool,
+            id="bool",
+        ),
+    ],
+)
+def test_tosa_select(
+    spec: str,
+    condition: torch.Tensor,
+    input1: torch.Tensor,
+    input2: torch.Tensor,
+    expected_shape: tuple[int, ...],
+    expected_dtype: torch.dtype,
+) -> None:
+    with TosaLoweringContext(
+        TosaSpecification.create_from_string(spec)
+    ), FakeTensorMode() as mode:
+        output = exir_ops.backend.tosa.SELECT.default(
+            *_to_fake(mode, condition, input1, input2)
+        )
+
+    assert output.dtype == expected_dtype
+    assert tuple(output.shape) == expected_shape
+
+
+@pytest.mark.parametrize("spec", ["TOSA-1.1+INT", "TOSA-1.1+FP"])
+def test_tosa_select_registered_for_all_profiles(spec: str) -> None:
+    with TosaLoweringContext(TosaSpecification.create_from_string(spec)):
+        registered_ops = get_registered_tosa_ops()
+
+    assert exir_ops.backend.tosa.SELECT.default in registered_ops
+
+
+def test_tosa_select_accepts_bfloat16_with_bf16_extension() -> None:
+    condition = torch.randint(0, 2, (2, 3), dtype=torch.bool)
+    input1 = torch.randn((2, 3), dtype=torch.bfloat16)
+    input2 = torch.randn((2, 3), dtype=torch.bfloat16)
+
+    with TosaLoweringContext(
+        TosaSpecification.create_from_string("TOSA-1.1+FP+bf16")
+    ), FakeTensorMode() as mode:
+        output = exir_ops.backend.tosa.SELECT.default(
+            *_to_fake(mode, condition, input1, input2)
+        )
+
+    assert output.dtype == torch.bfloat16
+    assert tuple(output.shape) == tuple(input1.shape)
+
+
+def test_tosa_select_rejects_non_bool_condition() -> None:
+    condition = torch.ones((2, 3), dtype=torch.int32)
+    input1 = torch.randn((2, 3), dtype=torch.float32)
+    input2 = torch.randn((2, 3), dtype=torch.float32)
+
+    with TosaLoweringContext(
+        TosaSpecification.create_from_string("TOSA-1.1+FP")
+    ), FakeTensorMode() as mode:
+        with pytest.raises(TosaValueError, match="requires bool condition"):
+            exir_ops.backend.tosa.SELECT.default(
+                *_to_fake(mode, condition, input1, input2)
+            )
+
+
+def test_tosa_select_rejects_mismatched_value_dtypes() -> None:
+    condition = torch.randint(0, 2, (2, 3), dtype=torch.bool)
+    input1 = torch.randn((2, 3), dtype=torch.float32)
+    input2 = torch.randn((2, 3), dtype=torch.float16)
+
+    with TosaLoweringContext(
+        TosaSpecification.create_from_string("TOSA-1.1+FP")
+    ), FakeTensorMode() as mode:
+        with pytest.raises(TosaValueError, match="Expected matching dtypes"):
+            exir_ops.backend.tosa.SELECT.default(
+                *_to_fake(mode, condition, input1, input2)
+            )
+
+
+def test_tosa_select_rejects_bfloat16_without_bf16_extension() -> None:
+    condition = torch.randint(0, 2, (2, 3), dtype=torch.bool)
+    input1 = torch.randn((2, 3), dtype=torch.bfloat16)
+    input2 = torch.randn((2, 3), dtype=torch.bfloat16)
+
+    with TosaLoweringContext(
+        TosaSpecification.create_from_string("TOSA-1.1+FP")
+    ), FakeTensorMode() as mode:
+        with pytest.raises(TosaValueError, match="doesn't support"):
+            exir_ops.backend.tosa.SELECT.default(
+                *_to_fake(mode, condition, input1, input2)
+            )

backends/arm/tosa/dialect/__init__.py

@@ -28,6 +28,7 @@
     scatter,
     shape_ops,
     table,
+    ternary_elementwise,
     transpose_conv2d,
     unary_elementwise,
 )

backends/arm/tosa/dialect/ops/ternary_elementwise.py

@@ -0,0 +1,77 @@
+# 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 torch
+from executorch.backends.arm.tosa.dialect.lib import TosaValueError
+from executorch.backends.arm.tosa.dialect.ops._common import (
+    broadcast_shape,
+    require_same_dtype,
+)
+from executorch.backends.arm.tosa.dialect.ops_registration import register_fake_tosa_op
+from executorch.backends.arm.tosa.specification import (
+    get_context_spec,
+    TosaSpecification,
+)
+
+
+def _raise_unsupported_dtype(dtype: torch.dtype, op: str) -> None:
+    raise TosaValueError(f"Unsupported dtype {dtype} for {op}", op=op)
+
+
+def _raise_unsupported_profile(dtype: torch.dtype, op: str) -> None:
+    raise TosaValueError(
+        f"TOSA spec {get_context_spec()} doesn't support {dtype} for {op}",
+        op=op,
+    )
+
+
+def _validate_condition_dtype(dtype: torch.dtype, op: str) -> None:
+    if dtype != torch.bool:
+        raise TosaValueError(f"{op} requires bool condition but got {dtype}", op=op)
+
+
+def _validate_select_value_dtype(dtype: torch.dtype) -> None:
+    tosa_spec = get_context_spec()
+
+    if dtype == torch.bool:
+        return
+
+    if dtype in {torch.int8, torch.int16, torch.int32}:
+        if not tosa_spec.support_integer():
+            _raise_unsupported_profile(dtype, "SELECT")
+        return
+
+    if dtype in {torch.float16, torch.float32}:
+        if not tosa_spec.support_float():
+            _raise_unsupported_profile(dtype, "SELECT")
+        return
+
+    if dtype == torch.bfloat16:
+        if not (tosa_spec.support_float() and tosa_spec.support_extension("bf16")):
+            _raise_unsupported_profile(dtype, "SELECT")
+        return
+
+    _raise_unsupported_dtype(dtype, "SELECT")
+
+
+@register_fake_tosa_op(
+    "SELECT(Tensor condition, Tensor input1, Tensor input2) -> Tensor",
+    TosaSpecification.all_versions_and_profiles(),
+)
+def SELECT(
+    condition: torch.Tensor,
+    input1: torch.Tensor,
+    input2: torch.Tensor,
+) -> torch.Tensor:
+    _validate_condition_dtype(condition.dtype, "SELECT")
+    require_same_dtype(input1, input2, "SELECT")
+    _validate_select_value_dtype(input1.dtype)
+    output_shape = broadcast_shape(condition, input1, "SELECT")
+    output_shape = broadcast_shape(
+        torch.empty(output_shape, dtype=input1.dtype),
+        input2,
+        "SELECT",
+    )
+    return torch.empty(output_shape, dtype=input1.dtype)