Arm backend: Add TOSA dialect unary elementwise node visitors

SaoirseARM · SaoirseARM · commit 9371b6d41c93 · 2026-06-29T16:09:09.000+01:00
Signed-off-by: Saoirse Stewart &lt;saoirse.stewart@arm.com&gt;
diff --git a/backends/arm/_passes/aten_to_tosa_tensor_operators.py b/backends/arm/_passes/aten_to_tosa_tensor_operators.py
@@ -82,3 +82,37 @@ def rewrite_binary_operator(
             return None
 
     return DialectNodeSpec(target, node.args, dict(node.kwargs))
+
+
+def rewrite_unary_operator(
+    node: Node, pass_: AtenToDialectPass
+) -> DialectNodeSpec | None:
+    match node.target:
+        case exir_ops.edge.aten.abs.default:
+            target = exir_ops.backend.tosa.ABS.default
+        case exir_ops.edge.aten.bitwise_not.default:
+            target = exir_ops.backend.tosa.BITWISE_NOT.default
+        case exir_ops.edge.aten.ceil.default:
+            target = exir_ops.backend.tosa.CEIL.default
+        case exir_ops.edge.aten.cos.default:
+            target = exir_ops.backend.tosa.COS.default
+        case exir_ops.edge.aten.exp.default:
+            target = exir_ops.backend.tosa.EXP.default
+        case exir_ops.edge.aten.floor.default:
+            target = exir_ops.backend.tosa.FLOOR.default
+        case exir_ops.edge.aten.log.default:
+            target = exir_ops.backend.tosa.LOG.default
+        case exir_ops.edge.aten.logical_not.default:
+            target = exir_ops.backend.tosa.LOGICAL_NOT.default
+        case exir_ops.edge.aten.neg.default:
+            target = exir_ops.backend.tosa.NEGATE.default
+        case exir_ops.edge.aten.reciprocal.default:
+            target = exir_ops.backend.tosa.RECIPROCAL.default
+        case exir_ops.edge.aten.rsqrt.default:
+            target = exir_ops.backend.tosa.RSQRT.default
+        case exir_ops.edge.aten.sin.default:
+            target = exir_ops.backend.tosa.SIN.default
+        case _:
+            return None
+
+    return DialectNodeSpec(target, node.args, dict(node.kwargs))
diff --git a/backends/arm/_passes/exir_to_tosa_pass.py b/backends/arm/_passes/exir_to_tosa_pass.py
@@ -13,6 +13,7 @@
     rewrite_argmax,
     rewrite_binary_operator,
     rewrite_rfft2,
+    rewrite_unary_operator,
 )
 from executorch.backends.transforms.aten_to_dialect_pass import (
     AtenToDialectPass,
@@ -87,6 +88,26 @@ def _get_binary_operator_replacement(
     return rewrite_binary_operator(node, pass_)
 
 
+@register_dialect_substitutions(
+    exir_ops.edge.aten.abs.default,
+    exir_ops.edge.aten.bitwise_not.default,
+    exir_ops.edge.aten.ceil.default,
+    exir_ops.edge.aten.cos.default,
+    exir_ops.edge.aten.exp.default,
+    exir_ops.edge.aten.floor.default,
+    exir_ops.edge.aten.log.default,
+    exir_ops.edge.aten.logical_not.default,
+    exir_ops.edge.aten.neg.default,
+    exir_ops.edge.aten.reciprocal.default,
+    exir_ops.edge.aten.rsqrt.default,
+    exir_ops.edge.aten.sin.default,
+)
+def _get_unary_operator_replacement(
+    node: Node, pass_: AtenToDialectPass
+) -> DialectNodeSpec | None:
+    return rewrite_unary_operator(node, pass_)
+
+
 @register_dialect_substitutions(
     exir_ops.edge.aten.clamp.default,
     exir_ops.edge.aten.erf.default,
diff --git a/backends/arm/operators/__init__.py b/backends/arm/operators/__init__.py
@@ -11,61 +11,62 @@
 
 from . import (  # noqa
     node_visitor,
-    op_abs,
     op_amax,
     op_amin,
     op_any,
-    op_bitwise_not,
     op_cat,
-    op_ceil,
     op_cond_if,
-    op_cos,
-    op_exp,
-    op_floor,
-    op_log,
-    op_logical_not,
-    op_neg,
     op_permute,
-    op_reciprocal,
     op_repeat,
-    op_rsqrt,
-    op_sin,
     op_sum,
     op_to_dim_order_copy,
+    op_tosa_abs,
     op_tosa_add,
     op_tosa_argmax,
     op_tosa_avg_pool2d,
     op_tosa_avg_pool2d_adaptive,
     op_tosa_binary_ops,
+    op_tosa_bitwise_not,
     op_tosa_cast_to_block_scaled,
+    op_tosa_ceil,
     op_tosa_clamp,
+    op_tosa_clz,
     op_tosa_conv2d,
     op_tosa_conv2d_block_scaled,
     op_tosa_conv3d,
+    op_tosa_cos,
     op_tosa_custom,
     op_tosa_depthwise_conv2d,
     op_tosa_eq,
     op_tosa_erf,
+    op_tosa_exp,
     op_tosa_fft,
+    op_tosa_floor,
     op_tosa_gather,
     op_tosa_ge,
     op_tosa_gt,
     op_tosa_identity,
+    op_tosa_log,
+    op_tosa_logical_not,
     op_tosa_matmul,
     op_tosa_matmul_t_block_scaled,
     op_tosa_max_pool2d,
     op_tosa_max_pool2d_adaptive,
     op_tosa_maximum,
     op_tosa_minimum,
     op_tosa_mul,
+    op_tosa_neg,
     op_tosa_pad,
     op_tosa_pow,
+    op_tosa_reciprocal,
     op_tosa_rescale,
     op_tosa_resize,
     op_tosa_rshift_tensor,
+    op_tosa_rsqrt,
     op_tosa_scatter,
     op_tosa_shapes,
     op_tosa_sigmoid,
+    op_tosa_sin,
     op_tosa_slice,
     op_tosa_sub,
     op_tosa_table,
diff --git a/backends/arm/operators/op_tosa_abs.py b/backends/arm/operators/op_tosa_abs.py
@@ -14,7 +14,7 @@
 
 @register_node_visitor
 class AbsVisitor(SimpleNodeVisitor):
-    target = "aten.abs.default"
+    target = "tosa.ABS.default"
 
     @classmethod
     def get_config(cls) -> SimpleNodeVisitorConfig:
diff --git a/backends/arm/operators/op_tosa_bitwise_not.py b/backends/arm/operators/op_tosa_bitwise_not.py
@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class BitwiseNotVisitor(SimpleNodeVisitor):
-    target = "aten.bitwise_not.default"
+    target = "tosa.BITWISE_NOT.default"
     tosa_specs = INT_SPECS
 
     @classmethod
diff --git a/backends/arm/operators/op_tosa_ceil.py b/backends/arm/operators/op_tosa_ceil.py
@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class CeilVisitor(SimpleNodeVisitor):
-    target = "aten.ceil.default"
+    target = "tosa.CEIL.default"
     tosa_specs = FP_SPECS
 
     @classmethod
diff --git a/backends/arm/operators/op_tosa_clz.py b/backends/arm/operators/op_tosa_clz.py
@@ -0,0 +1,30 @@
+# 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.
+
+import tosa_serializer as ts
+
+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 import TosaSpecification
+
+INT_SPECS = TosaSpecification.all_versions_for_profile("INT")
+
+
+@register_node_visitor
+class ClzVisitor(SimpleNodeVisitor):
+    target = "tosa.CLZ.default"
+    tosa_specs = INT_SPECS
+
+    @classmethod
+    def get_config(cls) -> SimpleNodeVisitorConfig:
+        return SimpleNodeVisitorConfig(
+            tosa_op=ts.Op.CLZ,
+            attr_method="ClzAttribute",
+            num_inputs=1,
+            input_dtypes=[ts.DType.INT32],
+        )
diff --git a/backends/arm/operators/op_tosa_cos.py b/backends/arm/operators/op_tosa_cos.py
@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class CosVisitor(SimpleNodeVisitor):
-    target = "aten.cos.default"
+    target = "tosa.COS.default"
     tosa_specs = FP_SPECS
 
     @classmethod
diff --git a/backends/arm/operators/op_tosa_exp.py b/backends/arm/operators/op_tosa_exp.py
@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class ExpVisitor(SimpleNodeVisitor):
-    target = "aten.exp.default"
+    target = "tosa.EXP.default"
     tosa_specs = FP_SPECS
 
     @classmethod
diff --git a/backends/arm/operators/op_tosa_floor.py b/backends/arm/operators/op_tosa_floor.py
@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class FloorVisitor(SimpleNodeVisitor):
-    target = "aten.floor.default"
+    target = "tosa.FLOOR.default"
     tosa_specs = FP_SPECS
 
     @classmethod
diff --git a/backends/arm/operators/op_tosa_log.py b/backends/arm/operators/op_tosa_log.py
@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class LogVisitor(SimpleNodeVisitor):
-    target = "aten.log.default"
+    target = "tosa.LOG.default"
     tosa_specs = FP_SPECS
 
     @classmethod
@@ -26,5 +26,5 @@ def get_config(cls) -> SimpleNodeVisitorConfig:
             tosa_op=ts.Op.LOG,
             attr_method="LogAttribute",
             num_inputs=1,
-            input_dtypes=[ts.DType.FP32, ts.DType.FP16, ts.DType.BF16],
+            input_dtypes=[ts.DType.FP16, ts.DType.FP32, ts.DType.BF16],
         )
diff --git a/backends/arm/operators/op_tosa_logical_not.py b/backends/arm/operators/op_tosa_logical_not.py
@@ -14,7 +14,7 @@
 
 @register_node_visitor
 class LogicalNotVisitor(SimpleNodeVisitor):
-    target = "aten.logical_not.default"
+    target = "tosa.LOGICAL_NOT.default"
 
     @classmethod
     def get_config(cls) -> SimpleNodeVisitorConfig:
diff --git a/backends/arm/operators/op_tosa_neg.py b/backends/arm/operators/op_tosa_neg.py
@@ -6,7 +6,6 @@
 from typing import Any, List
 
 import torch.fx
-
 import tosa_serializer as ts
 
 from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import (
@@ -26,11 +25,6 @@
 
 
 def get_negate_zero_points(node: torch.fx.Node, is_int8: bool) -> tuple[int, int]:
-    """Returns (input1_zp, output_zp) for TOSA NEGATE.
-
-    Must be zero for non-int8 types.
-
-    """
     if is_int8:
         return (
             get_input_qparams(node)[0].get_zp_per_tensor(),
@@ -41,10 +35,7 @@ def get_negate_zero_points(node: torch.fx.Node, is_int8: bool) -> tuple[int, int
 
 @register_node_visitor
 class NegVisitor(NodeVisitor):
-    target = "aten.neg.default"
-
-    def __init__(self, *args):
-        super().__init__(*args)
+    target = "tosa.NEGATE.default"
 
     def define_node(
         self,
@@ -71,14 +62,13 @@ def define_node(
         input_zp, output_zp = get_negate_zero_points(
             node, inputs[0].dtype == ts.DType.INT8
         )
-
         input_zp_tensor = tosa_graph.addConst(
             (1,), inputs[0].dtype, [input_zp], name=output.name + "_input_zp"
         )
-
         output_zp_tensor = tosa_graph.addConst(
             (1,), output.dtype, [output_zp], name=output.name + "_output_zp"
         )
+
         attr = ts.TosaSerializerAttribute()
         attr.NegateAttribute()
         self._serialize_operator(
diff --git a/backends/arm/operators/op_tosa_reciprocal.py b/backends/arm/operators/op_tosa_reciprocal.py
@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class ReciprocalVisitor(SimpleNodeVisitor):
-    target = "aten.reciprocal.default"
+    target = "tosa.RECIPROCAL.default"
     tosa_specs = FP_SPECS
 
     @classmethod
diff --git a/backends/arm/operators/op_tosa_rsqrt.py b/backends/arm/operators/op_tosa_rsqrt.py
@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class RsqrtVisitor(SimpleNodeVisitor):
-    target = "aten.rsqrt.default"
+    target = "tosa.RSQRT.default"
     tosa_specs = FP_SPECS
 
     @classmethod
diff --git a/backends/arm/operators/op_tosa_sin.py b/backends/arm/operators/op_tosa_sin.py
@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class SinVisitor(SimpleNodeVisitor):
-    target = "aten.sin.default"
+    target = "tosa.SIN.default"
     tosa_specs = FP_SPECS
 
     @classmethod
