Arm backend: Reject unsupported comparisons in FP profile

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):

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))