[GitHub] [tvm] ekalda commented on a change in pull request #9384: Arm(R) Ethos(TM)-U NPU Pooling operators support

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

ekalda commented on a change in pull request #9384:
URL: https://github.com/apache/tvm/pull/9384#discussion_r740147401



##########
File path: python/tvm/relay/op/contrib/ethosu.py
##########
@@ -331,6 +332,133 @@ def qnn_depthwise_conv2d_pattern() -> tvm.relay.dataflow_pattern.DFPattern:
     return clip_or_req
 
 
+class MaxPool2DParams:
+    """
+    This class will parse a call to a ethosu.maxpool2d composite function
+    and extract the parameter information.
+    """
+
+    composite_name = "ethosu.maxpool2d"
+    # The hardware only supports padding upto the numbers as follows
+    padding_bounds = [127, 127, 128, 128]
+
+    def __init__(self, func_body: Call):
+        clip = None
+        if str(func_body.op) == "clip":
+            clip = func_body
+            pool_op = clip.args[0]
+        else:
+            pool_op = func_body
+
+        attrs = pool_op.attrs
+        self.ifm = TensorParams(pool_op.args[MaxPoolArgs.ifm.value], attrs.layout)
+        self.ofm = TensorParams(pool_op, attrs.layout)
+        self.pool_shape = [int(i) for i in attrs.pool_size]
+        self.strides = attrs.strides
+        self.padding = attrs.padding
+        self.activation = clip
+        self.pooling_type = "MAX"
+
+    def is_valid(self):
+        """
+        This function checks whether MaxPool2D has compatible attributes with the NPU
+        """
+        tensor_params = [self.ifm, self.ofm]
+        if not check_valid_dtypes(tensor_params):
+            return False
+        if self.ifm.dtype != self.ofm.dtype:
+            return False
+        if not check_strides(self.strides):
+            return False
+        if not check_batch_size(self.ifm):
+            return False
+        if not check_padding(self.padding, self.padding_bounds):
+            return False
+        # Check pool size
+        if (
+            len(self.pool_shape) != 2
+            or self.pool_shape[1] > 256
+            or self.pool_shape[0] * self.pool_shape[1] > 256 * 256
+        ):
+            return False
+        return True

Review comment:
       I think this is just an `if` corresponding to `pool_shape` and hence similar to e.g. `check_padding`, so to me it makes sense to treat it the same way




-- 
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