[GitHub] [tvm] NicolaLancellotti commented on a change in pull request #9547: [microNPU] Add the infrastructure for lookup table and TANH

[GitBox <gi...@apache.org>]

NicolaLancellotti commented on a change in pull request #9547:
URL: https://github.com/apache/tvm/pull/9547#discussion_r756239310



##########
File path: python/tvm/relay/backend/contrib/ethosu/codegen.py
##########
@@ -22,6 +22,109 @@
 from tvm.relay.backend.contrib.ethosu.legalize import LegalizeEthosU
 from tvm.relay.backend.contrib.ethosu import tir_to_cs_translator
 from tvm.relay.backend.contrib.ethosu import util
+from tvm.relay.expr_functor import ExprMutator
+from tvm.ir.transform import Pass
+
+# pylint: disable=unused-import
+from tvm.relay.backend.contrib.ethosu.op import op_attrs
+from tvm.relay.backend.contrib.ethosu import op
+
+
+class OptimizeLUTs(ExprMutator):
+    """A pass to merge an identity operator with a LUT based activation function with
+    a preceding operator provided that operator can do a table lookup for the activation
+    in the hardware"""
+
+    def __init__(self):
+        super().__init__()
+        self.lut_ops = {
+            "contrib.ethosu.conv2d": op.ethosu_conv2d,
+            "contrib.ethosu.depthwise_conv2d": op.ethosu_depthwise_conv2d,
+            "contrib.ethosu.pooling": op.ethosu_pooling,
+        }
+
+    def create_op_with_lut(self, call):
+        """Extract the parameters and attributes from the NPU operator and create
+        a new operator with LUT.
+        ----------
+        call : tvm.relay.expr.Call
+            The current call node being visited.
+        Returns
+        -------
+        tvm.relay.expr.Call
+            The new operator with LUT.
+        """
+        identity = call
+        ethosu_op = call.args[0]
+        lut = identity.args[1]
+        activation = identity.attrs.activation
+
+        new_attrs = dict(ethosu_op.attrs)
+        new_attrs["activation"] = activation
+
+        # Assume that LUT is always the last argument
+        new_args = [ethosu_op.args[n] for n in range(len(ethosu_op.args) - 1)]
+        new_args.append(lut)
+        assert ethosu_op.op.name in self.lut_ops.keys()
+
+        return self.lut_ops[ethosu_op.op.name](*new_args, **new_attrs)
+
+    def visit_call(self, call: tvm.relay.expr.Call) -> tvm.relay.expr.Call:
+        """Recursively visit call nodes in the input graph and if an ethosu.identity
+        operator with LUT is found and the preceding operator has a LUT attribute, create
+        a new NPU operator.
+        Parameters
+        ----------
+        call : tvm.relay.expr.Call
+            The current call node being visited.
+        Returns

Review comment:
       ```suggestion
           a new NPU operator.
           
           Parameters
           ----------
           call : tvm.relay.expr.Call
               The current call node being visited.
           
           Returns
   ```
   ```suggestion
           a new NPU operator.
           Parameters
           ----------
           call : tvm.relay.expr.Call
               The current call node being visited.
           Returns
   ```




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org