[TVMScript] Parser int64 support (#10789)

## Context

When dealing with end-to-end models, we note that some tensors may have large shapes. Thus, when designing graph-level IR, we sometimes use `int64` instead of `int32` for the shape. Below is an dense GeMM example which has `int64` input tensor shape:

```python
@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(rxplaceholder: T.Buffer[(1, 512), "float32"], rxplaceholder_1: T.Buffer[(T.int64(1000), T.int64(512)), "float32"], T_matmul_NT: T.Buffer[(1, T.int64(1000)), "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "dense", "tir.noalias": True, "op_pattern": 3})
        # body
        # with T.block("root")
        for i0_0, i1_0, i0_1, i1_1, i2_0, i0_2, i1_2, i2_1, i0_3, i1_3 in T.grid(1, 4, 1, 25, 8, 1, 10, 64, 1, 1):
            with T.block("T_matmul_NT"):
                i = T.axis.spatial(1, 0)
                j = T.axis.spatial(T.int64(1000), i1_0 * T.int64(250) + i1_1 * T.int64(10) + i1_2)
                k = T.axis.reduce(512, i2_0 * 64 + i2_1)
                T.reads(T_matmul_NT[i, j], rxplaceholder[i, k], rxplaceholder_1[j, k])
                T.writes(T_matmul_NT[i, j])
                T.block_attr({"layout_free_placeholders":[rxplaceholder_1], "meta_schedule.tiling_structure":"SSRSRS"})
                with T.init():
                    T_matmul_NT[i, j] = T.float32(0)
                T_matmul_NT[i, j] = T_matmul_NT[i, j] + rxplaceholder[i, k] * rxplaceholder_1[j, k]
```

## Problem

Though our TVMScript printer can easily print `int64` constants, the parser had poor support for `int64`. So this PR introduces some parser support for `int64`, basically about the data type of loop variables, block iterators and block read/write regions.

Besides the parser, most of the TIR schedule primitives didn't take `int64` into account in their implementations. These schedule primitives will be fixed and updated in recent future, in followup PRs.

Author: MasterJH5574

python/tvm/script/tir/node.py

@@ -97,9 +97,9 @@ def check_index(index: Union[int, PrimExpr]):
                     report_error("Negative index is not allowed during buffer access", span)
             elif isinstance(index, PrimExpr):
                 element_dtype = index.dtype.split("x", maxsplit=1)[0]
-                if element_dtype != "int32":
+                if element_dtype[:3] != "int":
                     report_error(
-                        "index expected an int32 type PrimExpr but got " + str(index.dtype),
+                        "index expected an integer type PrimExpr but got " + str(index.dtype),
                         index.span,
                     )
             else:

python/tvm/script/tir/special_stmt.py

@@ -486,7 +486,6 @@ def axis(
         if var_name in [iter_var.var.name for iter_var in block_scope.iter_vars]:
             self.context.report_error("Duplicate block axis " + var_name, self.node.span)
 
-        block_var = tvm.tir.Var(var_name, dtype="int32")
         dom = tvm.runtime.convert(dom)
         if isinstance(dom, PrimExpr):
             dom = tvm.ir.Range(dom)
@@ -497,6 +496,7 @@ def axis(
                 f"Block axis domain expected PrimExpr or Range, but got {type(dom)}",
                 self.node.span,
             )
+        block_var = tvm.tir.Var(var_name, dtype=dom.extent.dtype)
         value = tvm.runtime.convert(value)
         if not isinstance(value, PrimExpr):
             self.context.report_error(

python/tvm/script/tir/utils.py

@@ -16,11 +16,12 @@
 # under the License.
 """Helper functions in TVM Script Parser"""
 
-from typing import List, Optional, Union
+from typing import List, Optional
 
 from tvm.arith import Analyzer
 from tvm.ir import Range
 from tvm.tir import PrimExpr, BufferRegion
+from tvm.tir.expr import IntImm
 from .node import BufferSlice
 
 
@@ -44,8 +45,8 @@ def buffer_slice_to_region(
     """
     region: List[Range] = []
     for s in buffer_slice.slices:
-        start: Union[PrimExpr, int] = s.start
-        extent: Union[PrimExpr, int] = 1 if s.stop is None else s.stop - s.start
+        start = s.start if isinstance(s.start, PrimExpr) else IntImm("int32", s.start)
+        extent = IntImm(start.dtype, 1) if s.stop is None else s.stop - s.start
         if not analyzer:
             analyzer = Analyzer()
         if isinstance(extent, PrimExpr):

src/tir/ir/expr.cc

@@ -147,6 +147,10 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
 IterVar::IterVar(Range dom, Var var, IterVarType t, String thread_tag, Span span) {
   ObjectPtr<IterVarNode> n = make_object<IterVarNode>();
   if (dom.defined() && dom->extent.defined()) {
+    CHECK(dom->extent.dtype().is_int())
+        << "The dtype of the domain of an IterVar must be an integer type. However, the domain's "
+           "dtype is "
+        << dom->extent.dtype();
     CHECK_EQ(dom->extent.dtype(), var.dtype())
         << "The dtype of the extent of an IterVar (" << dom->extent.dtype()
         << ") must match its associated Var's dtype (" << var.dtype() << ")";

tests/python/unittest/test_tvmscript_error_report.py

@@ -21,8 +21,6 @@
 from tvm import tir
 from tvm.testing import check_error
 from tvm.script import tir as T
-from tvm.ir.diagnostics import override_renderer
-import inspect
 
 
 def buffer_bind_missing_args(a: T.handle) -> None:
@@ -629,5 +627,14 @@ def test_floor_dtype():
     check_error(floor_dtype, 3)
 
 
+def non_integer_typed_block_iter():
+    with T.block():
+        i = T.axis.S(0.1, 0.1)  # error IterVar requires an integer dtype
+
+
+def test_non_integer_typed_block_iter():
+    check_error(non_integer_typed_block_iter, 3)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__] + sys.argv[1:]))

tests/python/unittest/test_tvmscript_roundtrip.py

@@ -3205,6 +3205,24 @@ def segment_sum(
     return segment_sum
 
 
+def int64_support():
+    @T.prim_func
+    def elementwise_shape_int64(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (T.int64(128), T.int64(128)), dtype="float32")
+        B = T.alloc_buffer((T.int64(128), T.int64(128)), dtype="float32")
+        C = T.match_buffer(c, (T.int64(128), T.int64(128)), dtype="float32")
+        for i, j in T.grid(128, 128):
+            with T.block("B"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                B[vi, vj] = A[vi, vj] * 2.0
+        for i, j in T.grid(T.int64(128), T.int64(128)):
+            with T.block("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                C[vi, vj] = B[vi, vj] + 1.0
+
+    return elementwise_shape_int64
+
+
 ir_generator = tvm.testing.parameter(
     opt_gemm_normalize,
     opt_gemm_lower,
@@ -3237,6 +3255,7 @@ def segment_sum(
     func_T_ptr_allocate,
     llvm_intrin_call,
     parse_bufferslice_as_range_bound,
+    int64_support,
 )