[GitHub] [tvm] yzh119 opened a new pull request, #14626: [TensorIR][Bugfix] `reindex_cache_write` do not mutate init statement

["yzh119 (via GitHub)" <gi...@apache.org>]

yzh119 opened a new pull request, #14626:
URL: https://github.com/apache/tvm/pull/14626

   # The Bug
   When applying `reindex_cache_write` to the write buffer of a reduction block, the init statement would not be mutated accordingly:
   
   ```python
   # 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.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] tvm-bot commented on pull request #14626: [TensorIR][Bugfix] `reindex_cache_write` do not mutate init statement

["tvm-bot (via GitHub)" <gi...@apache.org>]

tvm-bot commented on PR #14626:
URL: https://github.com/apache/tvm/pull/14626#issuecomment-1508311500

   <!---bot-comment-->
   
   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](https://github.com/apache/incubator-tvm/blob/master/CONTRIBUTORS.md#reviewers) by @-ing them in a comment.
   
   <!--bot-comment-ccs-start-->
    * No users to tag found in teams: `tensorir`, `bugfix` <sub>See [#10317](https://github.com/apache/tvm/issues/10317) for details</sub><!--bot-comment-ccs-end-->
   
   <sub>Generated by [tvm-bot](https://github.com/apache/tvm/blob/main/ci/README.md#github-actions)</sub>


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] Hzfengsy merged pull request #14626: [TensorIR][Bugfix] `reindex_cache_write` do not mutate init statement

["Hzfengsy (via GitHub)" <gi...@apache.org>]

Hzfengsy merged PR #14626:
URL: https://github.com/apache/tvm/pull/14626


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org