Arm: Add pass that converts operators to clamp to the Arm backend (#8538)

Add pass that converts operators to clamp to the Arm backend

* Add ConvertToClampPass that converts relu and hardtanh to clamp
* Remove op_relu and op_hardtanh visitors from backend

Signed-off-by: Tom Allsop <[email protected]>

Author: tom-arm

backends/arm/_passes/arm_pass_manager.py

@@ -27,6 +27,7 @@
 from executorch.backends.arm._passes.convert_squeezes_to_view import (  # type: ignore[import-not-found]
     ConvertSqueezesToViewPass,
 )
+from executorch.backends.arm._passes.convert_to_clamp import ConvertToClampPass
 from executorch.backends.arm._passes.decompose_batchnorm_pass import (
     DecomposeBatchNormPass,
 )
@@ -104,6 +105,7 @@ def _tosa_080_BI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(DecomposeLinearPass())
         self.add_pass(ConvertMeanDimToAveragePoolPass())
         self.add_pass(ConvertFullLikeToFullPass())
+        self.add_pass(ConvertToClampPass())
 
         self.add_pass(ReplaceScalarWithTensorArgPass())
         self.add_pass(AnnotateDecomposedMatmulPass())
@@ -144,6 +146,8 @@ def _tosa_080_MI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(DecomposeDivPass())
         self.add_pass(DecomposeSoftmaxesPass())
         self.add_pass(ConvertFullLikeToFullPass())
+        self.add_pass(ConvertToClampPass())
+
         self.add_pass(AnnotateDecomposedMatmulPass())
         self.add_pass(QuantizeOperatorArguments())
         self.add_pass(FoldAndAnnotateQParamsPass())  # type: ignore[call-arg]

backends/arm/_passes/convert_to_clamp.py

@@ -0,0 +1,36 @@
+# 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.
+
+from typing import Tuple
+
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass
+
+edge_operators = {
+    exir_ops.edge.aten.hardtanh.default,
+    exir_ops.edge.aten.relu.default,
+}
+
+
+def get_clamp_params(op, args) -> Tuple[float | None, float | None]:
+    if op == exir_ops.edge.aten.hardtanh.default:
+        return args[1], args[2]
+    elif op == exir_ops.edge.aten.relu.default:
+        return 0.0, None
+    else:
+        raise ValueError(f"Getting clamp parameters for op {op} is not implemented.")
+
+
+class ConvertToClampPass(ExportPass):
+    def call_operator(self, op, args, kwargs, meta):
+        if op not in edge_operators:
+            return super().call_operator(op, args, kwargs, meta)
+
+        return super().call_operator(
+            exir_ops.edge.aten.clamp.default,
+            (args[0], *get_clamp_params(op, args)),
+            {},
+            meta,
+        )

backends/arm/operators/__init__.py

@@ -21,7 +21,6 @@
     op_ge,
     op_get_item,
     op_gt,
-    op_hardtanh,
     op_le,
     op_log,
     op_lt,
@@ -31,7 +30,6 @@
     op_mul,
     op_permute,
     op_reciprocal,
-    op_relu,
     op_repeat,
     op_rescale,
     op_rshift,

backends/arm/operators/op_hardtanh.py

@@ -0,0 +1,80 @@
+# 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.
+
+import unittest
+
+import torch
+from executorch.backends.arm._passes.convert_to_clamp import ConvertToClampPass
+
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.arm_tester import ArmTester
+
+from executorch.backends.xnnpack.test.tester.tester import RunPasses
+
+
+class HardTanh(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.hardtanh = torch.nn.Hardtanh()
+
+    def forward(self, x):
+        return self.hardtanh(x)
+
+    def get_inputs(self):
+        return (torch.rand(1, 64, 64, 3),)
+
+
+class ReLU(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.relu = torch.nn.ReLU()
+
+    def forward(self, x):
+        return self.relu(x)
+
+    def get_inputs(self):
+        return (torch.rand(1, 64, 64, 3),)
+
+
+class TestConvertToClampPass(unittest.TestCase):
+    """
+    Tests the ConvertToClampPass which converts hardtanh.default and relu.default to clamp.default
+    """
+
+    def test_tosa_MI_hardtahn(self):
+        module = HardTanh()
+        test_pass_stage = RunPasses([ConvertToClampPass])
+        (
+            ArmTester(
+                module,
+                example_inputs=module.get_inputs(),
+                compile_spec=common.get_tosa_compile_spec("TOSA-0.80+MI"),
+            )
+            .export()
+            .to_edge()
+            .check(["executorch_exir_dialects_edge__ops_aten_hardtanh_default"])
+            .run_passes(test_pass_stage)
+            .check(["executorch_exir_dialects_edge__ops_aten_clamp_default"])
+            .check_not(["executorch_exir_dialects_edge__ops_aten_hardtanh_default"])
+        )
+
+    def test_tosa_MI_relu(self):
+        module = ReLU()
+        test_pass_stage = RunPasses([ConvertToClampPass])
+        (
+            ArmTester(
+                module,
+                example_inputs=module.get_inputs(),
+                compile_spec=common.get_tosa_compile_spec("TOSA-0.80+MI"),
+            )
+            .export()
+            .to_edge()
+            .check(["executorch_exir_dialects_edge__ops_aten_relu_default"])
+            .run_passes(test_pass_stage)
+            .check(["executorch_exir_dialects_edge__ops_aten_clamp_default"])
+            .check_not(["executorch_exir_dialects_edge__ops_aten_relu_default"])
+        )