Arm backend: Rewrite le/lt to ge/gt and simplify visitors

backends/arm/_passes/__init__.py

@@ -132,6 +132,7 @@
 )
 from .rewrite_conv_pass import RewriteConvPass  # noqa
 from .rewrite_index_put_pass import RewriteIndexPutPass  # noqa
+from .rewrite_le_lt_to_ge_gt_pass import RewriteLeLtToGeGtPass  # noqa
 from .rewrite_matmul import RewriteMatmulPass  # noqa
 from .rewrite_upsample import RewriteUpsamplePass  # noqa
 from .scalars_to_attribute_pass import ScalarsToAttributePass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -117,6 +117,7 @@
     RewriteBoolToFp32CastViaInt8Pass,
     RewriteConvPass,
     RewriteIndexPutPass,
+    RewriteLeLtToGeGtPass,
     RewriteMatmulPass,
     RewriteUpsamplePass,
     ScalarsToAttributePass,
@@ -310,6 +311,7 @@ def _tosa_pipeline(
         self.add_passes(
             [
                 ReplaceScalarWithTensorByProfilePass(),
+                RewriteLeLtToGeGtPass(),
                 ConvertFullLikeToFullPass(),
                 MatchArgDtypePass(),
                 UnsqueezeScalarPlaceholdersPass(exported_program),

backends/arm/_passes/decompose_div_tensor_mode.py

@@ -20,7 +20,7 @@
     "floor": exir_ops.edge.aten.floor.default,
     "ceil": exir_ops.edge.aten.ceil.default,
     "full": exir_ops.edge.aten.full.default,
-    "lt": exir_ops.edge.aten.lt.Tensor,
+    "gt": exir_ops.edge.aten.gt.Tensor,
     "where": exir_ops.edge.aten.where.self,
 }
 
@@ -29,7 +29,7 @@
     "floor": torch.ops.aten.floor.default,
     "ceil": torch.ops.aten.ceil.default,
     "full": torch.ops.aten.full.default,
-    "lt": torch.ops.aten.lt.Tensor,
+    "gt": torch.ops.aten.gt.Tensor,
     "where": torch.ops.aten.where.self,
 }
 
@@ -87,11 +87,13 @@ def call_operator(self, op, args, kwargs, meta):
                 meta=meta,
                 updated=True,
             )
