MLBEDSW-5554: Place MEAN op exceeding max height with axis==1 on CPU

Signed-off-by: Rickard Bolin <rickard.bolin@arm.com>
Change-Id: I87dc5963972a7ef91db467b2ff8e0261e9899372

2 files changed, 45 insertions, 2 deletions

diff --git a/ethosu/vela/tflite_graph_optimiser.py b/ethosu/vela/tflite_graph_optimiser.py
index e01433d..7b10f86 100644
--- a/ethosu/vela/tflite_graph_optimiser.py
+++ b/ethosu/vela/tflite_graph_optimiser.py
@@ -1301,8 +1301,8 @@ def convert_mean_to_depthwise_conv_or_avgpool(op, arch, nng):
             if dims == 2:
                 shape += [1]
 
-        # If height is greater than max kernel height, reshape to from HxW to 1x(HxW)
-        if h > 64:
+        # If height is greater than max kernel height, reshape from HxW to 1x(HxW)
+        if (h > 64 and op.type == Op.DepthwiseConv2DBias) or (h > 256 and op.type == Op.AvgPool):
             shape = [shape[0], 1, h * w, shape[3]]
             op.ifm_shapes[0] = Shape4D(shape)
             if h > 256 and op.type == Op.AvgPool:
diff --git a/ethosu/vela/tflite_supported_operators.py b/ethosu/vela/tflite_supported_operators.py
index 5c7fd51..60bc6fd 100644
--- a/ethosu/vela/tflite_supported_operators.py
+++ b/ethosu/vela/tflite_supported_operators.py
@@ -209,6 +209,8 @@ class TFLiteSupportedOperators:
         self.specific_constraints[Op.Mean].append(TFLiteSupportedOperators.constraint_mean_height_width_product_avgpool)
         self.specific_constraints[Op.Mean].append(TFLiteSupportedOperators.constraint_mean_height_width_product)
         self.specific_constraints[Op.Mean].append(TFLiteSupportedOperators.constraint_mean_height_width_product_int8)
+        self.specific_constraints[Op.Mean].append(TFLiteSupportedOperators.constraint_depthwise_conv_height_single_axis)
+        self.specific_constraints[Op.Mean].append(TFLiteSupportedOperators.constraint_avgpool_height_single_axis)
 
         # Reshape specific checks:
         self.specific_constraints[Op.Reshape].append(TFLiteSupportedOperators.constraint_reshape_shape_constant)
@@ -686,6 +688,47 @@ class TFLiteSupportedOperators:
         max_prod = cls.mean_kernel_product_int8
         return h * w <= max_prod, f"Product of height and width is {h * w}"
 
+    @classmethod
+    @docstring_format_args([dilated_height_range[1]])
+    def constraint_depthwise_conv_height_single_axis(cls, op):
+        """Height can be at most {} for single axis when axis is 1."""
+        inp, axis = op.inputs
+        if axis.shape == [] or axis.shape[0] == 1:  # single axis
+            axis = int(axis.values) if len(axis.shape) == 0 else int(axis.values[0])
+        else:
+            # Multiple axes, constraint does not apply
+            return True, ""
+
+        # Height and width axes have different index depending on dimensions
+        shape = inp.shape
+        h = shape[0] if len(shape) < 4 else shape[1]
+
+        # If quantization is the same across IFM and OFM op will become avgpool and this constraint does not apply.
+        ifm, ofm = op.get_ifm_ofm()
+        if check_quantized_tens_scaling_equal(ifm, ofm):
+            return True, ""
+
+        return h <= 64 or axis != 1, f"Height is {h} and axis is {axis}."
+
+    @classmethod
+    @docstring_format_args([filter_height_range[1]])
+    def constraint_avgpool_height_single_axis(cls, op):
+        """Avgpool height can be at most {} for single axis when axis is 1."""
+        inp, axis = op.inputs
+        if axis.shape == [] or axis.shape[0] == 1:  # single axis
+            axis = int(axis.values) if len(axis.shape) == 0 else int(axis.values[0])
+        else:
+            # Multiple axes, constraint does not apply
+            return True, ""
+
+        # Height and width axes have different index depending on dimensions
+        shape = inp.shape
+        h = shape[0] if len(shape) < 4 else shape[1]
+        ifm, ofm = op.get_ifm_ofm()
+        scaling_equal = check_quantized_tens_scaling_equal(ifm, ofm)
+
+        return h <= 256 or axis != 1 or not scaling_equal, f"Height is {h} and axis is {axis}"
+
     @staticmethod
     def constraint_reshape_shape_constant(op):
         "Shape must be constant"