[TOPI] Use integer arithmetic for topi.image.resize

[Lunderberg]

Prior to this commit, floating point expressions were used to map between different-sized pixel arrays. These floating point expressions are less aggressively optimized by RewriteSimplifier, which can prevent some optimizations

This was first noticed during investigation into issue #13508. Benchmarks of topi.image.resize showed 1000x and 50x performance improvements using the LLVM and CUDA backends, respectively, by using integer expressions instead of floating point. This performance improvement is partly driven by enabling
tir.transform.VectorizeLoops to recognize vectorizable indices, where the round-trip through floating point previously prevented that optimization.

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• No users to tag found in teams: topi _{See #10317 for details}

_{Generated by tvm-bot}

[Lunderberg]

The performance testing script and the results on the LLVM, CUDA, and Vulkan backends are included below.

Click to expand test script

First, apply the following diff to patch float-indexing back into topi.image.resize.

diff --git a/python/tvm/topi/image/resize.py b/python/tvm/topi/image/resize.py
index 0383dd7ae..25675b614 100644
--- a/python/tvm/topi/image/resize.py
+++ b/python/tvm/topi/image/resize.py
@@ -661,6 +661,13 @@ def _resize_2d(
             roi[2],
         )
 
+    import os
+
+    use_float_indexing = "float" in os.environ.get("RESIZE_INDEXING", "")
+    if use_float_indexing:
+        in_x = in_x.astype("float32")
+        in_y = in_y.astype("float32")
+
     if method == "nearest_neighbor":
         if rounding_method == "":
             if coordinate_transformation_mode == "align_corners":

Then, run the following script.

#!/usr/bin/env python3

import os
import sys
import time

import numpy as np
import pytest

import tvm
from tvm import topi, te


pytest_plugins = [
    "pytest-benchmark",
    "tvm.testing.plugin",
]


def resize2d_dx_compute(inp, dy):
    """compute definition for resize2d_dx op"""
    size = (64, 32)
    layout = "NCHW"
    method = "cubic"
    coord_trans = "half_pixel"
    rounding_method = ""
    cubic_alpha = -0.75
    cubic_exclude = 0
    out_dtype = "float32"

    out = topi.image.resize2d(
        inp,
        (None, None, None, None),
        size,
        layout,
        method,
        coord_trans,
        rounding_method,
        bicubic_alpha=cubic_alpha,
        bicubic_exclude=cubic_exclude,
        out_dtype=out_dtype,
    )
    grads = tvm.te.gradient(out, [inp], head=dy)
    return grads


resize_indexing = tvm.testing.parameter("resize_int_indexing", "resize_float_indexing")
constrained_booleans = tvm.testing.parameter(
    by_dict={"kApplyConstraintsToBooleanBranches": "true", "": ""}
)


def _run_benchmark(
    benchmark, stage, resize_indexing, constrained_booleans="", target=None, dev=None
):
    os.environ["RESIZE_INDEXING"] = resize_indexing
    os.environ["REMOVE_NO_OP_CONSTRAINED_BOOLEANS"] = constrained_booleans

    inp_shape = (32, 3, 32, 32)
    dy_shape = (32, 3, 64, 32)

    inp = tvm.te.placeholder(inp_shape, name="inp")
    dy = tvm.te.placeholder(dy_shape, name="dy")

    if stage == "topi":
        benchmark(resize2d_dx_compute, inp, dy)
        return
    else:
        grad = resize2d_dx_compute(inp, dy)[0]

    # if resize_indexing == "resize_float_indexing" and constrained_booleans:
    if constrained_booleans:
        pytest.skip("Runs too slowly to effectively benchmark")

    target = tvm.target.Target(target)
    with target:
        if "gpu" in target.keys:
            sch = topi.cuda.injective.schedule_injective(grad)
        else:
            sch = topi.x86.injective.schedule_injective(grad)

    if stage == "tvm.lower":
        benchmark(tvm.lower, sch, [inp, dy, grad], simple_mode=True)
        return
    else:
        tvm.lower(sch, [inp, dy, grad], simple_mode=True)

    if stage == "tvm.build":
        benchmark(tvm.build, sch, [inp, dy, grad], target=target)
        return
    else:
        func = tvm.build(sch, [inp, dy, grad], target=target)

    inp_np = np.random.uniform(size=inp_shape).astype(inp.dtype)
    dy_np = np.random.uniform(size=dy_shape).astype(inp.dtype)

    inp_tvm = tvm.nd.array(inp_np, dev)
    dy_tvm = tvm.nd.array(dy_np, dev)
    grad_tvm = tvm.nd.empty(grad.shape, grad.dtype, dev)

    def execute_read():
        func(inp_tvm, dy_tvm, grad_tvm)
        dev.sync()

    if stage == "execute":
        iterations = 10 if "llvm" == target.kind.name else 100
        benchmark.pedantic(execute_read, iterations=iterations, warmup_rounds=5, rounds=10)
        return


def test_benchmark_topi(benchmark, resize_indexing):
    _run_benchmark(benchmark, "topi", resize_indexing)


def test_benchmark_lowering(benchmark, resize_indexing, constrained_booleans, target):
    _run_benchmark(benchmark, "tvm.lower", resize_indexing, constrained_booleans, target)


def test_benchmark_build(benchmark, resize_indexing, constrained_booleans, target):
    _run_benchmark(benchmark, "tvm.build", resize_indexing, constrained_booleans, target)


def test_benchmark_execute(benchmark, resize_indexing, constrained_booleans, target, dev):
    _run_benchmark(benchmark, "execute", resize_indexing, constrained_booleans, target, dev)


if __name__ == "__main__":
    sys.exit(pytest.main(sys.argv))

[Lunderberg]

@comaniac @tqchen This PR should be the last step in closing issue #13508

@huanmei9 @wrongtest-intellif If I understand #12315 correctly, this PR is effectively a stronger form of the same type of simplification. The main difference is that where #12315 applied to cases where the input index was an integer multiple of the resized index, this PR applies to any of the integer to integer indexing computations in topi.image.resize. Since the two PRs are closely related, any review/feedback would be appreciated.

[vinx13]

this sentence is not complete

[Lunderberg]

Thank you, and updated.

[wrongtest-intellif]

this PR is effectively a stronger form

Yes. Many thanks for the extension!

For #12315 it only try handle limited cases due to lack of fraction representation.

It originally aims to solve tiling issue because backend interger analysis works poorly with float values. We could not tile the resize space properly before... I think after this PR every modes of resize can be well scheduled :).

[Lunderberg]

@wrongtest-intellif Thank you, and that makes sense! I wanted to make sure that I was extending, and not removing, existing functionality.

Good point on the fraction representation. I was originally uncertain whether it would be overengineering, as compared to writing out the explicit integer formulas. The fun surprise for me was realizing that this approach would automatically handle cases where a user-provided parameter could be inferred to be a fraction (e.g. resizing with a cubic spline based on user-provided spline stiffness alpha).

[Lunderberg]

Cleaning out old PRs, rebasing this onto main. If it passes CI, great, but if not this clearly isn't a critical PR to land.