diff --git a/python/tvm/tir/buffer.py b/python/tvm/tir/buffer.py index 726d5d1c988c..c2c158c77f78 100644 --- a/python/tvm/tir/buffer.py +++ b/python/tvm/tir/buffer.py @@ -179,7 +179,7 @@ def offset_of(self, indices): def __getitem__(self, indices): from ..arith import Analyzer # pylint: disable=import-outside-toplevel - from .expr import BufferLoad, Ramp # pylint: disable=import-outside-toplevel + from .expr import BufferLoad, Ramp, const # pylint: disable=import-outside-toplevel from .stmt import BufferRegion # pylint: disable=import-outside-toplevel if not isinstance(indices, (tuple, list)): @@ -195,7 +195,11 @@ def __getitem__(self, indices): stop = self.shape[i] if index.stop is None else index.stop region.append(Range.from_min_extent(start, analyzer.simplify(stop - start))) else: - region.append(Range.from_min_extent(index, 1)) + region.append( + Range.from_min_extent( + index, const(1, index.dtype) if isinstance(index, PrimExpr) else 1 + ) + ) return BufferRegion(self, region) else: expr_indices = [] diff --git a/src/tir/ir/data_type_rewriter.cc b/src/tir/ir/data_type_rewriter.cc index f0f0d84644fe..8da7cfdd5b97 100644 --- a/src/tir/ir/data_type_rewriter.cc +++ b/src/tir/ir/data_type_rewriter.cc @@ -574,7 +574,10 @@ PrimExpr IndexDataTypeNormalizer::VisitExpr_(const VarNode* op) { } PrimExpr IndexDataTypeNormalizer::VisitExpr_(const CastNode* op) { - if (is_enabled_) { + // Unwrap the cast only when the dtype of this cast is integer dtype. + // When the dtype of this cast is not integer dtype, it means that this cast + // has some other purpose, and we should not unwrap the cast. + if (is_enabled_ && op->dtype.is_int()) { PrimExpr value = IndexDataTypeNormalizer::VisitExpr(op->value); return value->dtype == target_data_type_ ? value : Cast(target_data_type_, value); } diff --git a/tests/python/unittest/test_te_create_primfunc.py b/tests/python/unittest/test_te_create_primfunc.py index f78dc458d9d3..4b8d857e8619 100644 --- a/tests/python/unittest/test_te_create_primfunc.py +++ b/tests/python/unittest/test_te_create_primfunc.py @@ -689,6 +689,72 @@ def test_argmax(): tvm.ir.assert_structural_equal(prim_func, argmax_expected) +def te_resize2d_symbolic(): + oh = tir.Var("oh", "int64") + ow = tir.Var("ow", "int64") + roi = (0.0, 0.0, 0.0, 0.0) + A = te.placeholder((2, 3, 128, 128), "float32", name="A") + B = topi.image.resize2d( + A, + roi, + size=(oh, ow), + method="nearest_neighbor", + coordinate_transformation_mode="asymmetric", + rounding_method="round", + ) + return [A, B] + + +@T.prim_func +def tir_resize2d_symbolic( + A: T.Buffer[(T.int64(2), T.int64(3), T.int64(128), T.int64(128)), "float32"], + var_resize: T.handle, +): + T.func_attr({"global_symbol": "main", "tir.noalias": True}) + oh = T.var("int64") + ow = T.var("int64") + resize = T.match_buffer(var_resize, [T.int64(2), T.int64(3), oh, ow], dtype="float32") + for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), oh, ow): + with T.block("resize"): + v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) + T.reads(A[v_i0, v_i1, T.int64(0) : T.int64(128), T.int64(0) : T.int64(128)]) + T.writes(resize[v_i0, v_i1, v_i2, v_i3]) + resize[v_i0, v_i1, v_i2, v_i3] = A[ + v_i0, + v_i1, + T.max( + T.min( + T.Cast( + "int64", + T.round( + T.float32(128) / T.Cast("float32", oh) * T.Cast("float32", v_i2), + dtype="float32", + ), + ), + T.int64(127), + ), + T.int64(0), + ), + T.max( + T.min( + T.Cast( + "int64", + T.round( + T.float32(128) / T.Cast("float32", ow) * T.Cast("float32", v_i3), + dtype="float32", + ), + ), + T.int64(127), + ), + T.int64(0), + ), + ] + + +def test_resize2d_symbolic(): + _check_workload(te_resize2d_symbolic, tir_resize2d_symbolic, index_dtype_override="int64") + + def test_extern_with_explicit_buffer_access(): def te_extern(): A = te.placeholder((128, 128), name="A") diff --git a/tests/python/unittest/test_tvmscript_regression.py b/tests/python/unittest/test_tvmscript_regression.py index d063c0fcab7f..44d3036596ba 100644 --- a/tests/python/unittest/test_tvmscript_regression.py +++ b/tests/python/unittest/test_tvmscript_regression.py @@ -17,6 +17,7 @@ import numpy import tvm +import tvm.testing from tvm.script import tir as T @@ -73,9 +74,17 @@ def func_ref(): tvm.ir.assert_structural_equal(test_case, func_ref) +def test_tir_buffer_region_extent_correct_dtype(): + @T.prim_func + def func(A: T.Buffer[(T.int64(16), T.int64(1)), "float32"]): + for i in T.grid(T.int64(16)): + with T.block("block"): + vi = T.axis.remap("S", [i]) + T.reads(A[vi, T.int64(0) : T.int64(1)]) + T.evaluate(0) + + assert func.body.block.body.body.block.reads[0].region[0].extent.dtype == "int64" + + if __name__ == "__main__": - a = numpy.zeros((10, 10), dtype="int8") - test_multi_element_array_in_outmost_namespace() - test_different_dtype_assignment_to_var() - b = 1 - test_var_capturing_order() + tvm.testing.main()