diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index c1cf33764b9..a1cdd630876 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -222,7 +222,7 @@ def _floating_profile_negative_checks( ) -> list[OperatorSupportBase]: checks: list[OperatorSupportBase] = [CheckMixedFloatingInputs(reporter)] if not tosa_spec.support_integer(): - checks.append(CheckInt32ComparisonInputs(reporter)) + checks.append(CheckFPComparisonInputs(reporter)) return checks @@ -1146,12 +1146,14 @@ def is_node_supported( return True -class CheckInt32ComparisonInputs(OperatorSupportBase): - """Reject int32 comparisons under the FP profile.""" +class CheckFPComparisonInputs(OperatorSupportBase): + """Reject unsupported comparison inputs under the FP profile.""" target_ops = { exir_ops.edge.aten.eq.Tensor, exir_ops.edge.aten.eq.Scalar, + exir_ops.edge.aten.ne.Tensor, + exir_ops.edge.aten.ne.Scalar, exir_ops.edge.aten.ge.Tensor, exir_ops.edge.aten.ge.Scalar, exir_ops.edge.aten.gt.Tensor, @@ -1161,6 +1163,8 @@ class CheckInt32ComparisonInputs(OperatorSupportBase): exir_ops.edge.aten.lt.Tensor, exir_ops.edge.aten.lt.Scalar, } + supported_dtypes = {torch.float16, torch.float32, torch.bfloat16} + castable_comparison_dtypes = {torch.int8, torch.int16} def __init__(self, reporter: WhyNoPartitionReporter) -> None: self.reporter = reporter @@ -1172,19 +1176,25 @@ def is_node_supported( if node.target not in self.target_ops: return True - for input_node in ( - input_node + input_dtypes = [ + get_first_fake_tensor(input_node).dtype for input_node in node.all_input_nodes if input_node.op != "get_attr" - ): - if get_first_fake_tensor(input_node).dtype == torch.int32: - self.reporter.report_reject( - node, - "FP profile does not support int32 comparison inputs.", - ) - return False + ] + if all(dtype in self.supported_dtypes for dtype in input_dtypes): + return True - return True + if all(dtype in self.castable_comparison_dtypes for dtype in input_dtypes): + return True + + unsupported_dtype = next( + dtype for dtype in input_dtypes if dtype not in self.supported_dtypes + ) + self.reporter.report_reject( + node, + f"FP profile does not support {unsupported_dtype} comparison inputs.", + ) + return False class CheckScalarReductionInputs(OperatorSupportBase): diff --git a/backends/arm/test/misc/test_tosa_operator_support.py b/backends/arm/test/misc/test_tosa_operator_support.py index a3dce64fefc..662b428a21f 100644 --- a/backends/arm/test/misc/test_tosa_operator_support.py +++ b/backends/arm/test/misc/test_tosa_operator_support.py @@ -3,8 +3,10 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import pytest import torch from executorch.backends.arm.operator_support.tosa_supported_operators import ( + CheckFPComparisonInputs, CheckKnownUnsupportedTOSASemantics, ) from executorch.exir.backend.utils import WhyNoPartitionReporter @@ -27,6 +29,49 @@ def _checker() -> CheckKnownUnsupportedTOSASemantics: return CheckKnownUnsupportedTOSASemantics(WhyNoPartitionReporter()) +def _fp_comparison_checker() -> CheckFPComparisonInputs: + return CheckFPComparisonInputs(WhyNoPartitionReporter()) + + +@pytest.mark.parametrize( + "target", + ( + exir_ops.edge.aten.eq.Tensor, + exir_ops.edge.aten.ne.Tensor, + exir_ops.edge.aten.ge.Tensor, + exir_ops.edge.aten.gt.Tensor, + exir_ops.edge.aten.le.Tensor, + exir_ops.edge.aten.lt.Tensor, + ), +) +@pytest.mark.parametrize( + "dtype", + (torch.bool, torch.uint8, torch.int32, torch.int64), +) +def test_fp_comparison_rejects_unsupported_inputs(target, dtype) -> None: + graph = torch.fx.Graph() + x = _placeholder(graph, "x", (3, 4), dtype) + y = _placeholder(graph, "y", (3, 4), dtype) + node = graph.call_function(target, (x, y)) + node.meta["val"] = _fake_tensor((3, 4), torch.bool) + + assert not _fp_comparison_checker().is_node_supported({}, node) + + +@pytest.mark.parametrize( + "dtype", + (torch.float16, torch.float32, torch.bfloat16, torch.int8, torch.int16), +) +def test_fp_comparison_accepts_supported_inputs(dtype) -> None: + graph = torch.fx.Graph() + x = _placeholder(graph, "x", (3, 4), dtype) + y = _placeholder(graph, "y", (3, 4), dtype) + node = graph.call_function(exir_ops.edge.aten.eq.Tensor, (x, y)) + node.meta["val"] = _fake_tensor((3, 4), torch.bool) + + assert _fp_comparison_checker().is_node_supported({}, node) + + def test_rejects_argmax_without_int32_cast_user() -> None: graph = torch.fx.Graph() x = _placeholder(graph, "x", (3, 4))