[TIR][Fix] IndexDataTypeNormalizer not unwrapping float casting

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 = []

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);
   }

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")

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()