[Performance] image.resize2d nearest (asymmetric, non-floor rounding) is ~2x slower on CUDA since v0.24.0 (int64 index math from #18838) #20956

Thanks for participating in the TVM community! We use https://discuss.tvm.ai for any general usage questions and discussions. The issue tracker is used for actionable items such as feature proposals discussion, roadmaps, and bug tracking. You are always welcomed to post on the forum first 😸

Issues that are inactive for a period of time may get closed. We adopt this policy so that we won't lose track of actionable issues that may fall at the bottom of the pile. Feel free to reopen a new one if you feel there is an additional problem that needs attention when an old one gets closed.

Expected behavior

image.resize2d with method="nearest_neighbor", coordinate_transformation_mode="asymmetric" and an integer
upscale factor should run about as fast for round / round_prefer_floor / round_prefer_ceil / ceil as it does
for floor: the source index is still x / factor rounded, which is cheap integer math.

Actual behavior

Since v0.24.0, the non-floor rounding modes run ~2x slower on CUDA than floor for this case, and ~1.9x slower
than v0.23.0. The cause is #18838 ("[BugFix][TOPI] Fix resize accuracy issue with non-floor rounding"). That fix is
correct and should stay: before it, an integer upscale used the integer-division index path for every rounding mode,
so round, round_prefer_ceil and ceil silently computed floor (the repro below shows this on v0.23.0). After it,
only floor keeps the integer path. The other modes go through

closest_x_index = te.round(in_x).astype("int32")   # and floor/ceil(in_x ± 0.5) etc.

on a float coordinate, and in the generated kernel the whole index expression becomes int64_t with a
float conversion, nearbyintf/roundf/floorf, and a clamp per output element (int64_t appears 65 times in the
kernel on main, 4 times for floor). resize2d nearest is a pure copy kernel, so the index math is most of its time.

Kernel time on CUDA (RTX 4060 Ti, sm_89, CUDA 12; ncu gpu__time_duration.sum, median of 10 fresh processes,
x float32[1,8,56,56,8] NCHW8c → 112×112, rounding_method="round"):

The repro script below shows the same in wall-clock time_evaluator on the same GPU:

                     v0.23.0   v0.24.0   main      main + fix
round                10.32 us  24.68 us  24.69 us  11.90 us
round_prefer_floor   10.31 us  22.06 us  22.06 us  10.08 us
round_prefer_ceil    10.31 us  22.18 us  22.18 us  11.20 us
ceil                 10.31 us  22.18 us  22.18 us  11.20 us
floor                10.31 us   9.41 us   9.41 us  10.37 us
v0.23.0: every mode returns the floor output.  main + fix: bit-identical to main for all modes.

Proposed fix

When the upscale factor k = target_size / image_size is a power of two, in_x = x / k is exact in float32, so the
rounded index can be computed with one integer floordiv, with the same tie rules as get_closest_index:

def _is_pow2_upscale(origin_size, scaled_size):
    """Whether scaled_size == origin_size * 2**k (k >= 1) and the int-div path applies."""
    if not can_convert_multiply_to_intdiv(origin_size, scaled_size):
        return False
    factor = int(scaled_size) // int(origin_size)
    return factor > 1 and factor & (factor - 1) == 0


def _round_div(x, factor, rounding_method):
    """Round x / factor for a non-negative integer x and a power-of-two factor, with the same
    tie rules as get_closest_index: "round" is half-to-even (te.round), "round_prefer_floor"
    rounds ties down, "round_prefer_ceil" rounds ties up."""
    half = factor // 2
    if rounding_method == "round":
        # ties go up when x // factor is odd, down when it is even
        offset = half - 1 + te.floormod(te.floordiv(x, factor), 2)
    elif rounding_method == "round_prefer_floor":
        offset = half - 1
    elif rounding_method == "round_prefer_ceil":
        offset = half
    else:  # ceil
        offset = factor - 1
    return te.floordiv(x + offset, factor).astype("int32")

and in _resize_2d, for nearest_neighbor + asymmetric with one of these four rounding modes and a power-of-two
upscale, closest_x_index = _round_div(x, target_width // image_width, rounding_method) (same for y) instead of
get_closest_index(...). Everything else, including the existing edge clamp, is unchanged, and other factors keep the
float path.

Correctness: on main the patched outputs are bit-identical to unpatched main for all five rounding modes on the
benchmark input (GPU), and on an LLVM sweep of 90 settings (5 rounding modes × 6 shape pairs × 3 layouts) that also
checks the result against a float32 model of the rounding rule. test_transform_legalize_ops_image.py and
test_op_image.py pass (27/27); test_image_resize2d covers exactly this case, so its expected TIR is regenerated.
I restricted the integer path to power-of-two factors on purpose: an integer formula for every integer factor would
differ from the current float path at ties (5,279 of 64,168 settings checked), because there x / k is not exact in
float32.

The remaining ~12% over v0.23.0 is the edge clamp and one extra floormod for ties-to-even, which correct rounding needs
(v0.23.0 was faster because it computed floor for every mode).

I have the patch ready and I'll open a PR with the fix — filing this issue first to track it.

Environment

• TVM: the slowdown first appears in v0.24.0 ([BugFix][TOPI] Fix resize accuracy issue with non-floor rounding #18838) and is unchanged on main (1fa4b8627d).
• Target: CUDA, GPU sm_89 (RTX 4060 Ti), CUDA 12.
• OS: Linux x86_64.

Steps to reproduce

Minimal script. For each rounding mode it builds a single image.resize2d relax function, compiles it with the relax
CUDA pipeline, counts int64_t and float rounding calls in the generated kernel, checks whether the output equals the
floor output, and times it. It runs standalone on one TVM version (v0.2x releases and current main) and prints
tvm.__version__, so the output is self-labeling.

# repro_resize2d.py
import os
import re
import numpy as np
import tvm
from tvm import relax
from tvm.relax.backend.cuda.pipeline import get_default_pipeline

ROUNDING = os.environ.get("ROUNDING", "round,round_prefer_floor,round_prefer_ceil,ceil,floor").split(",")
SHAPE = (1, 8, 56, 56, 8)  # NCHW8c, upsampled 56x56 -> 112x112

print(f"tvm.__version__ = {tvm.__version__}")
print(f"tvm.__file__    = {tvm.__file__}")

try:  # current main: target / runtime live in tvm_tirx
    import tvm_tirx

    # main compiles with NVRTC by default; TVM_CUDA_COMPILER=nvcc selects nvcc instead
    cc = os.environ.get("TVM_CUDA_COMPILER")
    target = tvm_tirx.target.Target(
        {"kind": "cuda", "backend_config": {"cuda": {"compiler": cc}}} if cc else "cuda")
    dev = tvm_tirx.runtime.cuda(0)
    to_nd = tvm_tirx.runtime.tensor
except ImportError:  # v0.2x releases
    target = tvm.target.Target("cuda")
    dev = tvm.cuda(0)
    to_nd = getattr(tvm.runtime, "tensor", None) or tvm.nd.array

rng = np.random.default_rng(0)  # fixed seed: same input across versions/patches
data = to_nd(rng.standard_normal(SHAPE).astype("float32"), dev)


def cuda_source(mod):
    # best effort: the generated CUDA source, searched through the module import tree
    try:
        src = mod.inspect_source() if hasattr(mod, "inspect_source") else mod.get_source()
    except Exception:
        src = ""
    if "__global__" in src:
        return src
    for m in getattr(mod, "imports", None) or getattr(mod, "imported_modules", None) or []:
        src = cuda_source(m)
        if src:
            return src
    return ""


def bench(rounding_method):
    bb = relax.BlockBuilder()
    tensor_type = getattr(relax, "TensorStructInfo", None) or relax.TensorType  # renamed on main
    x = relax.Var("x", tensor_type(SHAPE, "float32"))
    with bb.function("main", [x]):
        with bb.dataflow():
            out = bb.emit(relax.op.image.resize2d(
                x, (112, 112), layout="NCHW8c", method="nearest_neighbor",
                coordinate_transformation_mode="asymmetric", rounding_method=rounding_method))
            gv = bb.emit_output(out)
        bb.emit_func_output(gv)
    with target:
        ex = tvm.compile(get_default_pipeline(target)(bb.get()), target=target)
    src = cuda_source(ex.mod)
    vm = relax.VirtualMachine(ex, dev)
    out = vm["main"](data)  # warmup
    res = vm.time_evaluator("main", dev, number=200, repeat=20)(data)
    return res.mean * 1e6, out.numpy(), src.count("int64_t"), bool(re.search(r"roundf|nearbyintf|rintf|floorf|ceilf", src))


results = {r: bench(r) for r in ROUNDING}
print("\n===== resize2d nearest, asymmetric, NCHW8c 56x56 -> 112x112 =====")
floor_out = results["floor"][1] if "floor" in results else None
print(f"{'rounding':<20} {'time (us)':>10} {'int64_t in kernel':>18} {'float rounding':>15} {'same output as floor':>21}")
for r, (us, o, n64, fround) in results.items():
    same = "-" if floor_out is None or r == "floor" else str(bool(np.array_equal(o, floor_out)))
    print(f"{r:<20} {us:10.2f} {n64:18d} {str(fround):>15} {same:>21}")

# Save outputs so a stock-vs-patched run can verify bit-identical results.
out_path = os.environ.get("SAVE_OUT")
if out_path:
    np.savez(out_path, **{r: o for r, (_, o, _, _) in results.items()})
    print(f"\nsaved outputs -> {out_path}")

Run it once per TVM build, e.g.:

for V in v0.23.0 v0.24.0; do
  PYTHONPATH=/path/to/tvm-$V/python python repro_resize2d.py
done
# current main (NVRTC by default; TVM_CUDA_COMPILER=nvcc to use nvcc instead)
SAVE_OUT=main.npz python repro_resize2d.py

For clock-invariant numbers, wrap the same run in ncu --metrics gpu__time_duration.sum.

Triage

• needs-triage
• topi
• performance