From 2af13c72fa7a4fd1ca5f9e46cc4987ec383194ef Mon Sep 17 00:00:00 2001 From: Saoirse Stewart Date: Mon, 22 Jun 2026 15:31:22 +0100 Subject: [PATCH] Arm backend: Add TOSA elementwise ternerary node visitors Change-Id: I6c2a123a91d4bc7f331fe0abfae70792d8c0f57e --- backends/arm/_passes/aten_to_tosa_ternary.py | 25 +++ backends/arm/_passes/exir_to_tosa_pass.py | 12 ++ backends/arm/operators/op_where.py | 2 +- .../test_tosa_dialect_ternary_ops.py | 149 ++++++++++++++++++ backends/arm/tosa/dialect/__init__.py | 1 + .../tosa/dialect/ops/ternary_elementwise.py | 77 +++++++++ 6 files changed, 265 insertions(+), 1 deletion(-) create mode 100644 backends/arm/_passes/aten_to_tosa_ternary.py create mode 100644 backends/arm/test/misc/tosa_dialect/test_tosa_dialect_ternary_ops.py create mode 100644 backends/arm/tosa/dialect/ops/ternary_elementwise.py diff --git a/backends/arm/_passes/aten_to_tosa_ternary.py b/backends/arm/_passes/aten_to_tosa_ternary.py new file mode 100644 index 00000000000..adfc2287b1c --- /dev/null +++ b/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 diff --git a/backends/arm/_passes/exir_to_tosa_pass.py b/backends/arm/_passes/exir_to_tosa_pass.py index 942f88d517d..9dd66c3d866 100644 --- a/backends/arm/_passes/exir_to_tosa_pass.py +++ b/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, diff --git a/backends/arm/operators/op_where.py b/backends/arm/operators/op_where.py index 543f6b12d6b..4f0009bdd55 100644 --- a/backends/arm/operators/op_where.py +++ b/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) diff --git a/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_ternary_ops.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_ternary_ops.py new file mode 100644 index 00000000000..db3f6f99c16 --- /dev/null +++ b/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) + ) diff --git a/backends/arm/tosa/dialect/__init__.py b/backends/arm/tosa/dialect/__init__.py index 909847c4fb7..504323365ac 100644 --- a/backends/arm/tosa/dialect/__init__.py +++ b/backends/arm/tosa/dialect/__init__.py @@ -28,6 +28,7 @@ scatter, shape_ops, table, + ternary_elementwise, transpose_conv2d, unary_elementwise, ) diff --git a/backends/arm/tosa/dialect/ops/ternary_elementwise.py b/backends/arm/tosa/dialect/ops/ternary_elementwise.py new file mode 100644 index 00000000000..cce549226d6 --- /dev/null +++ b/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)