Add reciprocal op to Arm Backend

Implements node visitor and tests
Quantized by one_to_one annotator

Signed-off-by: Erik Lundell <[email protected]>
Change-Id: I3f25096a8b908d7c25b6cd83bd7edb6871b145ab

Author: Erik-Lundell

backends/arm/arm_partitioner.py

@@ -57,6 +57,7 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             exir_ops.edge.aten.sigmoid.default,
             exir_ops.edge.aten.mm.default,
             exir_ops.edge.aten.repeat.default,
+            exir_ops.edge.aten.reciprocal.default,
             exir_ops.edge.aten.relu.default,
             exir_ops.edge.aten._softmax.default,
             exir_ops.edge.aten.slice_copy.Tensor,

backends/arm/operators/__init__.py

@@ -26,6 +26,7 @@
     op_mul,
     op_permute,
     op_quant,
+    op_reciprocal,
     op_relu,
     op_repeat,
     op_sigmoid,

backends/arm/operators/op_reciprocal.py

@@ -0,0 +1,80 @@
+# Copyright 2023-2024 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 List
+
+import numpy as np
+
+import serializer.tosa_serializer as ts
+import torch
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from executorch.backends.arm.tosa_quant_utils import (
+    dequantize_value,
+    get_quant_node_args,
+    QuantArgs,
+    quantize_value,
+)
+from serializer.tosa_serializer import TosaOp
+
+
+@register_node_visitor
+class DivVisitor(NodeVisitor):
+    target = "aten.reciprocal.default"
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+        is_quant_node: bool,
+    ) -> None:
+        # 1/X
+
+        if is_quant_node:
+            input = inputs[0]
+            input_qargs = get_quant_node_args(node.args[0])
+            output_qargs = get_quant_node_args(list(node.users)[0])
+
+            div_table = div_table_8bit(input_qargs, output_qargs)
+
+            table_attr = ts.TosaSerializerAttribute()
+            table_attr.TableAttribute(div_table)
+            tosa_graph.addOperator(
+                TosaOp.Op().TABLE, [input.name], [output.name], table_attr
+            )
+
+        else:
+            tosa_graph.addOperator(
+                TosaOp.Op().RECIPROCAL, [inputs[0].name], [output.name]
+            )
+
+
+def div_table_8bit(in_quantargs: QuantArgs, out_quantargs: QuantArgs):
+    """
+    Returns a table mapping 256 entries to div([qmin,qmax])
+    """
+
+    def div(x):
+        # Convert quantized input to floating point div input space.
+        v1 = dequantize_value(x, in_quantargs)
+        # Compute div.
+        v2 = 1.0 / v1
+        # Convert div output back to quantized space.
+        v3 = quantize_value(v2, out_quantargs)
+
+        # print(f"{v1} -> {v2} -> {v3}")
+        return v3
+
+    return [
+        div(x)
+        for x in np.linspace(in_quantargs.qmin, in_quantargs.qmax, 256, dtype=np.int8)
+    ]

backends/arm/quantizer/quantization_annotation/one_to_one_annotator.py

