Skip to content

Commit 7e9ac9b

Browse files
SaoirseARMSaoirseARM
authored
Arm backend: Clean up of TOSA dialect operators (pytorch#22515)
- move test_tosa_* to tosa_dialect folder - move comparison operators out of binary operators - add aten->tosa for comparison operators Signed-off-by: Saoirse Stewart <saoirse.stewart@arm.com>
1 parent 1e3b7fb commit 7e9ac9b

20 files changed

Lines changed: 247 additions & 119 deletions

‎backends/arm/_passes/aten_to_tosa_comparison.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
# Copyright 2026 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
from executorch.backends.transforms.aten_to_dialect_pass import (
7+
AtenToDialectPass,
8+
DialectNodeSpec,
9+
)
10+
from executorch.exir.dialects._ops import ops as exir_ops
11+
from torch.fx import Node
12+
13+
14+
def rewrite_comparison_operator(
15+
node: Node, pass_: AtenToDialectPass
16+
) -> DialectNodeSpec | None:
17+
match node.target:
18+
case exir_ops.edge.aten.eq.Tensor:
19+
target = exir_ops.backend.tosa.EQUAL.default
20+
case exir_ops.edge.aten.ge.Tensor:
21+
target = exir_ops.backend.tosa.GREATER_EQUAL.default
22+
case exir_ops.edge.aten.gt.Tensor:
23+
target = exir_ops.backend.tosa.GREATER.default
24+
case _:
25+
return None
26+
27+
return DialectNodeSpec(target, node.args, dict(node.kwargs))

‎backends/arm/_passes/aten_to_tosa_tensor_operators.py‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -56,12 +56,6 @@ def rewrite_binary_operator(
5656
target = exir_ops.backend.tosa.ARITHMETIC_RIGHT_SHIFT.default
5757
case exir_ops.edge.aten.bitwise_xor.Tensor:
5858
target = exir_ops.backend.tosa.BITWISE_XOR.default
59-
case exir_ops.edge.aten.eq.Tensor:
60-
target = exir_ops.backend.tosa.EQUAL.default
61-
case exir_ops.edge.aten.ge.Tensor:
62-
target = exir_ops.backend.tosa.GREATER_EQUAL.default
63-
case exir_ops.edge.aten.gt.Tensor:
64-
target = exir_ops.backend.tosa.GREATER.default
6559
case exir_ops.edge.aten.logical_and.default:
6660
target = exir_ops.backend.tosa.LOGICAL_AND.default
6761
case exir_ops.edge.aten.logical_or.default:

‎backends/arm/_passes/exir_to_tosa_pass.py‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,9 @@
99
from executorch.backends.arm._passes.aten_to_tosa_activation_functions import (
1010
get_activation_replacement,
1111
)
12+
from executorch.backends.arm._passes.aten_to_tosa_comparison import (
13+
rewrite_comparison_operator,
14+
)
1215
from executorch.backends.arm._passes.aten_to_tosa_data_layout import (
1316
rewrite_data_layout_operator,
1417
)
@@ -79,9 +82,6 @@ def _get_fft_replacement(
7982
exir_ops.edge.aten.bitwise_or.Tensor,
8083
exir_ops.edge.aten.bitwise_right_shift.Tensor,
8184
exir_ops.edge.aten.bitwise_xor.Tensor,
82-
exir_ops.edge.aten.eq.Tensor,
83-
exir_ops.edge.aten.ge.Tensor,
84-
exir_ops.edge.aten.gt.Tensor,
8585
exir_ops.edge.aten.logical_and.default,
8686
exir_ops.edge.aten.logical_or.default,
8787
exir_ops.edge.aten.logical_xor.default,
@@ -97,6 +97,17 @@ def _get_binary_operator_replacement(
9797
return rewrite_binary_operator(node, pass_)
9898

9999

100+
@register_dialect_substitutions(
101+
exir_ops.edge.aten.eq.Tensor,
102+
exir_ops.edge.aten.ge.Tensor,
103+
exir_ops.edge.aten.gt.Tensor,
104+
)
105+
def _get_comparison_operator_replacement(
106+
node: Node, pass_: AtenToDialectPass
107+
) -> DialectNodeSpec | None:
108+
return rewrite_comparison_operator(node, pass_)
109+
110+
100111
@register_dialect_substitutions(
101112
exir_ops.edge.aten.abs.default,
102113
exir_ops.edge.aten.bitwise_not.default,

backends/arm/test/misc/test_tosa_data_layout_visitors.py renamed to backends/arm/test/misc/tosa_dialect/test_tosa_data_layout_visitors.py

File renamed without changes.

backends/arm/test/misc/test_tosa_dialect_activation.py renamed to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_activation.py

File renamed without changes.

backends/arm/test/misc/test_tosa_dialect_argmax.py renamed to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_argmax.py

File renamed without changes.

‎backends/arm/test/misc/tosa_dialect/test_tosa_dialect_binary_ops.py‎

Lines changed: 0 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -81,33 +81,6 @@ def _to_fake(mode: FakeTensorMode, *values):
8181
(2, 3),
8282
torch.int8,
8383
),
84-
pytest.param(
85-
"EQUAL",
86-
"TOSA-1.1+INT",
87-
torch.randint(1, 16, (2, 1, 3), dtype=torch.int32),
88-
torch.randint(1, 8, (1, 4, 3), dtype=torch.int32),
89-
{},
90-
(2, 4, 3),
91-
torch.bool,
92-
),
93-
pytest.param(
94-
"GREATER",
95-
"TOSA-1.1+FP",
96-
torch.randn((2, 1, 3), dtype=torch.float32),
97-
torch.randn((1, 4, 3), dtype=torch.float32),
98-
{},
99-
(2, 4, 3),
100-
torch.bool,
101-
),
102-
pytest.param(
103-
"GREATER_EQUAL",
104-
"TOSA-1.1+INT",
105-
torch.randint(1, 16, (2, 1, 3), dtype=torch.int32),
106-
torch.randint(1, 8, (1, 4, 3), dtype=torch.int32),
107-
{},
108-
(2, 4, 3),
109-
torch.bool,
110-
),
11184
pytest.param(
11285
"INTDIV",
11386
"TOSA-1.1+INT",
@@ -400,31 +373,6 @@ def test_intdiv_supports_int32_on_fp_profile() -> None:
400373
assert tuple(output.shape) == tuple(input1.shape)
401374

402375

403-
def test_equal_rejects_int8() -> None:
404-
input1 = torch.randint(-8, 8, (2, 3), dtype=torch.int8)
405-
input2 = torch.randint(-8, 8, (2, 3), dtype=torch.int8)
406-
407-
with TosaLoweringContext(
408-
TosaSpecification.create_from_string("TOSA-1.1+INT")
409-
), FakeTensorMode() as mode:
410-
with pytest.raises(TosaValueError, match="Unsupported dtype"):
411-
exir_ops.backend.tosa.EQUAL.default(*_to_fake(mode, input1, input2))
412-
413-
414-
@pytest.mark.parametrize("op_name", ["EQUAL", "GREATER", "GREATER_EQUAL"])
415-
def test_compare_ops_reject_int32_on_fp_profile(op_name: str) -> None:
416-
input1 = torch.randint(1, 16, (2, 3), dtype=torch.int32)
417-
input2 = torch.randint(1, 8, (2, 3), dtype=torch.int32)
418-
419-
with TosaLoweringContext(
420-
TosaSpecification.create_from_string("TOSA-1.1+FP")
421-
), FakeTensorMode() as mode:
422-
with pytest.raises(TosaValueError, match="doesn't support int32"):
423-
getattr(exir_ops.backend.tosa, op_name).default(
424-
*_to_fake(mode, input1, input2)
425-
)
426-
427-
428376
@pytest.mark.parametrize("op_name", ["MAXIMUM", "MINIMUM"])
429377
def test_extrema_ops_reject_int32_on_fp_profile(op_name: str) -> None:
430378
input1 = torch.randint(1, 16, (2, 3), dtype=torch.int32)

‎backends/arm/test/misc/tosa_dialect/test_tosa_dialect_comparison.py‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
# Copyright 2026 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
import executorch.backends.arm.tosa.dialect # noqa: F401
7+
import pytest
8+
import torch
9+
from executorch.backends.arm.tosa.dialect.lib import TosaValueError
10+
from executorch.backends.arm.tosa.specification import (
11+
TosaLoweringContext,
12+
TosaSpecification,
13+
)
14+
from executorch.exir.dialects._ops import ops as exir_ops
15+
from torch._subclasses.fake_tensor import FakeTensorMode
16+
17+
18+
def _to_fake(mode: FakeTensorMode, *values):
19+
return [
20+
mode.from_tensor(value) if isinstance(value, torch.Tensor) else value
21+
for value in values
22+
]
23+
24+
25+
@pytest.mark.parametrize(
26+
(
27+
"op_name",
28+
"spec",
29+
"input1",
30+
"input2",
31+
"expected_shape",
32+
),
33+
[
34+
pytest.param(
35+
"EQUAL",
36+
"TOSA-1.1+INT",
37+
torch.randint(1, 16, (2, 1, 3), dtype=torch.int32),
38+
torch.randint(1, 8, (1, 4, 3), dtype=torch.int32),
39+
(2, 4, 3),
40+
),
41+
pytest.param(
42+
"GREATER",
43+
"TOSA-1.1+FP",
44+
torch.randn((2, 1, 3), dtype=torch.float32),
45+
torch.randn((1, 4, 3), dtype=torch.float32),
46+
(2, 4, 3),
47+
),
48+
pytest.param(
49+
"GREATER_EQUAL",
50+
"TOSA-1.1+INT",
51+
torch.randint(1, 16, (2, 1, 3), dtype=torch.int32),
52+
torch.randint(1, 8, (1, 4, 3), dtype=torch.int32),
53+
(2, 4, 3),
54+
),
55+
],
56+
)
57+
def test_tosa_comparison_ops(
58+
op_name: str,
59+
spec: str,
60+
input1: torch.Tensor,
61+
input2: torch.Tensor,
62+
expected_shape: tuple[int, ...],
63+
) -> None:
64+
with TosaLoweringContext(
65+
TosaSpecification.create_from_string(spec)
66+
), FakeTensorMode() as mode:
67+
output = getattr(exir_ops.backend.tosa, op_name).default(
68+
*_to_fake(mode, input1, input2)
69+
)
70+
71+
assert output.dtype == torch.bool
72+
assert tuple(output.shape) == expected_shape
73+
74+
75+
def test_equal_rejects_int8() -> None:
76+
input1 = torch.randint(-8, 8, (2, 3), dtype=torch.int8)
77+
input2 = torch.randint(-8, 8, (2, 3), dtype=torch.int8)
78+
79+
with TosaLoweringContext(
80+
TosaSpecification.create_from_string("TOSA-1.1+INT")
81+
), FakeTensorMode() as mode:
82+
with pytest.raises(TosaValueError, match="Unsupported dtype"):
83+
exir_ops.backend.tosa.EQUAL.default(*_to_fake(mode, input1, input2))
84+
85+
86+
@pytest.mark.parametrize("op_name", ["EQUAL", "GREATER", "GREATER_EQUAL"])
87+
def test_compare_ops_reject_int32_on_fp_profile(op_name: str) -> None:
88+
input1 = torch.randint(1, 16, (2, 3), dtype=torch.int32)
89+
input2 = torch.randint(1, 8, (2, 3), dtype=torch.int32)
90+
91+
with TosaLoweringContext(
92+
TosaSpecification.create_from_string("TOSA-1.1+FP")
93+
), FakeTensorMode() as mode:
94+
with pytest.raises(TosaValueError, match="doesn't support int32"):
95+
getattr(exir_ops.backend.tosa, op_name).default(
96+
*_to_fake(mode, input1, input2)
97+
)

backends/arm/test/misc/test_tosa_dialect_fft.py renamed to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_fft.py

File renamed without changes.

backends/arm/test/misc/test_tosa_dialect_max_pool2d_adaptive.py renamed to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_max_pool2d_adaptive.py

File renamed without changes.

0 commit comments

Comments
 (0)