lhutton1 commented on a change in pull request #9561:
URL: https://github.com/apache/tvm/pull/9561#discussion_r755176517



##########
File path: python/tvm/relay/backend/contrib/ethosu/codegen.py
##########
@@ -52,29 +60,147 @@ def constant_updater(expr, symbol):  # pylint: 
disable=unused-argument
     return dict()
 
 
+class LayoutOptimization(ExprMutator):
+    """A pass to optimize the layout of NPU operations. If both the
+    producer and consumer of a tensor are NPU operators, then the
+    layout is converted from NHWC to NHCWB16.
+    """
+
+    def __init__(self):
+        self.children = {}
+        self.optimize_op = {
+            "contrib.ethosu.conv2d": ethosu_op.ethosu_conv2d,
+            "contrib.ethosu.depthwise_conv2d": 
ethosu_op.ethosu_depthwise_conv2d,
+            "contrib.ethosu.pooling": ethosu_op.ethosu_pooling,
+            "contrib.ethosu.binary_elementwise": 
ethosu_op.ethosu_binary_elementwise,
+        }

Review comment:
       Thanks will do, I realized this as soon as I pushed the PR :)




-- 
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: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


Reply via email to