[TOPI][OP]change float multiplication of resize op to integer division

python/tvm/topi/image/resize.py

@@ -23,6 +23,22 @@
 from .. import tag
 
 
+def can_convert_multiply_to_intdiv(origin_size, scaled_size):
+    """Check whether can convert multiplication to division"""
+    # Only support IntImm type
+    if not isinstance(scaled_size, tvm.tir.expr.IntImm):
+        return False
+
+    div = scaled_size / origin_size.astype("float")
+    if div.value % 1 != 0:
+        return False
+    epsilon = 1e-5
+    check = 1 / (epsilon * origin_size + epsilon)
+    if div > check:
+        return False
+    return True
+
+
 def get_1d_indices(indices, layout="NCW"):
     """Get 1d indices"""
     (cc, inum, ic) = (0, 0, 0)
@@ -119,15 +135,26 @@ def get_3d_pixel(data, layout, image_depth, image_height, image_width, n, c, z,
     return data(n, c, z, y, x, cc).astype("float")
 
 
-def get_inx(x, image_width, target_width, coordinate_transformation_mode, start_x=0, end_x=-1):
+def get_inx(
+    x,
+    image_width,
+    target_width,
+    coordinate_transformation_mode,
+    start_x=0,
+    end_x=-1,
+    use_int_div=False,
+):
     """Infer input x from output x with various coordinate transformation methods"""
     scale_x = te.div(image_width.astype("float"), target_width.astype("float"))
     if coordinate_transformation_mode == "half_pixel":
         in_x = (x + 0.5) * scale_x - 0.5
     elif coordinate_transformation_mode == "align_corners":
         in_x = (image_width - 1).astype("float") / (target_width - 1) * x
     elif coordinate_transformation_mode == "asymmetric":
-        in_x = scale_x * x
+        if use_int_div:
+            in_x = te.div(x, te.div(target_width, image_width))
+        else:
+            in_x = scale_x * x
     elif coordinate_transformation_mode == "pytorch_half_pixel":
         in_x = te.if_then_else(target_width > 1, (x + 0.5) * scale_x - 0.5, 0.0)
     elif coordinate_transformation_mode == "tf_half_pixel_for_nn":
@@ -146,8 +173,12 @@ def get_inx(x, image_width, target_width, coordinate_transformation_mode, start_
     return in_x
 
 
-def get_closest_index(in_x, rounding_method, boxes):
+def get_closest_index(in_x, rounding_method, boxes, use_int_div=False):
     """get the closest index to a value based on a certain rounding method"""
+    if use_int_div:
+        closest_x_index = in_x.astype("int32")
+        return closest_x_index
+
     if rounding_method == "round" or boxes is not None:
         closest_x_index = te.round(in_x).astype("int32")
     elif rounding_method == "round_prefer_floor":
@@ -595,6 +626,12 @@ def _cast_output(value, data_dtype="float32", out_dtype=None):
             dtype = data_dtype
         return value.astype(dtype)
 
+    height_use_int_div = False
+    width_use_int_div = False
+    if method == "nearest_neighbor" and coordinate_transformation_mode == "asymmetric":
+        height_use_int_div = can_convert_multiply_to_intdiv(image_height, target_height)
+        width_use_int_div = can_convert_multiply_to_intdiv(image_width, target_width)
+
     n, c, y, x, cc, inum, ic = get_2d_indices(indices, layout)
     box_idx = box_indices(n) if box_indices is not None else n
     if boxes is not None:
@@ -609,9 +646,23 @@ def _cast_output(value, data_dtype="float32", out_dtype=None):
         in_y = y1 * (image_height - 1) + h_scale * y
         in_x = x1 * (image_width - 1) + w_scale * x
     else:
-        in_x = get_inx(x, image_width, target_width, coordinate_transformation_mode, roi[1], roi[3])
+        in_x = get_inx(
+            x,
+            image_width,
+            target_width,
+            coordinate_transformation_mode,
+            roi[1],
+            roi[3],
+            width_use_int_div,
+        )
         in_y = get_inx(
-            y, image_height, target_height, coordinate_transformation_mode, roi[0], roi[2]
+            y,
+            image_height,
+            target_height,
+            coordinate_transformation_mode,
+            roi[0],
+            roi[2],
+            height_use_int_div,
         )
 
     if method == "nearest_neighbor":
@@ -621,8 +672,8 @@ def _cast_output(value, data_dtype="float32", out_dtype=None):
             else:
                 rounding_method = "floor"
 
-        closest_x_index = get_closest_index(in_x, rounding_method, boxes)
-        closest_y_index = get_closest_index(in_y, rounding_method, boxes)
+        closest_x_index = get_closest_index(in_x, rounding_method, boxes, width_use_int_div)
+        closest_y_index = get_closest_index(in_y, rounding_method, boxes, height_use_int_div)
 
         value = get_2d_pixel(
             data,

tests/python/topi/python/test_topi_image.py

@@ -86,6 +86,9 @@ def test_resize2d():
     verify_resize2d(6, 32, 64, 64, 20, 20, "NHWC")
     for layout in ["NCHW", "NHWC"]:
         verify_resize2d(4, 16, 32, 32, 50, 50, layout, "asymmetric", method="nearest_neighbor")
+        verify_resize2d(4, 16, 32, 32, 64, 50, layout, "asymmetric", method="nearest_neighbor")
+        verify_resize2d(4, 16, 32, 32, 50, 96, layout, "asymmetric", method="nearest_neighbor")
+        verify_resize2d(4, 16, 32, 32, 96, 96, layout, "asymmetric", method="nearest_neighbor")
         verify_resize2d(4, 16, 32, 32, 50, 50, layout, "align_corners", method="nearest_neighbor")
         verify_resize2d(4, 16, 32, 32, 50, 50, layout, "half_pixel", method="nearest_neighbor")
         verify_resize2d(4, 16, 32, 32, 50, 50, layout, "asymmetric", method="linear")

tests/python/topi/python/test_topi_upsampling.py

@@ -101,6 +101,46 @@ def check_target(target, dev):
         check_target(target, dev)
 
 
+def test_int_div_upsampling():
+    """Test whether upsampling op is tilable when scale_h and scale_w is integer.
+
+    Compute_at cannot work correctly in the original floating-point multiplication.
+    After using integer division,compute_at can work correctly and reduce the
+    capacity of cache buffer.
+
+    In this test case, scale_h and scale_w are set to integers, the size
+    of cache buffer should be equal to (h_i/scale_h * w_i/scale_w * c_i).
+    """
+    dtype = "int8"
+    scale_h = 2
+    scale_w = 2
+
+    x = te.placeholder([1, 32, 64, 64], dtype, "x")
+    y = topi.nn.upsampling(x, scale_h, scale_w)
+    func = te.create_prim_func([x, y])
+
+    s = tvm.tir.Schedule(func)
+    block = s.get_block("resize")
+    cache = s.cache_read(block, 0, "local")
+    n, c, h, w = s.get_loops(block)
+    s_factor = 8
+    c_o, c_i = s.split(c, factors=[None, s_factor])
+    h_o, h_i = s.split(h, factors=[None, s_factor])
+    w_o, w_i = s.split(w, factors=[None, s_factor])
+    s.reorder(n, c_o, h_o, w_o, h_i, w_i, c_i)
+    s.compute_at(cache, w_o)
+    wanted_rt = s_factor**3 / (scale_h * scale_w)
+
+    def analyze_upsampling_allocate(stmt):
+        if isinstance(stmt, tvm.tir.stmt.Allocate):
+            tvm.testing.assert_allclose(stmt.extents[0].value, wanted_rt)
+
+    lowerd_irmodule = tvm.lower(s.mod["main"])
+    tvm.tir.stmt_functor.post_order_visit(
+        lowerd_irmodule.functions.items()[0][1].body, analyze_upsampling_allocate
+    )
+
+
 @tvm.testing.uses_gpu
 def test_upsampling():
     # nearest_neighbor - NCHW
@@ -269,3 +309,4 @@ def test_upsampling3d():
 if __name__ == "__main__":
     test_upsampling()
     test_upsampling3d()
+    test_int_div_upsampling()