diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index a3e6140381a0..ee3d5867579b 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -3755,6 +3755,7 @@ def _emit_resize_topi_dynamic_roi( cubic_coeff_a: float, exclude_outside: int, extrapolation_value: float, + scales: list | None, ) -> relax.Expr: """Lower Resize with runtime ROI via TOPI, which supports Expr ROI.""" if rank == 3: @@ -3771,6 +3772,7 @@ def resize1d_dyn(d, r, s0): cubic_coeff_a, exclude_outside, extrapolation_value, + scales=scales, ) return bb.emit_te(resize1d_dyn, data, roi_spatial_vec, sizes_spatial[0]) @@ -3789,12 +3791,14 @@ def resize2d_dyn(d, r, s0, s1): bicubic_alpha=cubic_coeff_a, bicubic_exclude=exclude_outside, extrapolation_value=extrapolation_value, + scales=scales, ) return bb.emit_te(resize2d_dyn, data, roi_spatial_vec, sizes_spatial[0], sizes_spatial[1]) def resize3d_dyn(d, r, s0, s1, s2): # r is ONNX order (D,H,W) x2; TOPI expects (W,H,D) x2. + # NOTE: scales order must stay ONNX (D,H,W), unlike r which is reordered return topi.image.resize3d( d, (r[2], r[1], r[0], r[5], r[4], r[3]), @@ -3806,6 +3810,7 @@ def resize3d_dyn(d, r, s0, s1, s2): bicubic_alpha=cubic_coeff_a, bicubic_exclude=exclude_outside, extrapolation_value=extrapolation_value, + scales=scales, ) return bb.emit_te( @@ -3872,7 +3877,10 @@ def _impl_v18(cls, bb, inputs, attr, params): use_dynamic_roi = roi_dynamic_vec is not None - # Convert scales to sizes if needed. + # Convert scales to sizes if needed, preserving the original spatial scales so + # the coordinate transformation uses the exact ONNX scale value rather than the + # lossy ratio derived from floor(input * scale) / input. + original_spatial_scales = None if scales is not None: if isinstance(scales, tvm.ir.GenericConst): scales = scales.value.numpy() @@ -3880,6 +3888,7 @@ def _impl_v18(cls, bb, inputs, attr, params): scales = [int(val.value) for val in scales.values] else: raise ValueError(f"Type {type(scales)} for scale is currently unsupported.") + original_spatial_scales = [float(s) for s in scales[2:]] sizes = [] for i, dim in enumerate(x.ty.shape): @@ -3906,6 +3915,7 @@ def _impl_v18(cls, bb, inputs, attr, params): cubic_coeff_a, exclude_outside, extrapolation_value, + original_spatial_scales, ) if ndims == 3: @@ -3921,8 +3931,24 @@ def _impl_v18(cls, bb, inputs, attr, params): cubic_coeff_a, exclude_outside, extrapolation_value, + scales=original_spatial_scales, ) elif ndims == 4: + if original_spatial_scales is not None: + return bb.emit_te( + topi.image.resize2d, + x, + roi_static, + sizes, + "NCHW", + topi_mode, + coord_mode, + rounding_method, + cubic_coeff_a, + exclude_outside, + extrapolation_value, + scales=original_spatial_scales, + ) return relax.op.image.resize2d( x, size=relax.ShapeExpr(sizes), @@ -3937,6 +3963,21 @@ def _impl_v18(cls, bb, inputs, attr, params): ) else: # ndims == 5 roi3d = _topi_resize3d_roi_from_onnx_ncdhw_spatial(roi_static) + if original_spatial_scales is not None: + return bb.emit_te( + topi.image.resize3d, + x, + roi3d, + sizes, + "NCDHW", + topi_mode, + coord_mode, + rounding_method, + cubic_coeff_a, + exclude_outside, + extrapolation_value, + scales=original_spatial_scales, + ) return relax.op.image.resize3d( x, size=relax.ShapeExpr(sizes), diff --git a/python/tvm/topi/image/resize.py b/python/tvm/topi/image/resize.py index 1f4799c8ecc8..46235e1c2de9 100644 --- a/python/tvm/topi/image/resize.py +++ b/python/tvm/topi/image/resize.py @@ -41,6 +41,21 @@ def can_convert_multiply_to_intdiv(origin_size, scaled_size): return True +def _explicit_scale_matches_intdiv(origin_size, scaled_size, scale_override): + """Check whether an explicit scale (image_size / target_size) agrees with the exact + integer ratio the sizes alone would imply. + + When it does, integer division stays numerically exact and should still be used; + falling back to floating-point reciproal multiplication in that case only adds + rounding error for no benefit. + """ + if not can_convert_multiply_to_intdiv(origin_size, scaled_size): + return False + int_ratio = (scaled_size / origin_size.astype("float")).value + epsilon = 1e-5 + return abs(scale_override * int_ratio - 1.0) < epsilon + + def get_1d_indices(indices, layout="NCW"): """Get 1d indices""" (cc, inum, ic) = (0, 0, 0) @@ -145,9 +160,13 @@ def get_inx( start_x=0, end_x=-1, use_int_div=False, + scale_x_override=None, ): """Infer input x from output x with various coordinate transformation methods""" - scale_x = te.div(image_width.astype("float"), target_width.astype("float")) + if scale_x_override is not None: + scale_x = scale_x_override + else: + 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": @@ -237,6 +256,7 @@ def _resize_1d( alpha=-0.5, exclude_outside=0, out_dtype=None, + scale_x=None, ): """Perform resize operation on the data with selected method and options. @@ -315,7 +335,15 @@ def _cast_output(value, data_dtype="float32", out_dtype=None): if boxes is not None: # TODO(mbrookhart): Find an example of this raise NotImplementedError("resize1d with image boxes not yet implemented") - in_x = get_inx(x, image_width, target_width, coordinate_transformation_mode, roi[0], roi[1]) + in_x = get_inx( + x, + image_width, + target_width, + coordinate_transformation_mode, + roi[0], + roi[1], + scale_x_override=scale_x, + ) if method == "nearest_neighbor": if rounding_method == "": @@ -383,6 +411,7 @@ def resize1d( extrapolation_value=0.0, out_dtype=None, output_shape=None, + scales=None, ): """Perform resize operation on the data. @@ -438,6 +467,12 @@ def resize1d( Shape to return. If left None will be inferred (If shape is determined dynamically, pass out_dtype.shape as output_shape) + scales: tuple or None + Explicit scale factors for coordinate transformation. + If not None, scales=[scale_w] for 1D. + When provided, overrides the scale derived from input and output sizes. + Used by the ONNX frontend when a Resize node is driven by a "scales" input. + Returns ------- output : tvm.te.Tensor @@ -472,6 +507,8 @@ def resize1d( if isinstance(size[i], int): size[i] = tvm.tirx.IntImm("int32", size[i]) + scale_x = (1.0 / scales[0]) if scales is not None else None + def compute_func(*indices): return _resize_1d( indices, @@ -487,6 +524,7 @@ def compute_func(*indices): exclude_outside=bicubic_exclude, extrapolation_value=extrapolation_value, out_dtype=out_dtype, + scale_x=scale_x, ) return te.compute(output_shape, compute_func, name="resize", tag=tag.INJECTIVE) @@ -510,6 +548,8 @@ def _resize_2d( alpha=-0.5, exclude_outside=0, out_dtype=None, + scale_h=None, + scale_w=None, ): """Perform resize operation on the data with selected method and options. @@ -593,8 +633,18 @@ def _cast_output(value, data_dtype="float32", out_dtype=None): width_use_int_div = False if method == "nearest_neighbor" and coordinate_transformation_mode == "asymmetric": if rounding_method == "floor" or rounding_method == "": - 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) + if scale_h is None: + height_use_int_div = can_convert_multiply_to_intdiv(image_height, target_height) + else: + height_use_int_div = _explicit_scale_matches_intdiv( + image_height, target_height, scale_h + ) + if scale_w is None: + width_use_int_div = can_convert_multiply_to_intdiv(image_width, target_width) + else: + width_use_int_div = _explicit_scale_matches_intdiv( + image_width, target_width, scale_w + ) 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 @@ -618,6 +668,7 @@ def _cast_output(value, data_dtype="float32", out_dtype=None): roi[1], roi[3], width_use_int_div, + scale_x_override=scale_w, ) in_y = get_inx( y, @@ -627,6 +678,7 @@ def _cast_output(value, data_dtype="float32", out_dtype=None): roi[0], roi[2], height_use_int_div, + scale_x_override=scale_h, ) if method == "nearest_neighbor": @@ -756,6 +808,7 @@ def resize2d( extrapolation_value=0.0, out_dtype=None, output_shape=None, + scales=None, ): """Perform resize operation on the data. @@ -805,6 +858,12 @@ def resize2d( Shape to return. If left None will be inferred (If shape is determined dynamically, pass out_dtype.shape as output_shape) + scales: tuple or None + Explicit scale factors for coordinate transformation. + If not None, scales=[scale_h, scale_w] for 2D. + When provided, overrides the scale derived from input and output sizes. + Used by the ONNX frontend when a Resize node is driven by a "scales" input. + Returns ------- output : tvm.te.Tensor @@ -839,6 +898,9 @@ def resize2d( if isinstance(size[i], int): size[i] = tvm.tirx.IntImm("int32", size[i]) + scale_h = (1.0 / scales[0]) if scales is not None else None + scale_w = (1.0 / scales[1]) if scales is not None else None + def compute_func(*indices): return _resize_2d( indices, @@ -856,6 +918,8 @@ def compute_func(*indices): exclude_outside=bicubic_exclude, extrapolation_value=extrapolation_value, out_dtype=out_dtype, + scale_h=scale_h, + scale_w=scale_w, ) return te.compute(output_shape, compute_func, name="resize", tag=tag.INJECTIVE) @@ -976,6 +1040,9 @@ def _resize_3d( alpha=-0.5, exclude_outside=0, out_dtype=None, + scale_d=None, + scale_h=None, + scale_w=None, ): """Perform resize operation on the data with selected method and options. @@ -1066,9 +1133,33 @@ def _cast_output(value, data_dtype="float32", out_dtype=None): if boxes is not None: # TODO(mbrookhart): Find an example of this raise NotImplementedError("resize1d with image boxes not yet implemented") - in_z = get_inx(z, image_depth, target_depth, coordinate_transformation_mode, roi[2], roi[5]) - in_y = get_inx(y, image_height, target_height, coordinate_transformation_mode, roi[1], roi[4]) - in_x = get_inx(x, image_width, target_width, coordinate_transformation_mode, roi[0], roi[3]) + in_z = get_inx( + z, + image_depth, + target_depth, + coordinate_transformation_mode, + roi[2], + roi[5], + scale_x_override=scale_d, + ) + in_y = get_inx( + y, + image_height, + target_height, + coordinate_transformation_mode, + roi[1], + roi[4], + scale_x_override=scale_h, + ) + in_x = get_inx( + x, + image_width, + target_width, + coordinate_transformation_mode, + roi[0], + roi[3], + scale_x_override=scale_w, + ) if method == "nearest_neighbor": if rounding_method == "": @@ -1225,6 +1316,7 @@ def resize3d( extrapolation_value=0.0, out_dtype=None, output_shape=None, + scales=None, ): """Perform resize operation on the data. @@ -1302,6 +1394,10 @@ def resize3d( if isinstance(size[i], int): size[i] = tvm.tirx.IntImm("int32", size[i]) + scale_d = (1.0 / scales[0]) if scales is not None else None + scale_h = (1.0 / scales[1]) if scales is not None else None + scale_w = (1.0 / scales[2]) if scales is not None else None + def compute_func(*indices): return _resize_3d( indices, @@ -1321,6 +1417,9 @@ def compute_func(*indices): exclude_outside=bicubic_exclude, extrapolation_value=extrapolation_value, out_dtype=out_dtype, + scale_d=scale_d, + scale_h=scale_h, + scale_w=scale_w, ) return te.compute(output_shape, compute_func, name="resize", tag=tag.INJECTIVE) diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index d7aa987c0527..04944e9b0cc3 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -40,7 +40,7 @@ import tvm import tvm.testing -from tvm import relax +from tvm import relax, te, topi from tvm.relax.frontend.onnx import from_onnx from tvm.script import ir as I from tvm.script import relax as R @@ -9432,6 +9432,457 @@ def _visit(expr): assert seen_resize3d +@pytest.mark.parametrize( + "coord_mode, method", + [ + ("half_pixel", "nearest"), + ("pytorch_half_pixel", "nearest"), + ("asymmetric", "nearest"), + ("half_pixel", "linear"), + ], +) +def test_resize_noninteger_scales_2d(coord_mode, method): + """Non-integer scales must use the original scale in coordinate transformation. + + floor(3 * 2.5) = 7, so the recomputed ratio 3/7 = 0.4286 differs from 1/2.5 = 0.4, + causing wrong pixel mapping at boundary positions before the fix. + """ + nearest_mode_kwargs = {} + if method == "nearest": + nearest_mode_kwargs["nearest_mode"] = "round_prefer_floor" + resize_node = helper.make_node( + "Resize", + ["X", "", "scales"], + ["Y"], + mode=method, + coordinate_transformation_mode=coord_mode, + **nearest_mode_kwargs, + ) + graph = helper.make_graph( + [resize_node], + "resize_noninteger_2d", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 3, 3])], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [4], [1.0, 1.0, 2.5, 2.5])], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 7, 7])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def test_resize_noninteger_scales_1d(): + resize_node = helper.make_node( + "Resize", + ["X", "", "scales"], + ["Y"], + mode="nearest", + coordinate_transformation_mode="half_pixel", + nearest_mode="round_prefer_floor", + ) + graph = helper.make_graph( + [resize_node], + "resize_noninteger_1d", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 5])], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [3], [1.0, 1.0, 1.5])], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 7])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def test_resize_noninteger_scales_3d(): + resize_node = helper.make_node( + "Resize", + ["X", "", "scales"], + ["Y"], + mode="nearest", + coordinate_transformation_mode="asymmetric", + nearest_mode="floor", + ) + graph = helper.make_graph( + [resize_node], + "resize_noninteger_3d", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 3, 3, 3])], + initializer=[ + helper.make_tensor("scales", TensorProto.FLOAT, [5], [1.0, 1.0, 1.5, 1.5, 1.5]) + ], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 4, 4, 4])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def test_resize_dynamic_roi_noninteger_scales_1d(): + resize_node = helper.make_node( + "Resize", + ["X", "roi", "scales"], + ["Y"], + mode="linear", + coordinate_transformation_mode="half_pixel", + ) + graph = helper.make_graph( + [resize_node], + "resize_dynamic_roi_noninteger_1d", + inputs=[ + helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 3]), + helper.make_tensor_value_info("roi", TensorProto.FLOAT, [6]), + ], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [3], [1.0, 1.0, 2.5])], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 7])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def test_resize_dynamic_roi_noninteger_scales_2d(): + resize_node = helper.make_node( + "Resize", + ["X", "roi", "scales"], + ["Y"], + mode="linear", + coordinate_transformation_mode="half_pixel", + ) + graph = helper.make_graph( + [resize_node], + "resize_dynamic_roi_noninteger_2d", + inputs=[ + helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 3, 3]), + helper.make_tensor_value_info("roi", TensorProto.FLOAT, [8]), + ], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [4], [1.0, 1.0, 2.5, 2.5])], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 7, 7])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def test_resize_dynamic_roi_noninteger_scales_3d_anisotropic(): + resize_node = helper.make_node( + "Resize", + ["X", "roi", "scales"], + ["Y"], + mode="linear", + coordinate_transformation_mode="asymmetric", + ) + graph = helper.make_graph( + [resize_node], + "resize_dynamic_roi_noninteger_3d_anisotropic", + inputs=[ + helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 3, 5, 7]), + helper.make_tensor_value_info("roi", TensorProto.FLOAT, [10]), + ], + initializer=[ + helper.make_tensor("scales", TensorProto.FLOAT, [5], [1.0, 1.0, 1.5, 2.5, 3.5]) + ], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 4, 12, 24])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def test_resize_asymmetric_nearest_noninteger_scales_2d(): + """asymmetric + nearest_neighbor + floor must honor a non-integer scale override. + + Input 2x2, scale 2.4 -> output 4x4. The size-derived ratio (4/2 = 2.0) is a whole + number, but it disagrees with the actual scale, so the integer-division fast path + must not be used here -- it would give a different (wrong) pixel mapping. + """ + resize_node = helper.make_node( + "Resize", + ["X", "", "scales"], + ["Y"], + mode="nearest", + coordinate_transformation_mode="asymmetric", + nearest_mode="floor", + ) + graph = helper.make_graph( + [resize_node], + "resize_asymmetric_nearest_noninteger_scales_2d", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 2, 2])], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [4], [1.0, 1.0, 2.4, 2.4])], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 4, 4])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def test_resize_asymmetric_nearest_integer_scales_baseline(): + """Baseline: asymmetric+nearest+floor with an integer scale should work correctly. + + Which code path handles this (fast integer division vs. floating-point fallback) + is covered separately by the TOPI-level tests below; this only checks the result. + """ + resize_node = helper.make_node( + "Resize", + ["X", "", "scales"], + ["Y"], + mode="nearest", + coordinate_transformation_mode="asymmetric", + nearest_mode="floor", + ) + graph = helper.make_graph( + [resize_node], + "resize_asymmetric_nearest_integer_scales_2d", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 2, 2])], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [4], [1.0, 1.0, 2.0, 2.0])], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 4, 4])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +@pytest.mark.parametrize( + "input_shape,scales,output_shape", + [ + ([1, 1, 4, 4], [1.0, 1.0, 2.0, 2.0], [1, 1, 8, 8]), + ([1, 1, 3, 3], [1.0, 1.0, 3.0, 3.0], [1, 1, 9, 9]), + ], +) +def test_resize_integer_scales_regression(input_shape, scales, output_shape): + resize_node = helper.make_node( + "Resize", + ["X", "", "scales"], + ["Y"], + mode="nearest", + coordinate_transformation_mode="half_pixel", + nearest_mode="round_prefer_floor", + ) + graph = helper.make_graph( + [resize_node], + "resize_integer_scales", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [len(scales)], scales)], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, output_shape)], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def test_resize_asymmetric_nearest_floor_exact_integer_scale(): + resize_node = helper.make_node( + "Resize", + ["X", "", "scales"], + ["Y"], + mode="nearest", + coordinate_transformation_mode="asymmetric", + nearest_mode="floor", + ) + graph = helper.make_graph( + [resize_node], + "resize_asymmetric_nearest_floor_exact_integer_scale", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 1, 411])], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [4], [1.0, 1.0, 1.0, 41.0])], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 1, 16851])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def _build_topi_resize2d_nearest_asymmetric_prim_func(input_shape, size, scales=None): + """Build (without running) a TOPI resize2d TIR PrimFunc. + + Uses `method="nearest_neighbor"` and `coordinate_transformation_mode="asymmetric"` + with the default rounding method, the only combination in which the + integer-division fast path can be selected. + + Parameters + ---------- + input_shape : tuple + NCHW input shape, e.g. (1, 1, 2, 2). + size: tuple + Target (out_h, out_w) spatial size + scales: tuple or None + If not None, an explicit (scale_h, scale_w) override, mirroring what the ONNX + frontend passes when a "scales" input (rather than "sizes") drives the resize. + + Returns + ------- + (resized_te_tensor, prim_func) : tuple + The TE output tensor (for its concrete output shape) and the lowered TIR + PrimFunc (to run the kernel and inspect which code path was generated). + """ + data = te.placeholder(input_shape, dtype="float32", name="data") + roi = (0.0, 0.0, 0.0, 0.0) + resized = topi.image.resize2d( + data, + roi, + size=size, + layout="NCHW", + method="nearest_neighbor", + coordinate_transformation_mode="asymmetric", + rounding_method="", + scales=scales, + ) + prim_func = te.create_prim_func([data, resized]) + return resized, prim_func + + +def _run_prim_func(prim_func, input_np, output_shape, target="llvm"): + """Compile a 1-input/1-output PrimFunc for `target` and run it on `input_np`. + + The numerical fallout of the floating-point fallback path is target-dependent: + some backends contract `a * b + c` into a single fused multiply-add, which can + round a borderline case back to tbe exact integer result and hide a rounding + error that a non-fused backend (e.g. `target="c"`) would expose. Pass + `target="c"` when the numeric result itself, not just which branch fired, + is what's being checked. + """ + built = tvm.compile(prim_func, target=target) + data_nd = tvm.runtime.tensor(input_np.astype("float32")) + out_nd = tvm.runtime.tensor(np.zeros(output_shape, dtype="float32")) + built(data_nd, out_nd) + return out_nd.numpy() + + +def _fast_path_used(prim_func): + """Return True if the lowered TIR took the integer-division fast path. + + The fast path computes the source index purely with integer division + (`T.Div`, no rounding needed) and never calls `T.floor`. The float + fallback always computes `floor(scale * out_index + eps)` and casts it + to int, so it always contains `T.floor` and never emits a bare `T.Div` + for the index computation. Checking both signals (rather than only the + absence of one) keeps the assertion meaningful in both directions. + """ + script = prim_func.script() + has_int_div = "T.Div(" in script + has_float_floor = "T.floor(" in script + assert has_int_div != has_float_floor, ( + "expected exactly one of the integer-division or floating-point " + f"code paths to appear in the lowered TIR, got:\n{script}" + ) + return has_int_div + + +@pytest.mark.parametrize( + "input_hw,output_hw", + [ + ((2, 2), (4, 4)), + ((3, 3), (6, 6)), + ], +) +def test_topi_resize2d_int_div_optimization_fires_without_scale_override(input_hw, output_hw): + in_h, in_w = input_hw + out_h, out_w = output_hw + input_shape = (1, 1, in_h, in_w) + + resized, prim_func = _build_topi_resize2d_nearest_asymmetric_prim_func( + input_shape, (out_h, out_w), scales=None + ) + + assert _fast_path_used(prim_func), ( + "expected the integer-division fast path to fire when scale_override " + "is None and the size ratio is an integer" + ) + + input_np = np.arange(np.prod(input_shape), dtype="float32").reshape(input_shape) + actual = _run_prim_func(prim_func, input_np, tuple(int(s) for s in resized.shape)) + + # Reference: ONNX asymmetric + nearest_neighbor + floor rounding is + # `in_idx = floor(out_idx * in_size / out_size)`, clamped to + # [0, in_size - 1]. Since in_size / out_size is an integer ratio here, + # this is exactly `out_idx // (out_size // in_size)`. + def ref_index(out_size, in_size): + ratio = out_size // in_size + idx = np.arange(out_size) // ratio + return np.clip(idx, 0, in_size - 1) + + y_idx = ref_index(out_h, in_h) + x_idx = ref_index(out_w, in_w) + expected = input_np[:, :, y_idx][:, :, :, x_idx] + + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize( + "input_hw,output_hw,scale_override", + [ + ((2, 2), (4, 4), (2.4, 2.4)), + ((3, 3), (6, 6), (2.5, 2.5)), + ], +) +def test_topi_resize2d_int_div_optimization_disabled_by_scale_override( + input_hw, output_hw, scale_override +): + """A scale override that disagrees with the size ratio disables the int-div fast path.""" + in_h, in_w = input_hw + out_h, out_w = output_hw + input_shape = (1, 1, in_h, in_w) + + resized, prim_func = _build_topi_resize2d_nearest_asymmetric_prim_func( + input_shape, (out_h, out_w), scales=scale_override + ) + + assert not _fast_path_used(prim_func), ( + "expected the integer-division fast path to be disabled when an " + "explicit (non-integer) scale_override is supplied" + ) + + input_np = np.arange(np.prod(input_shape), dtype="float32").reshape(input_shape) + actual = _run_prim_func(prim_func, input_np, tuple(int(s) for s in resized.shape)) + + epsilon = 1e-5 + + def ref_index(out_size, in_size, scale): + idx = np.floor(np.arange(out_size) / scale + epsilon).astype(np.int64) + return np.clip(idx, 0, in_size - 1) + + y_idx = ref_index(out_h, in_h, scale_override[0]) + x_idx = ref_index(out_w, in_w, scale_override[1]) + expected = input_np[:, :, y_idx][:, :, :, x_idx] + + np.testing.assert_array_equal(actual, expected) + + ratio = out_h // in_h + naive_idx = np.clip(np.arange(out_h) // ratio, 0, in_h - 1) + assert not np.array_equal(naive_idx, y_idx), ( + "test setup issue: chosen scale_override does not actually differ " + "from the naive integer-ratio index mapping" + ) + + +def test_topi_resize2d_int_div_fast_path_matches_integer_scale_override(): + """A scale override that agrees with the size ratio still fires the fast path, + and gives the same result as omitting the override entirely.""" + input_shape = (1, 1, 3, 3) + output_hw = (6, 6) + input_np = np.arange(np.prod(input_shape), dtype="float32").reshape(input_shape) + + fast_resized, fast_prim_func = _build_topi_resize2d_nearest_asymmetric_prim_func( + input_shape, output_hw, scales=None + ) + assert _fast_path_used(fast_prim_func) + fast_out_shape = tuple(int(s) for s in fast_resized.shape) + fast_actual = _run_prim_func(fast_prim_func, input_np, fast_out_shape) + + override_resized, override_prim_func = _build_topi_resize2d_nearest_asymmetric_prim_func( + input_shape, output_hw, scales=(2.0, 2.0) + ) + assert _fast_path_used(override_prim_func), ( + "expected the integer-division fast path to still fire when the " + "explicit scale_override agrees with the integer size ratio" + ) + override_out_shape = tuple(int(s) for s in override_resized.shape) + override_actual = _run_prim_func(override_prim_func, input_np, override_out_shape) + + np.testing.assert_array_equal(fast_actual, override_actual) + + +def test_topi_resize2d_int_div_fast_path_exact_index_with_integer_scale(): + """An explicit scale exactly matching the integer size ratio must give an exact index.""" + in_w = 411 + scale = 41.0 + out_w = int(in_w * scale) + assert out_w == 16851 + + input_shape = (1, 1, 1, in_w) + resized, prim_func = _build_topi_resize2d_nearest_asymmetric_prim_func( + input_shape, (1, out_w), scales=(1.0, scale) + ) + + input_np = np.arange(in_w, dtype="float32").reshape(input_shape) + out_shape = tuple(int(s) for s in resized.shape) + actual = _run_prim_func(prim_func, input_np, out_shape, target="c") + + assert actual[0, 0, 0, 16810] == input_np[0, 0, 0, 410], ( + f"expected exact index 410, got value {actual[0, 0, 0, 16810]} " + f"(input[409]={input_np[0, 0, 0, 409]}, input[410]={input_np[0, 0, 0, 410]})" + ) + assert _fast_path_used(prim_func), ( + "expected the integer-division fast path to fire since scale=41 " + "agrees exactly with the derived output/input size ratio" + ) + + def test_einsum(): eqn = "ij->i" einsum_node = helper.make_node("Einsum", ["x"], ["y"], equation=eqn)