@@ -35,7 +35,11 @@ def _annotate_one_to_one(
     Typical ops are ops implemented with a lookup table.
     """
     annotated_partitions = []
-    one_to_one_ops = (torch.ops.aten.exp.default, torch.ops.aten.log.default)
+    one_to_one_ops = (
+        torch.ops.aten.exp.default,
+        torch.ops.aten.log.default,
+        torch.ops.aten.reciprocal.default,
+    )
     for node in gm.graph.nodes:
         if node.op != "call_function" or node.target not in one_to_one_ops:
             continue

backends/arm/test/ops/test_reciprocal.py

@@ -0,0 +1,118 @@
+# Copyright 2024 Arm Limited and/or its affiliates.
+# All rights reserved.
+#
+# 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.test import common
+from executorch.backends.arm.test.tester.arm_tester import ArmTester
+from parameterized import parameterized
+
+test_data_t = tuple[str, torch.Tensor]
+test_data_suite: list[test_data_t] = [
+    (
+        "op_reciprocal_rank1_ones",
+        torch.ones(5),
+    ),
+    (
+        "op_reciprocal_rank1_rand",
+        torch.rand(5) * 5,
+    ),
+    ("op_reciprocal_rank1_negative_ones", torch.ones(5) * (-1)),
+    ("op_reciprocal_rank4_ones", torch.ones(5, 10, 25, 20)),
+    ("op_reciprocal_rank4_negative_ones", (-1) * torch.ones(5, 10, 25, 20)),
+    ("op_reciprocal_rank4_ones_reciprocal_negative", torch.ones(5, 10, 25, 20)),
+    ("op_reciprocal_rank4_large_rand", 200 * torch.rand(5, 10, 25, 20)),
+    ("op_reciprocal_rank4_negative_large_rand", (-200) * torch.rand(5, 10, 25, 20)),
+    ("op_reciprocal_rank4_large_randn", 200 * torch.randn(5, 10, 25, 20) + 1),
+]
+
+
+class TestReciprocal(unittest.TestCase):
+    """Tests reciprocal"""
+
+    class Reciprocal(torch.nn.Module):
+
+        def forward(self, input_: torch.Tensor):
+            return input_.reciprocal()
+
+    def _test_reciprocal_tosa_MI_pipeline(
+        self, module: torch.nn.Module, test_data: tuple[torch.Tensor]
+    ):
+        (
+            ArmTester(
+                module,
+                example_inputs=test_data,
+                compile_spec=common.get_tosa_compile_spec(),
+            )
+            .export()
+            .check_count({"torch.ops.aten.reciprocal.default": 1})
+            .check_not(["torch.ops.quantized_decomposed"])
+            .to_edge()
+            .partition()
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .run_method_and_compare_outputs(inputs=test_data)
+        )
+
+    def _test_reciprocal_tosa_BI_pipeline(
+        self, module: torch.nn.Module, test_data: tuple[torch.Tensor]
+    ):
+        (
+            ArmTester(
+                module,
+                example_inputs=test_data,
+                compile_spec=common.get_tosa_compile_spec(),
+            )
+            .quantize()
+            .export()
+            .check_count({"torch.ops.aten.reciprocal.default": 1})
+            .check(["torch.ops.quantized_decomposed"])
+            .to_edge()
+            .partition()
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .run_method_and_compare_outputs(inputs=test_data)
+        )
+
+    def _test_reciprocal_u55_BI_pipeline(
+        self, module: torch.nn.Module, test_data: tuple[torch.Tensor]
+    ):
+        (
+            ArmTester(
+                module,
+                example_inputs=test_data,
+                compile_spec=common.get_u55_compile_spec(),
+            )
+            .quantize()
+            .export()
+            .check_count({"torch.ops.aten.reciprocal.default": 1})
+            .check(["torch.ops.quantized_decomposed"])
+            .to_edge()
+            .partition()
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+        )
+
+    @parameterized.expand(test_data_suite)
+    def test_reciprocal_tosa_MI(self, test_name: str, input_: torch.Tensor):
+        test_data = (input_,)
+        self._test_reciprocal_tosa_MI_pipeline(self.Reciprocal(), test_data)
+
+    # Expected to fail since ArmQuantizer cannot quantize a Reciprocal layer
+    # TODO(MLETORCH-129)
+    @parameterized.expand(test_data_suite)
+    def test_reciprocal_tosa_BI(self, test_name: str, input_: torch.Tensor):
+
+        test_data = (input_,)
+        self._test_reciprocal_tosa_BI_pipeline(self.Reciprocal(), test_data)
+
+    # Expected to fail since Vela does not support TABLE
+    @parameterized.expand(test_data_suite)
+    @unittest.expectedFailure
+    def test_reciprocal_u55_BI(self, test_name: str, input_: torch.Tensor):
+        test_data = (input_,)
+        self._test_reciprocal_u55_BI_pipeline(self.Reciprocal(), test_data)