[GitHub] [tvm] MasterJH5574 commented on pull request #9205: [Arith][TensorIR][Bugfix] Add IterRangeSanityCheck in DetectIterMap

[GitBox <gi...@apache.org>]

MasterJH5574 commented on pull request #9205:
URL: https://github.com/apache/tvm/pull/9205#issuecomment-938541756


   After this fix, in some sparse workloads the real reduction block becomes invalid. Here's an IR of BSR workload:
   ```python
   @T.prim_func
   def bsr(x: T.handle, data: T.handle, indices: T.handle, indptr: T.handle, bsrmm: T.handle, N_blocks: T.int32) -> None:
       M = T.var("int32")
       K = T.var("int32")
       num_blocks = T.var("int32")
       bs_r = T.var("int32")
       bs_c = T.var("int32")
   
       X = T.match_buffer(x, [M, K], "float32")
       W_data = T.match_buffer(data, [num_blocks, bs_r, bs_c], "float32")
       W_indices = T.match_buffer(indices, [num_blocks], "int32")
       W_indptr = T.match_buffer(indptr, [N_blocks + 1], "int32")
       BSRmm = T.match_buffer(bsrmm, [M, N_blocks*bs_r], "float32")
   
       for i, j in T.grid(M, N_blocks*bs_r):
           for block_offset, k in T.grid(W_indptr[T.floordiv(j, bs_r) + 1] - W_indptr[T.floordiv(j, bs_r)], bs_c):
               with T.block([M, N_blocks*bs_r, T.reduce_axis(0, W_indptr[T.floordiv(j, bs_r) + 1] - W_indptr[T.floordiv(j, bs_r)]), T.reduce_axis(0, bs_c)], "sparse_dense") as [m, n, offset, kk]:
                   with T.init():
                       BSRmm[m, n] = T.float32(0)
                   BSRmm[m, n] = BSRmm[m, n] + W_data[offset + W_indptr[T.floordiv(n, bs_r)], T.floormod(n, bs_r), kk]*X[m, bs_c*W_indices[offset+W_indptr[T.floordiv(n, bs_r)]]+kk]
   ```
   
   In this IR, block `"sparse_dense"` is indeed doing a reduction, and users are supposed to be allowed to apply some transformation like `rfactor` to the block. However, due to the fix of this PR, the block is no longer viewed as a reduction block, which disables the possibility of the transformation.
   
   IMO we should still view this block as a reduction block. So here comes the contradiction:
   * if we stick to this fix (prohibiting any iter whose extent depends on other iters), the block cannot be viewed as a reduction block;
   * if we still want to view the block as a reduction block, we should allow the iters whose extents depend on others.
   
   Or another option is to remove the "affine binding" condition for reduction blocks.
   
   @junrushao1994 @spectrometerHBH Want to here your opinions on how to get this issue fixed :-)


-- 
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