rebase and add support for no-upsampling case.

Change-Id: I840d8ee3671a40c5c99f22119442c349dbed39cf

Author: lhutton1

python/tvm/relay/backend/contrib/ethosu/legalize.py

@@ -1274,6 +1274,8 @@ class Resize2dRewriter(DFPatternCallback):
     Convert ethos-u.resize2d composite function to an equivalent operation that
     performs the relevant upsampling operation.
 
+    Case 1: No upsampling (upscale factor of 1):
+        Identity.
     Case 1: Nearest neighbor upsampling:
         1x1 pooling with 2x2 nearest neighbor upsampling.
     Case 2: Bilinear upsampling:
@@ -1296,23 +1298,32 @@ def callback(
         ifm_shape = params.ifm.shape
         in_channels = ifm_shape[-1]
         reduced_op = params.ifm.tensor
+        current_size = np.array(ifm_shape[1:3])
+        output_size = np.array(params.size)
 
-        if params.method == "nearest_neighbor":
-            pool_shape = [1, 1]
-            padding = [0, 0, 0, 0]
-            rounding_mode = "TFL"
-        else:  # Bilinear upsampling
+        if (current_size == output_size).all():
+            return ethosu_ops.ethosu_identity(
+                reduced_op,
+                lut,
+                ifm_scale=float(params.ifm.q_params.scale_f32),
+                ifm_zero_point=int(params.ifm.q_params.zero_point),
+                ofm_scale=float(params.ofm.q_params.scale_f32),
+                ofm_zero_point=int(params.ofm.q_params.zero_point),
+            )
+
+        padding = [0, 0, 0, 0]
+        rounding_mode = "TFL"
+        pool_shape = [1, 1]
+        if params.method == "linear":
             pool_shape = [2, 2]
             rounding_mode = "NATURAL"
-            # For align_corners use VALID padding, otherwise use SAME padding
-            if params.coordinate_transformation_mode == "align_corners":
-                padding = [0, 0, 0, 0]
-            else:
+            if params.coordinate_transformation_mode == "asymmetric":
+                # Use SAME padding.
                 ypad = Resize2dRewriter.get_required_padding(ifm_shape[1])
                 xpad = Resize2dRewriter.get_required_padding(ifm_shape[2])
                 padding = [ypad // 2, xpad // 2, (ypad + 1) // 2, (xpad + 1) // 2]
 
-        reduced_op = ethosu_ops.ethosu_pooling(
+        return ethosu_ops.ethosu_pooling(
             ifm=reduced_op,
             lut=lut,
             pooling_type="AVG",
@@ -1328,8 +1339,6 @@ def callback(
             rounding_mode=rounding_mode,
         )
 
-        return reduced_op
-
     @staticmethod
     def get_required_padding(input_size: int, pool_size: int = 2) -> int:
         """Gets the amount of padding required needed to achieve

python/tvm/relay/backend/contrib/ethosu/te/binary_elementwise.py

@@ -288,7 +288,10 @@ def match_ethosu_binary_elementwise(output_tensor, device_config):
     pad = binary_elementwise.op.input_tensors[0]
     if pad.op.name != "ethosu_pad":
         return None
-    convert_to_nhwc = pad.op.input_tensors[0]
+    upscale = pad.op.input_tensors[0]
+    if upscale.op.name != "ethosu_upscale":
+        return None
+    convert_to_nhwc = upscale.op.input_tensors[0]
     if convert_to_nhwc.op.name != "ethosu_convert_to_nhwc":
         return None
     read = convert_to_nhwc.op.input_tensors[0]
@@ -297,7 +300,10 @@ def match_ethosu_binary_elementwise(output_tensor, device_config):
     pad2 = binary_elementwise.op.input_tensors[1]
     if pad2.op.name != "ethosu_pad":
         return None
-    convert_to_nhwc2 = pad2.op.input_tensors[0]
+    upscale2 = pad2.op.input_tensors[0]
+    if upscale2.op.name != "ethosu_upscale":
+        return None
+    convert_to_nhwc2 = upscale2.op.input_tensors[0]
     if convert_to_nhwc2.op.name != "ethosu_convert_to_nhwc":
         return None
     read2 = convert_to_nhwc2.op.input_tensors[0]

python/tvm/relay/backend/contrib/ethosu/te/depthwise.py

@@ -267,7 +267,10 @@ def match_ethosu_depthwise_conv2d(output_tensor, device_config):
     pad = depthwise2d.op.input_tensors[0]
     if pad.op.name != "ethosu_pad":
         return None
-    convert_to_nhwc = pad.op.input_tensors[0]
+    upscale = pad.op.input_tensors[0]
+    if upscale.op.name != "ethosu_upscale":
+        return None
+    convert_to_nhwc = upscale.op.input_tensors[0]
     if convert_to_nhwc.op.name != "ethosu_convert_to_nhwc":
         return None
     read = convert_to_nhwc.op.input_tensors[0]

python/tvm/relay/backend/contrib/ethosu/te/pooling.py

@@ -231,7 +231,10 @@ def match_ethosu_pooling(output_tensor, device_config):
     pad = pool2d.op.input_tensors[0]
     if pad.op.name != "ethosu_pad":
         return None
-    convert_to_nhwc = pad.op.input_tensors[0]
+    upscale = pad.op.input_tensors[0]
+    if upscale.op.name != "ethosu_upscale":
+        return None
+    convert_to_nhwc = upscale.op.input_tensors[0]
     if convert_to_nhwc.op.name != "ethosu_convert_to_nhwc":
         return None
     read = convert_to_nhwc.op.input_tensors[0]

python/tvm/relay/backend/contrib/ethosu/te/unary_elementwise.py

@@ -205,7 +205,10 @@ def match_ethosu_unary_elementwise(output_tensor, device_config):
     pad = unary_elementwise.op.input_tensors[0]
     if pad.op.name != "ethosu_pad":
         return None
-    convert_to_nhwc = pad.op.input_tensors[0]
+    upscale = pad.op.input_tensors[0]
+    if upscale.op.name != "ethosu_upscale":
+        return None
+    convert_to_nhwc = upscale.op.input_tensors[0]
     if convert_to_nhwc.op.name != "ethosu_convert_to_nhwc":
         return None
     read = convert_to_nhwc.op.input_tensors[0]

python/tvm/relay/backend/contrib/ethosu/util.py

@@ -47,20 +47,6 @@ class QConv2DArgs(Enum):
     WEIGHTS_SCALE = 5
 
 
-class QConv2DTransposeArgs(Enum):
-    """
-    This is a helper enum to obtain the correct index of
-    qnn.conv2d_transpose aruments.
-    """
-
-    IFM = 0
-    WEIGHTS = 1
-    IFM_ZERO_POINT = 2
-    WEIGHTS_ZERO_POINT = 3
-    IFM_SCALE = 4
-    WEIGHTS_SCALE = 5
-
-
 class RequantArgs(Enum):
     """
     This is a helper enum to obtain the correct index

python/tvm/relay/op/contrib/ethosu.py

@@ -17,8 +17,8 @@
 # pylint: disable=ungrouped-imports, import-outside-toplevel
 """Arm(R) Ethos(TM)-U NPU supported operators."""
 import functools
-
 from typing import Dict, List, Tuple, Callable, Optional
+
 import numpy as np  # type: ignore
 
 import tvm  # type: ignore
@@ -1239,16 +1239,17 @@ def is_valid(self) -> bool:
         Checks whether image.resize2d has compatible attributes with HW.
         """
 
-        def check_compatible_size(mode, method, upscale_size, height, width):
+        def check_compatible_size(mode, method, upscale_size, ifm_size):
             """Checking the provided upscale_size is compatible with the NPU. The NPU only
             supports upsampling when the upsampling size is 2 * input_size, or when there is
             no upsampling to be done, so check that this is the case. In the special case of
             resize_bilinear with align_corners=True, the NPU only supports an upsampling
             size of 2 * input_size - 1."""
             delta = 1 if mode == "align_corners" and method == "linear" else 0
-            upscale_height = height * 2 - delta
-            upscale_width = width * 2 - delta
-            return upscale_height == int(upscale_size[0]) and upscale_width == int(upscale_size[1])
+            upscale_size = np.array(upscale_size)
+            ifm_size = np.array(ifm_size)
+            ifm_upscaled = ifm_size * 2 - delta
+            return (ifm_upscaled == upscale_size).all() or (ifm_size == upscale_size).all()
 
         tensor_params = [self.ifm, self.ofm]
         if not check_valid_dtypes(tensor_params, supported_dtypes=[np.int8]):
@@ -1265,8 +1266,7 @@ def check_compatible_size(mode, method, upscale_size, height, width):
             self.coordinate_transformation_mode,
             self.method,
             self.size,
-            self.ifm.shape[1],
-            self.ifm.shape[2],
+            self.ifm.shape[1:3],
         ):
             return False
         if self.rounding_method != "":

tests/python/contrib/test_ethosu/test_codegen.py

@@ -1023,10 +1023,7 @@ def create_model():
 @pytest.mark.parametrize("accel_type", ACCEL_TYPES)
 @pytest.mark.parametrize(
     "ifm_shape,size",
-    [
-        [(1, 2, 2, 1), (4, 4)],
-        [(1, 4, 7, 3), (8, 14)],
-    ],
+    [[(1, 2, 2, 1), (4, 4)], [(1, 4, 7, 3), (8, 14)], [(1, 3, 5, 3), (3, 5)]],
 )
 def test_tflite_resize2d_nearest_neighbor(accel_type, ifm_shape, size):
     align_corners = False
@@ -1048,6 +1045,7 @@ def resize_model(x):
         [(1, 4, 7, 3), (8, 14), False],
         [(1, 2, 2, 1), (3, 3), True],
         [(1, 4, 7, 3), (7, 13), True],
+        [(1, 3, 5, 3), (3, 5), False],
     ],
 )
 def test_tflite_resize2d_bilinear(accel_type, ifm_shape, size, align_corners):

tests/python/contrib/test_ethosu/test_legalize.py

@@ -1608,6 +1608,7 @@ def verify(ext_func):
     [
         [(1, 2, 2, 1), (4, 4)],
         [(1, 4, 7, 3), (8, 14)],
+        [(1, 3, 5, 3), (3, 5)],
     ],
 )
 def test_tflite_resize2d_nearest_neighbor(ifm_shape, size):
@@ -1647,21 +1648,25 @@ def representative_dataset():
         return mod
 
     def verify(ext_func):
-        identity = ext_func.body
-        in_var = identity.args[0]
+        op = ext_func.body
+        in_var = op.args[0]
 
         # check IFM
         assert tuple(in_var.checked_type.shape) == ifm_shape
         assert in_var.checked_type.dtype == dtype
 
         # check OFM
-        attrs = dict(identity.attrs)
+        attrs = dict(op.attrs)
         out_shape = (ifm_shape[0], size[0], size[1], ifm_shape[3])
-        assert tuple(identity.checked_type.shape) == out_shape
-        assert identity.checked_type.dtype == dtype
+        assert tuple(op.checked_type.shape) == out_shape
+        assert op.checked_type.dtype == dtype
 
         # Check Op attributes
-        assert attrs["upscale"] == "NEAREST"
+        if size[0] == ifm_shape[1] and size[1] == ifm_shape[2]:
+            assert op.op.name == "contrib.ethosu.identity"
+        else:
+            assert attrs["pooling_type"] == "AVG"
+            assert attrs["upscale"] == "NEAREST"
 
     rewriter = legalize.Resize2dRewriter()
     pattern_table = [
@@ -1687,6 +1692,7 @@ def verify(ext_func):
         [(1, 4, 7, 3), (8, 14), False],
         [(1, 2, 2, 1), (3, 3), True],
         [(1, 4, 7, 3), (7, 13), True],
+        [(1, 3, 5, 3), (3, 5), False],
     ],
 )
 def test_tflite_resize2d_bilinear(ifm_shape, size, align_corners):
@@ -1725,28 +1731,31 @@ def representative_dataset():
         return mod
 
     def verify(ext_func):
-        avg_pool = ext_func.body
-        in_var = avg_pool.args[0]
+        op = ext_func.body
+        in_var = op.args[0]
 
         # check IFM
         assert tuple(in_var.checked_type.shape) == ifm_shape
         assert in_var.checked_type.dtype == dtype
 
         # check OFM
-        attrs = dict(avg_pool.attrs)
+        attrs = dict(op.attrs)
         out_shape = (ifm_shape[0], size[0], size[1], ifm_shape[3])
-        assert tuple(avg_pool.checked_type.shape) == out_shape
-        assert avg_pool.checked_type.dtype == dtype
+        assert tuple(op.checked_type.shape) == out_shape
+        assert op.checked_type.dtype == dtype
 
         # Check Op attributes
-        assert attrs["pooling_type"] == "AVG"
-        assert attrs["upscale"] == "NEAREST"
-
-        # Check padding
-        if align_corners:
-            assert list(attrs["padding"]) == [0, 0, 0, 0]
+        if size[0] == ifm_shape[1] and size[1] == ifm_shape[2]:
+            assert op.op.name == "contrib.ethosu.identity"
         else:
-            assert list(attrs["padding"]) == [0, 0, 1, 1]
+            assert attrs["pooling_type"] == "AVG"
+            assert attrs["upscale"] == "NEAREST"
+
+            # Check padding
+            if align_corners:
+                assert list(attrs["padding"]) == [0, 0, 0, 0]
+            else:
+                assert list(attrs["padding"]) == [0, 0, 1, 1]
 
     rewriter = legalize.Resize2dRewriter()
     pattern_table = [