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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
