Skip to content

Commit aeb4224

Browse files
Sebastian-LarssonSebastian-Larsson
authored
Arm backend: Share data layout dtype validation (pytorch#22823)
Signed-off-by: Sebastian Larsson <sebastian.larsson@arm.com>
1 parent 8e43991 commit aeb4224

10 files changed

Lines changed: 57 additions & 133 deletions

‎backends/arm/operators/op_tosa_concat.py‎

Lines changed: 4 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
register_node_visitor,
1414
)
1515
from executorch.backends.arm.operators.operator_validation_utils import (
16+
supported_data_layout_dtypes,
1617
validate_num_inputs,
1718
validate_same_dtype,
1819
validate_valid_dtype,
@@ -25,29 +26,16 @@
2526
class CatVisitor(NodeVisitor):
2627
target = "tosa.CONCAT.default"
2728

28-
def __init__(self, *args):
29-
super().__init__(*args)
30-
3129
def define_node(
3230
self,
3331
node: Node,
3432
tosa_graph: Any,
3533
inputs: List[TosaArg],
3634
output: TosaArg,
3735
) -> None:
38-
supported_dtypes = [ts.DType.BOOL]
39-
if self.tosa_spec.support_integer():
40-
supported_dtypes.extend([ts.DType.INT8, ts.DType.INT32])
41-
if self.tosa_spec.support_extension("int16"):
42-
supported_dtypes.append(ts.DType.INT16)
43-
if self.tosa_spec.support_float():
44-
supported_dtypes.extend([ts.DType.FP16, ts.DType.FP32])
45-
if self.tosa_spec.support_extension("bf16"):
46-
supported_dtypes.append(ts.DType.BF16)
47-
if self.tosa_spec.support_extension("fp8e4m3"):
48-
supported_dtypes.append(ts.DType.FP8E4M3)
49-
if self.tosa_spec.support_extension("fp8e5m2"):
50-
supported_dtypes.append(ts.DType.FP8E5M2)
36+
supported_dtypes = supported_data_layout_dtypes(
37+
self.tosa_spec, allow_int16_without_extension=False
38+
)
5139
validate_num_inputs(self.target, inputs, [1, 2])
5240
input_tosa_args = [TosaArg(arg, self.tosa_spec) for arg in inputs[0].special]
5341
validate_same_dtype(self.target, [*input_tosa_args, output], ts)

‎backends/arm/operators/op_tosa_gather.py‎

Lines changed: 2 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
register_node_visitor,
1313
)
1414
from executorch.backends.arm.operators.operator_validation_utils import (
15+
supported_data_layout_dtypes,
1516
validate_num_inputs,
1617
validate_same_dtype,
1718
validate_valid_dtype,
@@ -56,24 +57,7 @@ def define_node(
5657
validate_valid_dtype(
5758
self.target,
5859
[values, output],
59-
[
60-
ts.DType.INT8,
61-
ts.DType.INT16,
62-
ts.DType.INT32,
63-
ts.DType.FP16,
64-
ts.DType.FP32,
65-
ts.DType.BF16,
66-
*(
67-
[ts.DType.FP8E4M3]
68-
if self.tosa_spec.support_extension("fp8e4m3")
69-
else []
70-
),
71-
*(
72-
[ts.DType.FP8E5M2]
73-
if self.tosa_spec.support_extension("fp8e5m2")
74-
else []
75-
),
76-
],
60+
supported_data_layout_dtypes(self.tosa_spec, include_bool=False),
7761
self.tosa_spec,
7862
)
7963

‎backends/arm/operators/op_tosa_pad.py‎

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
register_node_visitor,
1515
)
1616
from executorch.backends.arm.operators.operator_validation_utils import (
17+
supported_data_layout_dtypes,
1718
validate_num_inputs,
1819
validate_same_dtype,
1920
validate_valid_dtype,
@@ -34,17 +35,7 @@ def define_node(
3435
inputs: List[TosaArg],
3536
output: TosaArg,
3637
) -> None:
37-
supported_dtypes = [ts.DType.BOOL]
38-
if self.tosa_spec.support_integer():
39-
supported_dtypes.extend([ts.DType.INT8, ts.DType.INT16, ts.DType.INT32])
40-
if self.tosa_spec.support_float():
41-
supported_dtypes.extend([ts.DType.FP16, ts.DType.FP32])
42-
if self.tosa_spec.support_extension("bf16"):
43-
supported_dtypes.append(ts.DType.BF16)
44-
if self.tosa_spec.support_extension("fp8e4m3"):
45-
supported_dtypes.append(ts.DType.FP8E4M3)
46-
if self.tosa_spec.support_extension("fp8e5m2"):
47-
supported_dtypes.append(ts.DType.FP8E5M2)
38+
supported_dtypes = supported_data_layout_dtypes(self.tosa_spec)
4839

4940
validate_num_inputs(self.target, inputs, 2)
5041
validate_same_dtype(self.target, [inputs[0], output], ts)

‎backends/arm/operators/op_tosa_reshape.py‎

Lines changed: 5 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
register_node_visitor,
1515
)
1616
from executorch.backends.arm.operators.operator_validation_utils import (
17+
supported_data_layout_dtypes,
1718
validate_num_inputs,
1819
validate_same_dtype,
1920
validate_valid_dtype,
@@ -25,36 +26,23 @@
2526
class ViewVisitor(NodeVisitor):
2627
target = "tosa.RESHAPE.default"
2728

28-
def __init__(self, *args):
29-
super().__init__(*args)
30-
3129
def define_node(
3230
self,
3331
node: torch.fx.Node,
3432
tosa_graph: Any,
3533
inputs: List[TosaArg],
3634
output: TosaArg,
3735
) -> None:
38-
supported_dtypes = {ts.DType.BOOL}
39-
if self.tosa_spec.support_integer():
40-
supported_dtypes.update([ts.DType.INT8, ts.DType.INT16, ts.DType.INT32])
41-
if self.tosa_spec.support_float():
42-
supported_dtypes.update([ts.DType.FP16, ts.DType.FP32])
43-
if self.tosa_spec.support_extension("bf16"):
44-
supported_dtypes.add(ts.DType.BF16)
45-
if self.tosa_spec.support_extension("fp8e4m3"):
46-
supported_dtypes.add(ts.DType.FP8E4M3)
47-
if self.tosa_spec.support_extension("fp8e5m2"):
48-
supported_dtypes.add(ts.DType.FP8E5M2)
49-
if self.tosa_spec.support_extension("mxfp"):
50-
supported_dtypes.update([ts.DType.FP8E4M3, ts.DType.FP8E5M2])
36+
supported_dtypes = supported_data_layout_dtypes(
37+
self.tosa_spec, include_mxfp=True
38+
)
5139

5240
validate_num_inputs(self.target, inputs, 2)
5341
validate_same_dtype(self.target, [inputs[0], output], ts)
5442
validate_valid_dtype(
5543
self.target,
5644
[inputs[0], output],
57-
list(supported_dtypes),
45+
supported_dtypes,
5846
self.tosa_spec,
5947
)
6048

‎backends/arm/operators/op_tosa_reverse.py‎

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
register_node_visitor,
1414
)
1515
from executorch.backends.arm.operators.operator_validation_utils import (
16+
supported_data_layout_dtypes,
1617
validate_num_inputs,
1718
validate_same_dtype,
1819
validate_valid_dtype,
@@ -31,17 +32,7 @@ def define_node(
3132
inputs: list[TosaArg],
3233
output: TosaArg,
3334
) -> None:
34-
supported_dtypes = [ts.DType.BOOL]
35-
if self.tosa_spec.support_integer():
36-
supported_dtypes.extend([ts.DType.INT8, ts.DType.INT16, ts.DType.INT32])
37-
if self.tosa_spec.support_float():
38-
supported_dtypes.extend([ts.DType.FP16, ts.DType.FP32])
39-
if self.tosa_spec.support_extension("bf16"):
40-
supported_dtypes.append(ts.DType.BF16)
41-
if self.tosa_spec.support_extension("fp8e4m3"):
42-
supported_dtypes.append(ts.DType.FP8E4M3)
43-
if self.tosa_spec.support_extension("fp8e5m2"):
44-
supported_dtypes.append(ts.DType.FP8E5M2)
35+
supported_dtypes = supported_data_layout_dtypes(self.tosa_spec)
4536

4637
validate_num_inputs(self.target, inputs, 1)
4738
validate_same_dtype(self.target, [inputs[0], output], ts)

‎backends/arm/operators/op_tosa_scatter.py‎

Lines changed: 2 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
register_node_visitor,
1313
)
1414
from executorch.backends.arm.operators.operator_validation_utils import (
15+
supported_data_layout_dtypes,
1516
validate_num_inputs,
1617
validate_same_dtype,
1718
validate_valid_dtype,
@@ -43,24 +44,7 @@ def define_node(
4344
validate_valid_dtype(
4445
self.target,
4546
[inputs[0], inputs[2], output],
46-
[
47-
ts.DType.INT8,
48-
ts.DType.INT16,
49-
ts.DType.INT32,
50-
ts.DType.FP32,
51-
ts.DType.FP16,
52-
ts.DType.BF16,
53-
*(
54-
[ts.DType.FP8E4M3]
55-
if self.tosa_spec.support_extension("fp8e4m3")
56-
else []
57-
),
58-
*(
59-
[ts.DType.FP8E5M2]
60-
if self.tosa_spec.support_extension("fp8e5m2")
61-
else []
62-
),
63-
],
47+
supported_data_layout_dtypes(self.tosa_spec, include_bool=False),
6448
self.tosa_spec,
6549
)
6650

‎backends/arm/operators/op_tosa_slice.py‎

Lines changed: 2 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
register_node_visitor,
1414
)
1515
from executorch.backends.arm.operators.operator_validation_utils import (
16+
supported_data_layout_dtypes,
1617
validate_num_inputs,
1718
validate_same_dtype,
1819
validate_valid_dtype,
@@ -25,27 +26,14 @@
2526
class SliceVisitor(NodeVisitor):
2627
target = "tosa.SLICE.default"
2728

28-
def __init__(self, *args):
29-
super().__init__(*args)
30-
3129
def define_node(
3230
self,
3331
node: Node,
3432
tosa_graph: Any,
3533
inputs: List[TosaArg],
3634
output: TosaArg,
3735
) -> None:
38-
supported_dtypes = [ts.DType.BOOL]
39-
if self.tosa_spec.support_integer():
40-
supported_dtypes.extend([ts.DType.INT8, ts.DType.INT16, ts.DType.INT32])
41-
if self.tosa_spec.support_float():
42-
supported_dtypes.extend([ts.DType.FP16, ts.DType.FP32])
43-
if self.tosa_spec.support_extension("bf16"):
44-
supported_dtypes.append(ts.DType.BF16)
45-
if self.tosa_spec.support_extension("fp8e4m3"):
46-
supported_dtypes.append(ts.DType.FP8E4M3)
47-
if self.tosa_spec.support_extension("fp8e5m2"):
48-
supported_dtypes.append(ts.DType.FP8E5M2)
36+
supported_dtypes = supported_data_layout_dtypes(self.tosa_spec)
4937

5038
validate_num_inputs(self.target, inputs, 3)
5139
validate_same_dtype(self.target, [inputs[0], output], ts)

‎backends/arm/operators/op_tosa_tile.py‎

Lines changed: 2 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
register_node_visitor,
1515
)
1616
from executorch.backends.arm.operators.operator_validation_utils import (
17+
supported_data_layout_dtypes,
1718
validate_num_inputs,
1819
validate_same_dtype,
1920
validate_valid_dtype,
@@ -25,27 +26,14 @@
2526
class RepeatVisitor(NodeVisitor):
2627
target = "tosa.TILE.default"
2728

28-
def __init__(self, *args):
29-
super().__init__(*args)
30-
3129
def define_node(
3230
self,
3331
node: torch.fx.Node,
3432
tosa_graph: Any,
3533
inputs: list[TosaArg],
3634
output: TosaArg,
3735
) -> None:
38-
supported_dtypes = [ts.DType.BOOL]
39-
if self.tosa_spec.support_integer():
40-
supported_dtypes.extend([ts.DType.INT8, ts.DType.INT16, ts.DType.INT32])
41-
if self.tosa_spec.support_float():
42-
supported_dtypes.extend([ts.DType.FP16, ts.DType.FP32])
43-
if self.tosa_spec.support_extension("bf16"):
44-
supported_dtypes.append(ts.DType.BF16)
45-
if self.tosa_spec.support_extension("fp8e4m3"):
46-
supported_dtypes.append(ts.DType.FP8E4M3)
47-
if self.tosa_spec.support_extension("fp8e5m2"):
48-
supported_dtypes.append(ts.DType.FP8E5M2)
36+
supported_dtypes = supported_data_layout_dtypes(self.tosa_spec)
4937

5038
validate_num_inputs(self.target, inputs, 2)
5139
validate_same_dtype(self.target, [inputs[0], output], ts)

‎backends/arm/operators/op_tosa_transpose.py‎

Lines changed: 2 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
register_node_visitor,
1616
)
1717
from executorch.backends.arm.operators.operator_validation_utils import (
18+
supported_data_layout_dtypes,
1819
validate_num_inputs,
1920
validate_same_dtype,
2021
validate_valid_dtype,
@@ -26,27 +27,14 @@
2627
class PermuteVisitor(NodeVisitor):
2728
target = "tosa.TRANSPOSE.default"
2829

29-
def __init__(self, *args):
30-
super().__init__(*args)
31-
3230
def define_node(
3331
self,
3432
node: torch.fx.Node,
3533
tosa_graph: Any,
3634
inputs: List[TosaArg],
3735
output: TosaArg,
3836
) -> None:
39-
supported_dtypes = [ts.DType.BOOL]
40-
if self.tosa_spec.support_integer():
41-
supported_dtypes.extend([ts.DType.INT8, ts.DType.INT16, ts.DType.INT32])
42-
if self.tosa_spec.support_float():
43-
supported_dtypes.extend([ts.DType.FP16, ts.DType.FP32])
44-
if self.tosa_spec.support_extension("bf16"):
45-
supported_dtypes.append(ts.DType.BF16)
46-
if self.tosa_spec.support_extension("fp8e4m3"):
47-
supported_dtypes.append(ts.DType.FP8E4M3)
48-
if self.tosa_spec.support_extension("fp8e5m2"):
49-
supported_dtypes.append(ts.DType.FP8E5M2)
37+
supported_dtypes = supported_data_layout_dtypes(self.tosa_spec)
5038

5139
validate_num_inputs(self.target, inputs, 2)
5240
validate_same_dtype(self.target, [inputs[0], output], ts)

‎backends/arm/operators/operator_validation_utils.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,44 @@
1212
from typing import Any, List, Optional
1313

1414
import torch
15+
import tosa_serializer as ts
1516

1617
from executorch.backends.arm.tosa.specification import Tosa_1_00, TosaSpecification
1718

1819

20+
def supported_data_layout_dtypes(
21+
tosa_spec: TosaSpecification,
22+
*,
23+
include_bool: bool = True,
24+
allow_int16_without_extension: bool = True,
25+
include_mxfp: bool = False,
26+
) -> List[Any]:
27+
"""Return serializer dtypes supported by TOSA data-layout ops."""
28+
supported_dtypes = []
29+
if include_bool:
30+
supported_dtypes.append(ts.DType.BOOL)
31+
if tosa_spec.support_integer():
32+
if allow_int16_without_extension:
33+
supported_dtypes.extend([ts.DType.INT8, ts.DType.INT16, ts.DType.INT32])
34+
else:
35+
supported_dtypes.extend([ts.DType.INT8, ts.DType.INT32])
36+
if tosa_spec.support_extension("int16"):
37+
supported_dtypes.append(ts.DType.INT16)
38+
if tosa_spec.support_float():
39+
supported_dtypes.extend([ts.DType.FP16, ts.DType.FP32])
40+
if tosa_spec.support_extension("bf16"):
41+
supported_dtypes.append(ts.DType.BF16)
42+
if tosa_spec.support_extension("fp8e4m3"):
43+
supported_dtypes.append(ts.DType.FP8E4M3)
44+
if tosa_spec.support_extension("fp8e5m2"):
45+
supported_dtypes.append(ts.DType.FP8E5M2)
46+
if include_mxfp and tosa_spec.support_extension("mxfp"):
47+
for dtype in (ts.DType.FP8E4M3, ts.DType.FP8E5M2):
48+
if dtype not in supported_dtypes:
49+
supported_dtypes.append(dtype)
50+
return supported_dtypes
51+
52+
1953
def validate_num_inputs(op_name: str, inputs: List[Any], expected: int | List[int]):
2054
"""Validate the number of inputs against expected values.
2155

0 commit comments

Comments
 (0)