Replacing unary ops with LookUpTable and Take op to improve performance

python/tvm/contrib/hexagon/generate_take_op.py

@@ -0,0 +1,98 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-docstring, invalid-name, unnecessary-comprehension, unused-argument
+
+import tvm
+import tvm.testing
+from tvm import relax
+from tvm.contrib.hexagon import hexagon_unary_ops
+
+
+def op_replace(call_node, func) -> bool:
+    if not isinstance(call_node, relax.Call):
+        return False
+    call_tir_op = tvm.ir.Op.get("relax.call_tir")
+    if call_node.op != call_tir_op:
+        return False
+    ops = [
+        "qnn.tanh",
+        "qnn.sqrt",
+        "qnn.rsqrt",
+        "qnn.exp",
+        "qnn.erf",
+        "qnn.sigmoid",
+        "qnn.hardswish",
+        "qnn.log",
+        "qnn.abs",
+    ]
+    if func.attrs["op_attrs"]["op_name"] in ops:
+        return True
+    return False
+
+
+@relax.expr_functor.mutator
+class Tanh2TakeReplace(tvm.relax.PyExprMutator):
+    def __init__(self, mod: tvm.IRModule) -> None:
+        super().__init__(mod)
+        self.mod_ = mod
+
+    def transform(self) -> tvm.IRModule:
+        # Iterate over all the nodes to check for the node replaceable
+        for global_var, func in self.mod_.functions.items():
+            # Skip non-relax functions
+            if not isinstance(func, relax.Function):
+                continue
+            updated_func = self.visit_expr(func)
+            self.builder_.normalize(updated_func)
+            self.builder_.update_func(global_var, updated_func)
+        # At the end of the transformation we return the updated IRModule from the BlockBuilder.
+        return self.builder_.get()
+
+    def visit_call_(self, call_node: relax.Call) -> relax.Call:
+        call_tir_op = tvm.ir.Op.get("relax.call_tir")
+        if call_node.op != call_tir_op:
+            return call_node
+
+        var = call_node.args[0]
+        func = self.mod_[var]
+
+        if call_node.args[1][0].struct_info.dtype == "uint8":
+            if op_replace(call_node, func):
+                inp, inp_scale, inp_zp, out_scale, out_zp = [x for x in call_node.args[1]]
+                # LUT node creation
+                LUT = hexagon_unary_ops.LUT_generation(
+                    inp_scale, inp_zp, out_scale, out_zp, call_node.args[0].name_hint
+                )
+                # Take operation node creation
+                take_func = hexagon_unary_ops.generate_take_primfunc(inp, call_node.struct_info)
+                take_func = take_func.without_attr("global_symbol")
+                take_func_gv = self.builder_.add_func(take_func, "take")
+                take_node = relax.call_tir(
+                    take_func_gv,
+                    relax.expr.Tuple(
+                        [call_node.args[1][0], relax.expr.Constant(tvm.nd.array(LUT))]
+                    ),
+                    call_node.struct_info,
+                )
+                return take_node
+        return call_node
+
+
+@tvm.ir.transform.module_pass(opt_level=2, name="replace_tanh_take")
+class PassReplaceWithTakeOpPrimFuncs:
+    def transform_module(self, mod, ctx):
+        return Tanh2TakeReplace(mod).transform()

python/tvm/contrib/hexagon/hexagon_unary_ops.py

@@ -0,0 +1,97 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-docstring, invalid-name
+import logging
+import numpy as np
+from scipy import special
+from tvm import te
+
+logger = logging.getLogger(__name__)
+
+######################################################################
+#################### PRIMFUNC FOR LUT and Take Op ####################
+######################################################################
+
+
+def saturate(x: te.Tensor, dtype: str):
+    """Saturate value for the specified data type"""
+    return te.max(te.min_value(dtype), te.min(x, te.max_value(dtype)))
+
+
+def hardswish_func(x):
+    x_2 = np.add(x, 3.0)
+    x_2 = np.clip(x_2, 0.0, 6.0)
+    return x * x_2 / 6.0
+
+
+def LUT_generation(inp_scale, inp_zp, out_scale, out_zp, op_name) -> None:
+    LUT = []
+    for i in range(256):
+        i = np.int32(i)
+        # converting the constants to the numpy value
+        if inp_zp.data.shape == ():
+            i_zp = inp_zp.data.numpy()[()]
+        if inp_scale.data.shape == ():
+            i_scale = inp_scale.data.numpy()[()]
+        if out_zp.data.shape == ():
+            o_zp = out_zp.data.numpy()[()]
+        if out_scale.data.shape == ():
+            o_scale = out_scale.data.numpy()[()]
+        # Dequantization followed by computing the op value
+        dequant = (i - i_zp) * i_scale
+        if "tanh" in op_name:
+            op_val = np.tanh(dequant)
+        elif "rsqrt" in op_name:
+            op_val = 1 / np.sqrt(dequant)
+        elif "sqrt" in op_name:
+            op_val = np.sqrt(dequant)
+        elif "exp" in op_name:
+            op_val = np.exp(dequant)
+        elif "erf" in op_name:
+            op_val = special.erf(dequant)
+        elif "sigmoid" in op_name:
+            op_val = 1 / (1 + np.exp(np.negative(dequant)))
+        elif "hardswish" in op_name:
+            op_val = hardswish_func(dequant)
+        elif "log" in op_name:
+            op_val = np.log(dequant)
+        elif "abs" in op_name:
+            op_val = np.abs(dequant)
+        else:
+            logger.error("Error op is other than unary op")
+
+        # Quantizing the value generated and appending in the Look Up Table
+        quant = np.round((op_val) / o_scale) + o_zp
+        val = np.maximum(0, np.minimum(quant, 255)).astype(np.uint8)
+        LUT.append(val)
+    return LUT
+
+
+def generate_take_primfunc(inp, struct_info):
+    # Generating the take op
+    N, H, W, C = inp.struct_info.shape
+    data = te.placeholder((N, H, W, C), dtype=struct_info.dtype, name="data")
+    LUT_func = te.placeholder((256,), dtype="uint8", name="LUT")
+    take = te.compute(
+        struct_info.shape,
+        lambda *indices: saturate(
+            (LUT_func[data[indices].astype("uint8")]), struct_info.dtype
+        ).astype(struct_info.dtype),
+        name="take_op",
+    )
+    mod = te.create_prim_func([data, LUT_func, take])
+    return mod