[TensorIR][Bugfix] reindex_cache_write do not mutate init statement

[yzh119]

The Bug

When applying reindex_cache_write to the write buffer of a reduction block, the init statement would not be mutated accordingly:

# original program
@T.prim_func
def reduce(A: T.Buffer((128, 128, 128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
    B = T.alloc_buffer((128, 128, 128), dtype="float32")
    for i, j, k in T.grid(128, 128, 128):
        for l in range(128):
            with T.block("B"):
                vi, vj, vk, vl = T.axis.remap("SSSR", [i, j, k, l])
                with T.init():
                    B[vi, vj, vk] = T.float32(0)
                B[vi, vj, vk] = B[vi, vj, vk] + A[vi, vj, vk, vl]
        with T.block("C"):
            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
            with T.init():
                C[vi, vj] = T.float32(0)
            C[vi, vj] = C[vi, vj] + B[vi, vj, vk]

# schedule
sch = tir.Schedule(reduce, debug_mask="all")
sch.reindex_cache_write("B", 0, "shared", lambda i, j, k, l: (j, i, k))

# after schedule
@T.prim_func
def reduce_after_reindex_cache_write(
    A: T.Buffer((128, 128, 128, 128), "float32"), C: T.Buffer((128, 128), "float32")
):
    B = T.alloc_buffer((128, 128, 128))
    B_shared = T.alloc_buffer((128, 128, 128), scope="shared")
    for i, j, k in T.grid(128, 128, 128):
        for l in range(128):
            with T.block("B"):
                vi, vj, vk, vl = T.axis.remap("SSSR", [i, j, k, l])
                T.reads(A[vi, vj, vk, vl])
                T.writes(B_shared[vj, vi, vk])
                with T.init():
                    B[vj, vi, vk] = T.float32(0)
                B_shared[vj, vi, vk] = B_shared[vj, vi, vk] + A[vi, vj, vk, vl]
        with T.block("B_shared"):
            vi, vj, vk = T.axis.remap("SSS", [i, j, k])
            T.reads(B_shared[vj, vi, vk])
            T.writes(B[vi, vj, vk])
            B[vi, vj, vk] = B_shared[vj, vi, vk]
        with T.block("C"):
            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
            T.reads(B[vi, vj, vk])
            T.writes(C[vi, vj])
            with T.init():
                C[vi, vj] = T.float32(0)
            C[vi, vj] = C[vi, vj] + B[vi, vj, vk]

The init statement inside block "B" should be transformed to B[vj, vi, vk] = T.float32(0)

The Fix

In our previous implementation, we mistakenly specify the consumer block to be the block itself, which is not necessary and would cause the later ReindexCacheWriteRewriter to skip rewriting the init statement.

cc @Hzfengsy @vinx13 @junrushao

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• No users to tag found in teams: tensorir, bugfix _{See #10317 for details}

_{Generated by tvm-bot}