-            lt0 = super().call_operator(opset["lt"], (q, zero), {}, meta, updated=True)
+            is_neg = super().call_operator(
+                opset["gt"], (zero, q), {}, meta, updated=True
+            )
             ceilq = super().call_operator(opset["ceil"], (q,), {}, meta, updated=True)
             floorq = super().call_operator(opset["floor"], (q,), {}, meta, updated=True)
             return super().call_operator(
-                opset["where"], (lt0, ceilq, floorq), {}, meta, updated=True
+                opset["where"], (is_neg, ceilq, floorq), {}, meta, updated=True
             )
 
         raise RuntimeError(

backends/arm/_passes/rewrite_le_lt_to_ge_gt_pass.py

@@ -0,0 +1,42 @@
+# Copyright 2026 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.
+
+from typing import Set, Type
+
+import torch
+
+from executorch.backends.arm._passes import ArmPass
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass
+
+OP_MAP = {
+    exir_ops.edge.aten.le.Tensor: exir_ops.edge.aten.ge.Tensor,
+    exir_ops.edge.aten.lt.Tensor: exir_ops.edge.aten.gt.Tensor,
+    torch.ops.aten.le.Tensor: torch.ops.aten.ge.Tensor,
+    torch.ops.aten.lt.Tensor: torch.ops.aten.gt.Tensor,
+}
+
+
+class RewriteLeLtToGeGtPass(ArmPass):
+    """Rewrite le/lt into ge/gt with swapped inputs."""
+
+    _passes_required_after: Set[Type[ExportPass]] = set()
+
+    def call_operator(self, op, args, kwargs, meta):
+        if not self.allowed_to_transform(meta):
+            return super().call_operator(op, args, kwargs, meta)
+
+        target_op = OP_MAP.get(op)
+        if target_op is None:
+            return super().call_operator(op, args, kwargs, meta)
+
+        lhs, rhs = args
+        return super().call_operator(
+            target_op,
+            (rhs, lhs),
+            kwargs,
+            meta,
+            updated=True,
+        )

backends/arm/operators/__init__.py

@@ -32,10 +32,8 @@
     op_ge,
     op_gt,
     op_index_tensor,
-    op_le,
     op_log,
     op_logical_not,
-    op_lt,
     op_max_pool2d,
     op_maximum,
     op_minimum,

backends/arm/operators/op_eq.py

@@ -3,56 +3,27 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-
-from typing import Any, List
-
 import tosa_serializer as ts
 
-from executorch.backends.arm.operators.node_visitor import (
-    NodeVisitor,
-    register_node_visitor,
+from executorch.backends.arm.operators.node_visitor import register_node_visitor
+from executorch.backends.arm.operators.simple_node_visitor import (
+    SimpleNodeVisitor,
+    SimpleNodeVisitorConfig,
 )
-from executorch.backends.arm.operators.operator_validation_utils import (
-    validate_num_inputs,
-    validate_same_dtype,
-    validate_valid_dtype,
-)
-from executorch.backends.arm.tosa.mapping import TosaArg
-
-from torch.fx import Node
 
 
 @register_node_visitor
-class EqualVisitor(NodeVisitor):
+class EqualVisitor(SimpleNodeVisitor):
     target = "aten.eq.Tensor"
 
-    def __init__(self, *args):
-        super().__init__(*args)
-
-    def define_node(
-        self,
-        node: Node,
-        tosa_graph: Any,
-        inputs: List[TosaArg],
-        output: TosaArg,
-    ) -> None:
-        validate_num_inputs(self.target, inputs, 2)
-        validate_same_dtype(self.target, inputs, ts)
-        validate_valid_dtype(
-            self.target,
-            inputs,
-            [ts.DType.INT32, ts.DType.FP16, ts.DType.FP32, ts.DType.BF16],
-            self.tosa_spec,
-        )
-        validate_valid_dtype(self.target, output, ts.DType.BOOL, self.tosa_spec)
-
-        attr = ts.TosaSerializerAttribute()
-        attr.EqualAttribute()
-        self._serialize_operator(
-            node,
-            tosa_graph,
-            ts.Op.EQUAL,
-            [inputs[0].name, inputs[1].name],
-            [output.name],
-            attr,
+    @classmethod
+    def get_config(cls) -> SimpleNodeVisitorConfig:
+        return SimpleNodeVisitorConfig(
+            tosa_op=ts.Op.EQUAL,
+            attr_method="EqualAttribute",
+            num_inputs=2,
+            input_dtypes=[ts.DType.INT32, ts.DType.FP16, ts.DType.FP32, ts.DType.BF16],
+            output_dtypes=[ts.DType.BOOL],
+            same_dtype_with_output=False,
+            dtype_check_inputs_only=True,
         )

backends/arm/operators/op_ge.py

@@ -3,56 +3,29 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-
-from typing import Any, List
-
 import tosa_serializer as ts
 
-from executorch.backends.arm.operators.node_visitor import (
-    NodeVisitor,
-    register_node_visitor,
-)
-from executorch.backends.arm.operators.operator_validation_utils import (
-    validate_num_inputs,
-    validate_same_dtype,
-    validate_valid_dtype,
+from executorch.backends.arm.operators.node_visitor import register_node_visitor
+from executorch.backends.arm.operators.simple_node_visitor import (
+    SimpleNodeVisitor,
+    SimpleNodeVisitorConfig,
 )
-from executorch.backends.arm.tosa.mapping import TosaArg
 
-from torch.fx import Node
+COMPARE_INPUT_DTYPES = [ts.DType.INT32, ts.DType.FP32, ts.DType.BF16, ts.DType.FP16]
 
 
 @register_node_visitor
-class GreaterEqualVisitor(NodeVisitor):
+class GreaterEqualVisitor(SimpleNodeVisitor):
     target = "aten.ge.Tensor"
 
-    def __init__(self, *args):
-        super().__init__(*args)
-
-    def define_node(
-        self,
-        node: Node,
-        tosa_graph: Any,
-        inputs: List[TosaArg],
-        output: TosaArg,
-    ) -> None:
-        validate_num_inputs(self.target, inputs, 2)
-        validate_same_dtype(self.target, inputs, ts)
-        validate_valid_dtype(
-            self.target,
-            inputs,
-            [ts.DType.INT32, ts.DType.FP32, ts.DType.BF16, ts.DType.FP16],
-            self.tosa_spec,
-        )
-        validate_valid_dtype(self.target, output, ts.DType.BOOL, self.tosa_spec)
-
-        attr = ts.TosaSerializerAttribute()
-        attr.GreaterEqualAttribute()
-        self._serialize_operator(
-            node,
-            tosa_graph,
-            ts.Op.GREATER_EQUAL,
-            [inputs[0].name, inputs[1].name],
-            [output.name],
-            attr,
+    @classmethod
+    def get_config(cls) -> SimpleNodeVisitorConfig:
+        return SimpleNodeVisitorConfig(
+            tosa_op=ts.Op.GREATER_EQUAL,
+            attr_method="GreaterEqualAttribute",
+            num_inputs=2,
+            input_dtypes=COMPARE_INPUT_DTYPES,
+            output_dtypes=[ts.DType.BOOL],
+            same_dtype_with_output=False,
+            dtype_check_inputs_only=True,
         )

backends/arm/operators/op_gt.py

@@ -3,56 +3,29 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-
-from typing import Any, List
-
 import tosa_serializer as ts
 
-from executorch.backends.arm.operators.node_visitor import (
-    NodeVisitor,
-    register_node_visitor,
-)
-from executorch.backends.arm.operators.operator_validation_utils import (
-    validate_num_inputs,
-    validate_same_dtype,
-    validate_valid_dtype,
+from executorch.backends.arm.operators.node_visitor import register_node_visitor
+from executorch.backends.arm.operators.simple_node_visitor import (
+    SimpleNodeVisitor,
+    SimpleNodeVisitorConfig,
 )
-from executorch.backends.arm.tosa.mapping import TosaArg
 
-from torch.fx import Node
+COMPARE_INPUT_DTYPES = [ts.DType.INT32, ts.DType.FP32, ts.DType.BF16, ts.DType.FP16]
 
 
 @register_node_visitor
-class GreaterThanVisitor(NodeVisitor):
+class GreaterThanVisitor(SimpleNodeVisitor):
     target = "aten.gt.Tensor"
 
-    def __init__(self, *args):
-        super().__init__(*args)
-
-    def define_node(
-        self,
-        node: Node,
-        tosa_graph: Any,
-        inputs: List[TosaArg],
-        output: TosaArg,
-    ) -> None:
-        validate_num_inputs(self.target, inputs, 2)
-        validate_same_dtype(self.target, inputs, ts)
-        validate_valid_dtype(
-            self.target,
-            inputs,
-            [ts.DType.INT32, ts.DType.FP32, ts.DType.BF16, ts.DType.FP16],
-            self.tosa_spec,
-        )
-        validate_valid_dtype(self.target, output, ts.DType.BOOL, self.tosa_spec)
-
-        attr = ts.TosaSerializerAttribute()
-        attr.GreaterAttribute()
-        self._serialize_operator(
-            node,
-            tosa_graph,
-            ts.Op.GREATER,
-            [inputs[0].name, inputs[1].name],
-            [output.name],
-            attr,
+    @classmethod
+    def get_config(cls) -> SimpleNodeVisitorConfig:
+        return SimpleNodeVisitorConfig(
+            tosa_op=ts.Op.GREATER,
+            attr_method="GreaterAttribute",
+            num_inputs=2,
+            input_dtypes=COMPARE_INPUT_DTYPES,
+            output_dtypes=[ts.DType.BOOL],
+            same_dtype_with_output=False,
+            dtype_check_inputs_only=True,
         )