[WIP] Type hinting for TIR scripts

[yzh119]

A prototype of #420 .

We can write a TIR gemm as:

import tvm.tir.keywords as T

@tvm.script.tir('T')
def matmul(a: ty.handle, b: ty.handle, c: ty.handle) -> None:
    C = T.match_buffer(c, (2048, 2048), "float32")
    A = T.match_buffer(a, (2048, 2048), "float32")
    B = T.match_buffer(b, (2048, 2048), "float32")

    with T.block([2048, 2048, T.reduce_axis(0, 2048)], "C") as [vi, vj, vk]:
        with T.init():
            C[vi, vj] = T.float32(0)
        C[vi, vj] = C[vi, vj] + A[vk, vi] * B[vk, vj]

and IDEs could provide auto completion and type hints.

Note that we provide an argument for tvm.script.tir decorator so that the parser would recognize T as the prefix for tir built-in's. If we do not specify this argument, the parser would assume tir. as the default prefix.

Some choices:

1. loops returns a Iterable[IterVar] instead of For to be compatible with python semantics.
2. block is a ContextManager w/ a __init__ function and a __enter__ function which returns iter vars.
3. env_threads are viewed as iter vars as well.

However, I'm not sure if we have covered enough the coverage of keywords (e.g. should we expose tir.ramp/realize here as they were rarely used in user's original scripts).

Also, I did not directly reuse the types defined under tvm.tir such as PrimExpr/BufferRegion/etc but simplify and redefine them.

TODOs

• Docstrings.
• Let the parser recognize T.
• Use mypy for type check.

@MasterJH5574 @Hzfengsy @spectrometerHBH @jinhongyii @junrushao1994 @vinx13 I'm glad to hear your comments if you have time.

[junrushao]

Let's move the discussion to #471

[shingjan]

@yzh119 There is a follow-up work for this WIP (#9432) and I wonder what is the reason that we need to re-define PrimExpr type with so many private methods here. Thankds!