ArmBackend: Add support for comparison operators (#7885)

Add support for the following operators:
 * aten.eq.Tensor
 * aten.le.Tensor
 * aten.lt.Tensor
 * aten.ge.Tensor
 * aten.gt.Tensor

Signed-off-by: Georgios Pinitas <[email protected]>
Signed-off-by: Per Åstrand <[email protected]>
Co-authored-by: Georgios Pinitas <[email protected]>

Author: per

backends/arm/operator_support/tosa_supported_operators.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Arm Limited and/or its affiliates.
+# Copyright 2024-2025 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.
@@ -81,11 +81,16 @@ def is_node_supported(self, submodules, node: fx.Node) -> bool:
             exir_ops.edge.aten.hardtanh.default,
             exir_ops.edge.aten.convolution.default,
             exir_ops.edge.aten.div.Tensor,
+            exir_ops.edge.aten.eq.Tensor,
             exir_ops.edge.aten.exp.default,
             exir_ops.edge.aten.log.default,
             exir_ops.edge.aten.linear.default,
             exir_ops.edge.aten.split_with_sizes_copy.default,
             exir_ops.edge.aten.full.default,
+            exir_ops.edge.aten.ge.Tensor,
+            exir_ops.edge.aten.gt.Tensor,
+            exir_ops.edge.aten.le.Tensor,
+            exir_ops.edge.aten.lt.Tensor,
             exir_ops.edge.aten.mul.Tensor,
             exir_ops.edge.aten._native_batch_norm_legit_no_training.default,
             exir_ops.edge.aten.native_layer_norm.default,

backends/arm/operators/__init__.py

@@ -13,11 +13,16 @@
     op_bmm,
     op_cat,
     op_conv2d,
+    op_eq,
     op_exp,
     op_full,
+    op_ge,
     op_get_item,
+    op_gt,
     op_hardtanh,
+    op_le,
     op_log,
+    op_lt,
     op_max,
     op_max_pool2d,
     op_min,

backends/arm/operators/op_eq.py

@@ -0,0 +1,58 @@
+# Copyright 2025 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.
+
+# pyre-unsafe
+
+from typing import List
+
+import executorch.backends.arm.tosa_quant_utils as tqutils
+
+import serializer.tosa_serializer as ts
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from serializer.tosa_serializer import TosaOp
+
+from torch.fx import Node
+
+
+@register_node_visitor
+class EqualVisitor(NodeVisitor):
+    target = "aten.eq.Tensor"
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+        assert (
+            inputs[0].dtype == inputs[1].dtype
+        ), "EQ must have the same dtypes as input"
+
+        input_nodes = inputs
+        # Handle quantization
+        if inputs[0].dtype == ts.DType.INT8:
+            # Rescale inputs to 32 bit
+            rescaled_inputs, _ = tqutils.insert_rescale_ops_to_int32(
+                tosa_graph, inputs, node
+            )
+
+            # Update IO
+            input_nodes = rescaled_inputs
+
+        # Do the equal comparison
+        tosa_graph.addOperator(
+            TosaOp.Op().EQUAL,
+            [input_nodes[0].name, input_nodes[1].name],
+            output.name,
+            None,
+        )

backends/arm/operators/op_ge.py

@@ -0,0 +1,57 @@
+# Copyright 2025 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.
+
+# pyre-unsafe
+
+from typing import List
+
+import executorch.backends.arm.tosa_quant_utils as tqutils
+
+import serializer.tosa_serializer as ts
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from serializer.tosa_serializer import TosaOp
+
+from torch.fx import Node
+
+
+@register_node_visitor
+class GreaterEqualVisitor(NodeVisitor):
+    target = "aten.ge.Tensor"
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+        assert (
+            inputs[0].dtype == inputs[1].dtype
+        ), "GE must have the same dtypes as input"
+
+        input_nodes = inputs
+        # Handle quantization
+        if inputs[0].dtype == ts.DType.INT8:
+            # Rescale inputs to 32 bit
+            rescaled_inputs, _ = tqutils.insert_rescale_ops_to_int32(
+                tosa_graph, inputs, node
+            )
+
+            # Update IO
+            input_nodes = rescaled_inputs
+
+        tosa_graph.addOperator(
+            TosaOp.Op().GREATER_EQUAL,
+            [input_nodes[0].name, input_nodes[1].name],
+            [output.name],
+            None,
+        )

backends/arm/operators/op_gt.py

@@ -0,0 +1,57 @@
+# Copyright 2025 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.
+
+# pyre-unsafe
+
+from typing import List
+
+import executorch.backends.arm.tosa_quant_utils as tqutils
+
+import serializer.tosa_serializer as ts
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from serializer.tosa_serializer import TosaOp
+
+from torch.fx import Node
+
+
+@register_node_visitor
+class GreaterThanVisitor(NodeVisitor):
+    target = "aten.gt.Tensor"
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+        assert (
+            inputs[0].dtype == inputs[1].dtype
+        ), "GT must have the same dtypes as input"
+
+        input_nodes = inputs
+        # Handle quantization
+        if inputs[0].dtype == ts.DType.INT8:
+            # Rescale inputs to 32 bit
+            rescaled_inputs, _ = tqutils.insert_rescale_ops_to_int32(
+                tosa_graph, inputs, node
+            )
+
+            # Update IO
+            input_nodes = rescaled_inputs
+
+        tosa_graph.addOperator(
+            TosaOp.Op().GREATER,
+            [input_nodes[0].name, input_nodes[1].name],
+            [output.name],
+            None,
+        )

backends/arm/operators/op_le.py

@@ -0,0 +1,57 @@
+# Copyright 2025 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.
+
+# pyre-unsafe
+
+from typing import List
+
+import executorch.backends.arm.tosa_quant_utils as tqutils
+
+import serializer.tosa_serializer as ts
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from serializer.tosa_serializer import TosaOp
+
+from torch.fx import Node
+
+
+@register_node_visitor
+class LessEqualVisitor(NodeVisitor):
+    target = "aten.le.Tensor"
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+        assert (
+            inputs[0].dtype == inputs[1].dtype
+        ), "LE must have the same dtypes as input"
+
+        input_nodes = inputs
+        # Handle quantization
+        if inputs[0].dtype == ts.DType.INT8:
+            # Rescale inputs to 32 bit
+            rescaled_inputs, _ = tqutils.insert_rescale_ops_to_int32(
+                tosa_graph, inputs, node
+            )
+
+            # Update IO
+            input_nodes = rescaled_inputs
+
+        tosa_graph.addOperator(
+            TosaOp.Op().GREATER_EQUAL,
+            [input_nodes[1].name, input_nodes[0].name],
+            [output.name],
+            None,
+        )

backends/arm/operators/op_lt.py

@@ -0,0 +1,57 @@
+# Copyright 2025 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.
+
+# pyre-unsafe
+
+from typing import List
+
+import executorch.backends.arm.tosa_quant_utils as tqutils
+
+import serializer.tosa_serializer as ts
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from serializer.tosa_serializer import TosaOp
+
+from torch.fx import Node
+
+
+@register_node_visitor
+class LessThanVisitor(NodeVisitor):
+    target = "aten.lt.Tensor"
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+        assert (
+            inputs[0].dtype == inputs[1].dtype
+        ), "LT must have the same dtypes as input"
+
+        input_nodes = inputs
+        # Handle quantization
+        if inputs[0].dtype == ts.DType.INT8:
+            # Rescale inputs to 32 bit
+            rescaled_inputs, _ = tqutils.insert_rescale_ops_to_int32(
+                tosa_graph, inputs, node
+            )
+
+            # Update IO
+            input_nodes = rescaled_inputs
+
+        tosa_graph.addOperator(
+            TosaOp.Op().GREATER,
+            [input_nodes[1].name, input_nodes[0].name],
+            [output.name],
+            None,
+        )

backends/arm/quantizer/quantization_annotator.py

@@ -304,6 +304,21 @@ def any_or_hardtanh_min_zero(n: Node):
         quant_properties.quant_output = _QuantProperty(
             0, SharedQuantizationSpec((node.args[0], node))
         )
+    elif node.target in [
+        torch.ops.aten.eq.Tensor,
+        torch.ops.aten.ge.Tensor,
+        torch.ops.aten.gt.Tensor,
+        torch.ops.aten.le.Tensor,
+        torch.ops.aten.lt.Tensor,
+    ]:
+        shared_qspec = SharedQuantizationSpec((node.args[0], node))
+        quant_properties.quant_inputs = [
+            _QuantProperty(0, input_act_qspec),
+            _QuantProperty(
+                1, input_act_qspec if node.args[0] == node.args[1] else shared_qspec
+            ),
+        ]
+        quant_properties.quant_output = None
     elif node.target in _parent_shared_qspec:
         if not isinstance(node.args[0], Node):
             return None