[TIR][TVMScript] specialize

[Hzfengsy]

This PR enables meta_programming for TVM Functions (usually used for TVMScript), which is useful to define only one function(TVMScript) but used in different shapes.

Example:
Mate fucntion

@tvm.script.tir
def mem_copy(a: ty.handle, b: ty.handle, m: ty.int32, n: ty.int32) -> None:
    A = tir.match_buffer(a, (m, n), "float32")
    B = tir.match_buffer(b, (m, n), "float32")

    with tir.block([m, n], "") as [vi, vj]:
        B[vi, vj] = A[vi, vj]

Instruction:

a, _, m, n = mem_copy.params
func = mem_copy.specialize({a: tir.decl_buffer((16, 16))})
# or
func = mem_copy.specialize({n: 16, m: 16})

Specialized function:

@tvm.script.tir
def mem_copy_16_16(a: ty.handle, b: ty.handle) -> None:
    A = tir.match_buffer(a, (16, 16), "float32")
    B = tir.match_buffer(b, (16, 16), "float32")

    with tir.block([16, 16], "") as [vi, vj]:
        B[vi, vj] = A[vi, vj]

cc @tqchen @junrushao1994 @MasterJH5574 @comaniac

[jcf94]

LGTM. Thanks!

[tqchen]

cc @vinx13 @junrushao1994 please help to take a look

[comaniac]

The logic LGTM. One miner point: is it possible to work on param_map directly instead of one-by-one? Specifically, can we build var mapping in Specialize with several parameters so that we can run PrimFuncSpecializer::Specialize just once and make the time complexity decouple to the parameter number?

[Hzfengsy]

Thanks, @comaniac. Just have fixed it

[MasterJH5574]

I saw this TODO item today. Is it intended to leave the TODO here? Or it was because we forgot to write the unittest?