[TOPI] upsample operator 'NCHWinic' format support. (#4791)

* [TOPI] upsample operator 'NCHWinic' format support.

some hardware accelerator ask packed format data like NCHWinic to fit the
hardware resource, here add upsample NCHWinic format support to help
such requirement.

* address review comments, add assert for 'else must be NCHWxc' logic.

topi/python/topi/testing/bilinear_resize_python.py

@@ -18,14 +18,20 @@
 """Bilinear Scale in python"""
 import math
 import numpy as np
+from topi.util import nchw_pack_layout
 
 def bilinear_resize_python(image, out_size, layout, coordinate_transformation_mode="align_corners"):
     """ Bilinear scaling using python"""
     (new_h, new_w) = out_size
+    (ib, ic) = (1, 1)
 
     if layout == 'NHWC':
         (batch, h, w, channel) = image.shape
         scaled_image = np.ones((batch, new_h, new_w, channel))
+    # NCHWinic
+    elif nchw_pack_layout(layout):
+        (batch, channel, h, w, ib, ic) = image.shape
+        scaled_image = np.ones((batch, channel, new_h, new_w, ib, ic))
     else:
         (batch, channel, h, w) = image.shape
         scaled_image = np.ones((batch, channel, new_h, new_w))
@@ -40,47 +46,59 @@ def bilinear_resize_python(image, out_size, layout, coordinate_transformation_mo
     def _lerp(A, B, t):
         return A * (1.0 - t) + B * t
 
-    for b in range(batch):
-        for i in range(channel):
-            for j in range(new_h):
-                for k in range(new_w):
-                    if coordinate_transformation_mode == "half_pixel":
-                        in_y = (j + 0.5) * height_scale - 0.5
-                    else:
-                        in_y = j * height_scale
-                    y0 = int(math.floor(in_y))
-                    y1 = max(min(y0 + 1, h - 1), 0)
-                    y0 = max(y0, 0)
-                    y_lerp = in_y - math.floor(in_y)
+    def _img_scale(b, m, i, n):
+        for j in range(new_h):
+            for k in range(new_w):
+                if coordinate_transformation_mode == "half_pixel":
+                    in_y = (j + 0.5) * height_scale - 0.5
+                else:
+                    in_y = j * height_scale
+                y0 = int(math.floor(in_y))
+                y1 = max(min(y0 + 1, h - 1), 0)
+                y0 = max(y0, 0)
+                y_lerp = in_y - math.floor(in_y)
+
+                if coordinate_transformation_mode == "half_pixel":
+                    in_x = (k + 0.5) * width_scale - 0.5
+                else:
+                    in_x = k * width_scale
+                x0 = int(math.floor(in_x))
+                x1 = max(min(x0 + 1, w - 1), 0)
+                x0 = max(x0, 0)
+                x_lerp = in_x - math.floor(in_x)
 
-                    if coordinate_transformation_mode == "half_pixel":
-                        in_x = (k + 0.5) * width_scale - 0.5
-                    else:
-                        in_x = k * width_scale
-                    x0 = int(math.floor(in_x))
-                    x1 = max(min(x0 + 1, w - 1), 0)
-                    x0 = max(x0, 0)
-                    x_lerp = in_x - math.floor(in_x)
+                if layout == 'NHWC':
+                    A = image[b][y0][x0][i]
+                    B = image[b][y0][x1][i]
+                    C = image[b][y1][x0][i]
+                    D = image[b][y1][x1][i]
+                elif nchw_pack_layout(layout):
+                    A = image[b][i][y0][x0][m][n]
+                    B = image[b][i][y0][x1][m][n]
+                    C = image[b][i][y1][x0][m][n]
+                    D = image[b][i][y1][x1][m][n]
+                else:
+                    A = image[b][i][y0][x0]
+                    B = image[b][i][y0][x1]
+                    C = image[b][i][y1][x0]
+                    D = image[b][i][y1][x1]
 
-                    if layout == 'NHWC':
-                        A = image[b][y0][x0][i]
-                        B = image[b][y0][x1][i]
-                        C = image[b][y1][x0][i]
-                        D = image[b][y1][x1][i]
-                    else:
-                        A = image[b][i][y0][x0]
-                        B = image[b][i][y0][x1]
-                        C = image[b][i][y1][x0]
-                        D = image[b][i][y1][x1]
+                top = _lerp(A, B, x_lerp)
+                bottom = _lerp(C, D, x_lerp)
 
-                    top = _lerp(A, B, x_lerp)
-                    bottom = _lerp(C, D, x_lerp)
+                pixel = np.float32(_lerp(top, bottom, y_lerp))
 
-                    pixel = np.float32(_lerp(top, bottom, y_lerp))
+                if layout == 'NHWC':
+                    scaled_image[b][j][k][i] = pixel
+                elif nchw_pack_layout(layout):
+                    scaled_image[b][i][j][k][m][n] = pixel
+                else:
+                    scaled_image[b][i][j][k] = pixel
 
-                    if layout == 'NHWC':
-                        scaled_image[b][j][k][i] = pixel
-                    else:
-                        scaled_image[b][i][j][k] = pixel
+    for b in range(batch):
+        for m in range(ib):
+            for i in range(channel):
+                for n in range(ic):
+                    _img_scale(b, m, i, n)
 
     return scaled_image

topi/python/topi/testing/upsampling_python.py

@@ -18,6 +18,8 @@
 """Upsampling in python"""
 import math
 import numpy as np
+from topi.util import nchw_pack_layout
+
 
 def upsample_nearest(arr, scale):
     """ Populate the array by scale factor"""
@@ -44,6 +46,18 @@ def upsampling_python(data, scale, layout='NCHW'):
             for c in range(oshape[1]):
                 output_np[b, c, :, :] = upsample_nearest(data[b, c, :, :], scale)
         return output_np
+    # NCHWinic
+    if nchw_pack_layout(layout):
+        oshape = (ishape[0], ishape[1], int(round(ishape[2]*scale[0])),
+                  int(round(ishape[3]*scale[1])), ishape[4], ishape[5])
+        output_np = np.zeros(oshape, dtype=data.dtype)
+        for b in range(oshape[0]):
+            for ib in range(oshape[4]):
+                for c in range(oshape[1]):
+                    for ic in range(oshape[5]):
+                        output_np[b, c, :, :, ib, ic] = upsample_nearest(data[b, c, :, :, ib, ic], scale)
+        return output_np
+
     if layout == 'NHWC':
         oshape = (ishape[0], int(round(ishape[1]*scale[0])),
                   int(round(ishape[2]*scale[1])), ishape[3])

topi/python/topi/util.py

@@ -27,6 +27,14 @@ class InvalidShapeError(ValueError):
     """Invalid shape for a topi function. i.e. call winograd template for non-3x3 kernel)"""
     pass
 
+def nchw_pack_layout(layout_info):
+    """Check whether the layout type is NCHWinic"""
+    return layout_info[:4] == 'NCHW' and 'c' in layout_info and 'n' in layout_info
+
+def nchw_xc_layout(layout_info):
+    """Check whether the layout type is NCHWxc"""
+    return layout_info[:4] == 'NCHW' and 'c' in layout_info and layout_info[4:-1].isnumeric()
+
 def traverse_inline(s, final_op, callback):
     """Traverse computation graph and do auto inline
 

topi/tests/python/test_topi_upsampling.py

@@ -20,16 +20,26 @@ import tvm
 import topi
 import topi.testing
 import math
+from topi.util import nchw_pack_layout
 
 from common import get_all_backend
 
 def verify_upsampling(batch, in_channel, in_height, in_width, scale_h, scale_w,
-                      layout='NCHW', method="nearest_neighbor"):
+                      layout='NCHW', method="nearest_neighbor",
+                      in_batch_block = 0, in_channel_block = 0):
     if layout == 'NCHW':
         A = tvm.placeholder((batch, in_channel, in_height, in_width), name='A')
         dtype = A.dtype
         out_shape = (batch, in_channel, int(round(in_height*scale_h)), int(round(in_width*scale_w)))
         a_np = np.random.uniform(size=(batch, in_channel, in_height, in_width)).astype(dtype)
+    elif nchw_pack_layout(layout):
+        A = tvm.placeholder((batch, in_channel, in_height, in_width, in_batch_block, in_channel_block),
+                             name='A')
+        dtype = A.dtype
+        out_shape = (batch, in_channel, int(round(in_height*scale_h)), int(round(in_width*scale_w)),
+                     in_batch_block, in_channel_block)
+        a_np = np.random.uniform(size=(batch, in_channel, in_height, in_width,
+                                 in_batch_block, in_channel_block)).astype(dtype)
     elif layout == 'NHWC':
         A = tvm.placeholder((batch, in_height, in_width, in_channel), name='A')
         dtype = A.dtype
@@ -81,6 +91,22 @@ def test_upsampling():
     verify_upsampling(2, 2, 32, 32, 3.0, 3.0, method="bilinear")
     verify_upsampling(1, 64, 22, 32, 1.954545497894287, 2.0, method="bilinear")
 
+    # nearest_neighbor - NCHWinic
+    verify_upsampling(2, 2, 32, 32, in_batch_block=4, in_channel_block=8,
+                      scale_h=2.0, scale_w=2.0)
+    verify_upsampling(2, 2, 64, 64, in_batch_block=1, in_channel_block=16,
+                      scale_h=3.0, scale_w=3.0)
+    verify_upsampling(1, 4, 22, 32, in_batch_block=1, in_channel_block=16,
+                      scale_h=1.954545497894287, scale_w=2.0)
+
+    # bilinear - NCHWinic
+    verify_upsampling(2, 2, 32, 32, in_batch_block=1, in_channel_block=1,
+                      scale_h=2.0, scale_w=2.0, method="bilinear")
+    verify_upsampling(2, 2, 32, 32, in_batch_block=1, in_channel_block=1,
+                      scale_h=3.0, scale_w=3.0, method="bilinear")
+    verify_upsampling(2, 4, 22, 32, in_batch_block=1, in_channel_block=16,
+                      scale_h=1.954545497894287, scale_w=2.0, layout="NCHW1n16c", method="bilinear")
+
     # bilinear - NHWC
     verify_upsampling(2, 2, 32, 32, 2.0, 2.0, layout="NHWC", method="bilinear")
     verify_upsampling(2, 2, 32, 32, 3.0, 3.0, layout="NHWC", method="bilinear")