This is an automated email from the ASF dual-hosted git repository. junrushao pushed a commit to branch unity-staging in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 959b7e5e09315a3e8cb9679e6fcf4d896d598a97 Author: Ruihang Lai <[email protected]> AuthorDate: Mon Jul 17 17:39:54 2023 -0700 [Unity][Dlight] Fix DecodeGeMV rule for spatial-inner with grouping (#15340) This PR fixes a bug of DecodeGeMV dlight rule when the innermost tensor dimension is spatial with `unroll_factor` (for example, the grouping used in group quantization). Prior to this PR, a reduction loop that is bound to threadIdx was reordered to reside outside a split spatial loop, which prevents the TIR LowerCrossThreadReduction pass to successfully apply due to some safety-guard requirement. This PR fixes this issue by not reordering the split spatial loop after the reduction loop, so that the pass can be applied. Note that we can do this as the order of thread-binding loops does not matter. --- python/tvm/dlight/gpu/decode_gemv.py | 3 +-- tests/python/dlight/test_gpu_decode_gemv.py | 4 ++-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/python/tvm/dlight/gpu/decode_gemv.py b/python/tvm/dlight/gpu/decode_gemv.py index 1aa5d68fc5..5566f3248c 100644 --- a/python/tvm/dlight/gpu/decode_gemv.py +++ b/python/tvm/dlight/gpu/decode_gemv.py @@ -220,8 +220,7 @@ class DecodeGEMV(ScheduleRule): s = sch.fuse(*s) sch.reorder(s, r) if unroll_spatial_factor: - s, inner = sch.split(s, factors=[None, unroll_spatial_factor]) - sch.reorder(s, r, inner) + s, _ = sch.split(s, factors=[None, unroll_spatial_factor]) sch.bind(s, "threadIdx.x") sch.bind(r, "threadIdx.y") # Schedule epilogue diff --git a/tests/python/dlight/test_gpu_decode_gemv.py b/tests/python/dlight/test_gpu_decode_gemv.py index 971f5f4d09..d037ffa3ee 100644 --- a/tests/python/dlight/test_gpu_decode_gemv.py +++ b/tests/python/dlight/test_gpu_decode_gemv.py @@ -259,8 +259,8 @@ def test_decode_gemv_4(): vk_fused_0 = T.axis.reduce(256, k_fused_0) C_rf_local[vk_fused_1, 0, 0, v_i2] = C_rf_local[vk_fused_1, 0, 0, v_i2] + V[0, 0, vk_fused_0 * 16 + vk_fused_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(W[vk_fused_0 * 16 + vk_fused_1, v_i2 // 8], T.Cast("uint32", v_i2 % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[vk_fused_0 * 16 + vk_fused_1, v_i2 // 32]) for ax1_ax2_ax3_fused_0 in T.thread_binding(16, thread="threadIdx.x"): - for ax0_fused in T.thread_binding(16, thread="threadIdx.y"): - for ax1_ax2_ax3_fused_1 in range(8): + for ax1_ax2_ax3_fused_1 in range(8): + for ax0_fused in T.thread_binding(16, thread="threadIdx.y"): with T.block("matmul"): vk_fused_1 = T.axis.reduce(16, ax0_fused) v_i2 = T.axis.spatial(4096, i2_0_i0_i1_fused_0 * 128 + ax1_ax2_ax3_fused_0 * 8 + ax1_ax2_ax3_fused_1)
