[TIR] Enhance and fix tensorize schedule for some case

[LeiWang1999]

To optimize i4_to_f16 decoding, we can use some advanced hardware instructions to do fast type conversion to alleviate the cost of decoding, we can do that by tensorize in tvm.

To tensorize decoding, this pr extends the call component of ir_comparator, which is necessary because the decode block comprises call expressions.

Moreover, currently comparator do simplification on the lhs expr, however, the tensor intrin descs are not simplified, which will be inconsistent and will fail at comparation,
see this pr: #14108.

For example, we provide a test case for this situation:

def test_tensorize_arith_simplification():
    # fmt: off
    @T.prim_func
    def decode_i4s_to_int32_to_f16():
        B_decode_local = T.alloc_buffer((16384, 16384), "float16", scope="local")
        B_local = T.alloc_buffer((16384, 2048), "int32", scope="local")
        for ax0_0 in T.thread_binding(8192, thread="blockIdx.x"):
            for ax0_1 in T.thread_binding(2, thread="threadIdx.y"):
                for ax1_0 in range(32):
                    for ax1_1 in T.thread_binding(64, thread="threadIdx.x"):
                        for ax0, ax1 in T.grid(1, 8):
                            with T.block("B_decode_local"):
                                v0 = T.axis.spatial(16384, ax0_0 * 2 + ax0_1 + ax0)
                                v1 = T.axis.spatial(16384, ax1_0 * 512 + ax1_1 * 8 + ax1)
                                T.reads(B_local[v0, v1 // 8])
                                T.writes(B_decode_local[v0, v1])
                                B_decode_local[v0, v1] = T.Cast("float16", T.shift_right(T.shift_left(T.bitwise_and(T.shift_right(B_local[v0, v1 // 8], v1 % 8 * 4), 15), 28), 28))

The desc should be simplified from [v1 // 8] and [v1 % 8] to [0], [v1] to match the simplified lhs expr.

To do simplification for tensor intrin's desc, we warp and reuse tir::transform::simplify to support simplification for single stmt.

[LeiWang1999]

cc @Hzfengsy and @Lunderberg , looks like pr #13299 provides a stmt_simplify declaration but do not provide an implementation.

[tqchen]

cc @vinx13

[Lunderberg]

Instead of simplifying the body of the PrimFunc, can we instead simplify the entire PrimFunc? That way, dynamic expressions that are used in shapes are exposed to the analyzer as non-negative. (e.g. Using buffer of shape [n,m] implies that n >= 0 && m >= 0.)

[LeiWang1999]

as u mentioned in #13299 (comment) , perhaps its better to simplify in prim_func level, I chose to implement a stmt simplifier because it may be more useful. The rationale is that stmt is more fine-grained.

Moreover, in the context of tensor desc in tensorize schedule, prim_func typically encompasses a single block without dynamic symbolic. I think for this issue a stmt simplifier is enough.
But we can implement a prim_func one as well, should we keep both stmt and primfunc simplifier or just maintain only one of them?

[Lunderberg]

as u mentioned in #13299 (comment) , perhaps its better to simplify in prim_func level, I chose to implement a stmt simplifier because it may be more useful. The rationale is that stmt is more fine-grained.

Good point. Thinking on it again in the morning, I think we should avoid having the simplify function for tir::Stmt altogether, because it is more fine-grained. That is, its existence would encourage simplifications to be performed for specific statements, even though those statements might not be the outer-most.

But we can implement a prim_func one as well, should we keep both stmt and primfunc simplifier or just maintain only one of them?

I think having a simplifier for a PrimFunc would be better, because it encourages developers to simplify with the full context of a statement. The functionality already exists here, and would just need a wrapper function to expose StmtSimplifier::Apply.

[LeiWang1999]

@Lunderberg hi, I think this pr is ready for review.

[tqchen]

@vinx13 @spectrometerHBH do you mind take a look at this PR?