# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # ruff: noqa: E501 import tvm import tvm.testing from tvm import relax from tvm.relax.transform import LegalizeOps from tvm.script import relax as R from tvm.script import tirx as T def test_image_resize2d(): # fmt: off @tvm.script.ir_module class Resize2D: @R.function def main(x: R.Tensor((2, 8, 8, 3), "float32")) -> R.Tensor((2, 16, 16, 3), "float32"): gv: R.Tensor((2, 16, 16, 3), "float32") = R.image.resize2d(x, size=(16, 16), layout="NHWC", method="nearest_neighbor", coordinate_transformation_mode="asymmetric") return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 8, 8, 3), "float32")) -> R.Tensor((2, 16, 16, 3), "float32"): gv = R.call_tir(Expected.resize2d, (x,), R.Tensor((2, 16, 16, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def resize2d(rxplaceholder: T.Buffer((T.int64(2), T.int64(8), T.int64(8), T.int64(3)), "float32"), resize: T.Buffer((T.int64(2), T.int64(16), T.int64(16), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(16), T.int64(16), T.int64(3)): with T.sblock("resize"): i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[i0_1, T.int64(0):T.int64(8), T.int64(0):T.int64(8), i3_1]) T.writes(resize[i0_1, i1_1, i2_1, i3_1]) resize[i0_1, i1_1, i2_1, i3_1] = rxplaceholder[i0_1, T.max(T.min(T.Cast("int64", T.round(T.float32(0.5) * T.Cast("float32", i1_1))), T.int64(7)), T.int64(0)), T.max(T.min(T.Cast("int64", T.round(T.float32(0.5) * T.Cast("float32", i2_1))), T.int64(7)), T.int64(0)), i3_1] # fmt: on mod = LegalizeOps()(Resize2D) tvm.ir.assert_structural_equal(mod, Expected) def test_image_resize2d_symbolic(): # fmt: off @tvm.script.ir_module class Resize2D: @R.function def main(dumb_param: R.Tensor(("oh", "ow")), x: R.Tensor(("n", "c", "h", "w", 16), "float32")) -> R.Tensor(("n", "c", "oh", "ow", 16), "float32"): n = T.int64() c = T.int64() oh = T.int64() ow = T.int64() gv: R.Tensor((n, c, oh, ow, 16), "float32") = R.image.resize2d(x, size=(oh, ow), layout="NCHW16c", method="nearest_neighbor", coordinate_transformation_mode="asymmetric") return gv @tvm.script.ir_module class Expected: @R.function def main(dumb_param: R.Tensor(("oh", "ow")), x: R.Tensor(("n", "c", "h", "w", 16), "float32")) -> R.Tensor(("n", "c", "oh", "ow", 16), "float32"): n = T.int64() c = T.int64() oh = T.int64() ow = T.int64() gv = R.call_tir(Expected.resize2d, (x,), R.Tensor((n, c, oh, ow, 16), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def resize2d(var_rxplaceholder: T.handle, var_resize: T.handle): T.func_attr({"tirx.noalias": True}) c = T.int64() h = T.int64() n = T.int64() oh = T.int64() ow = T.int64() w = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, [n, c, h, w, T.int64(16)], dtype="float32") resize = T.match_buffer(var_resize, [n, c, oh, ow, T.int64(16)], dtype="float32") for i0, i1, i2, i3, i4 in T.grid(n, c, oh, ow, T.int64(16)): with T.sblock("resize"): i0_1, i1_1, i2_1, i3_1, i4_1 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4]) T.reads(rxplaceholder[i0_1, i1_1, T.int64(0) : T.max(h, T.int64(1)), T.int64(0) : T.max(w, T.int64(1)), i4_1]) T.writes(resize[i0_1, i1_1, i2_1, i3_1, i4_1]) resize[i0_1, i1_1, i2_1, i3_1, i4_1] = rxplaceholder[i0_1, i1_1, T.max(T.min(T.Cast("int64", T.round(T.Cast("float32", h) / T.Cast("float32", oh) * T.Cast("float32", i2_1), dtype="float32")), h - T.int64(1)), T.int64(0)), T.max(T.min(T.Cast("int64", T.round(T.Cast("float32", w) / T.Cast("float32", ow) * T.Cast("float32", i3_1), dtype="float32")), w - T.int64(1)), T.int64(0)), i4_1] # fmt: on mod = LegalizeOps()(Resize2D) tvm.ir.assert_structural_equal(mod, Expected) def test_image_affine_grid(): # fmt: off @tvm.script.ir_module class AffineGrid: @R.function def main(theta: R.Tensor((2, 2, 3), "float32")) -> R.Tensor((2, 2, 16, 16), "float32"): gv: R.Tensor((2, 2, 16, 16), "float32") = R.image.affine_grid(theta, size=(16, 16)) return gv @tvm.script.ir_module class Expected: @R.function def main(theta: R.Tensor((2, 2, 3), "float32")) -> R.Tensor((2, 2, 16, 16), "float32"): gv = R.call_tir(Expected.affine_grid, (theta,), R.Tensor((2, 2, 16, 16), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def affine_grid(var_theta: T.handle, var_compute: T.handle): T.func_attr({"tirx.noalias": True}) theta = T.match_buffer(var_theta, (T.int64(2), T.int64(2), T.int64(3))) compute = T.match_buffer(var_compute, (T.int64(2), T.int64(2), T.int64(16), T.int64(16))) with T.sblock("root"): T.reads() T.writes() for n, dim, i0, i1 in T.grid(T.int64(2), T.int64(2), T.int64(16), T.int64(16)): with T.sblock("compute"): v_n, v_dim, v_i0, v_i1 = T.axis.remap("SSSS", [n, dim, i0, i1]) T.reads(theta[v_n, v_dim, T.int64(0):T.int64(3)]) T.writes(compute[v_n, v_dim, v_i0, v_i1]) compute[v_n, v_dim, v_i0, v_i1] = theta[v_n, v_dim, T.int64(2)] + theta[v_n, v_dim, T.int64(1)] * (T.float32(-1.0) + T.Cast("float32", v_i0) * T.float32(0.13333332666666667)) + theta[v_n, v_dim, T.int64(0)] * (T.float32(-1.0) + T.Cast("float32", v_i1) * T.float32(0.13333332666666667)) # fmt: on mod = LegalizeOps()(AffineGrid) tvm.ir.assert_structural_equal(mod, Expected) def test_image_resize3d(): # fmt: off @tvm.script.ir_module class Resize3D: @R.function def main(x: R.Tensor((2, 3, 8, 8, 8), "float32")) -> R.Tensor((2, 3, 4, 6, 7), "float32"): gv: R.Tensor((2, 3, 4, 6, 7), "float32") = R.image.resize3d( x, size=(4, 6, 7), layout="NCDHW", method="nearest_neighbor", coordinate_transformation_mode="asymmetric", rounding_method="floor", ) return gv # fmt: on mod = LegalizeOps()(Resize3D) seen_call_tir = False seen_resize3d_relax_op = False def _visit(expr): nonlocal seen_call_tir, seen_resize3d_relax_op if isinstance(expr, relax.Call): if isinstance(expr.op, tvm.ir.Op): if expr.op.name == "relax.call_tir": seen_call_tir = True if expr.op.name == "relax.image.resize3d": seen_resize3d_relax_op = True relax.analysis.post_order_visit(mod["main"].body, _visit) assert seen_call_tir assert not seen_resize3d_relax_op if __name__ == "__main__": tvm.testing.main()