This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 5e5b6602e4 [REFACTOR][DLight] Remove obsolete vectorization pragmas
(#20561)
5e5b6602e4 is described below
commit 5e5b6602e4b1fd9905593f9c755468069735daba
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 6 12:19:00 2026 -0400
[REFACTOR][DLight] Remove obsolete vectorization pragmas (#20561)
---
python/tvm/s_tir/dlight/gpu/gemv.py | 3 ---
python/tvm/s_tir/dlight/gpu/low_batch_gemv.py | 3 ---
tests/python/s_tir/dlight/test_gpu_gemv.py | 8 ++++----
3 files changed, 4 insertions(+), 10 deletions(-)
diff --git a/python/tvm/s_tir/dlight/gpu/gemv.py
b/python/tvm/s_tir/dlight/gpu/gemv.py
index 50c0635ed6..60c8944a90 100644
--- a/python/tvm/s_tir/dlight/gpu/gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/gemv.py
@@ -274,9 +274,6 @@ class GEMV(GPUScheduleRule):
ann_key="unroll_explicit",
ann_val=unroll_factor,
)
- sch.annotate(
- block_or_loop=sch.get_loops(V_shared)[-4],
ann_key="pragma_vectorize", ann_val=1
- )
# Schedule epilogue
if epilogue_info is not None:
diff --git a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
index 0c4c29a27a..f41a7227c0 100644
--- a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
@@ -488,9 +488,6 @@ class LowBatchGEMV(GPUScheduleRule):
ann_key="unroll_explicit",
ann_val=unroll_factor,
)
- sch.annotate(
- block_or_loop=sch.get_loops(V_shared)[-4],
ann_key="pragma_vectorize", ann_val=1
- )
epilogue = sch.get_consumers(gemv)
# Schedule epilogue
diff --git a/tests/python/s_tir/dlight/test_gpu_gemv.py
b/tests/python/s_tir/dlight/test_gpu_gemv.py
index 3a6c415e9d..d204ccf331 100644
--- a/tests/python/s_tir/dlight/test_gpu_gemv.py
+++ b/tests/python/s_tir/dlight/test_gpu_gemv.py
@@ -114,7 +114,7 @@ def test_gemv_basic():
for ax0_fused_ax1_fused_fused_1 in T.thread_binding(1,
thread="threadIdx.y"):
for ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 in
T.thread_binding(64, thread="threadIdx.x"):
for ax0, ax1, ax2 in T.grid(1, 1, 1):
- for ax3_0 in T.serial(1,
annotations={"unroll_explicit": 256, "pragma_vectorize": 1}):
+ for ax3_0 in T.serial(1,
annotations={"unroll_explicit": 256}):
for ax3_1 in T.thread_binding(1,
thread="threadIdx.y"):
for ax3_2 in T.thread_binding(64,
thread="threadIdx.x"):
for ax3_3 in T.vectorized(2):
@@ -334,7 +334,7 @@ def test_decode_gemv1():
for u_fused_ax0_fused_fused_1 in T.thread_binding(16,
thread="threadIdx.y"):
for
ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 in
T.thread_binding(32, thread="threadIdx.x"):
for ax0, ax1 in T.grid(1, 1):
- for ax2_0 in T.serial(1,
annotations={"unroll_explicit": 256, "pragma_vectorize": 1}):
+ for ax2_0 in T.serial(1,
annotations={"unroll_explicit": 256}):
for ax2_1 in T.thread_binding(16,
thread="threadIdx.y"):
for ax2_2 in T.thread_binding(32,
thread="threadIdx.x"):
for ax2_3 in T.vectorized(8):
@@ -450,7 +450,7 @@ def test_decode_gemv2():
for u_fused_ax0_fused_fused_1 in T.thread_binding(16,
thread="threadIdx.y"):
for
ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 in
T.thread_binding(32, thread="threadIdx.x"):
for ax0, ax1 in T.grid(1, 1):
- for ax2_0 in T.serial(1,
annotations={"unroll_explicit": 256, "pragma_vectorize": 1}):
+ for ax2_0 in T.serial(1,
annotations={"unroll_explicit": 256}):
for ax2_1 in T.thread_binding(16,
thread="threadIdx.y"):
for ax2_2 in T.thread_binding(32,
thread="threadIdx.x"):
for ax2_3 in T.vectorized(8):
@@ -573,7 +573,7 @@ def test_decode_gemv3():
for u_fused_ax0_fused_fused_1 in T.thread_binding(T.int64(16),
thread="threadIdx.y"):
for
ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 in
T.thread_binding(T.int64(32), thread="threadIdx.x"):
for ax0, ax1 in T.grid(T.int64(1), T.int64(1)):
- for ax2_0 in T.serial(T.int64(22),
annotations={"unroll_explicit": 256, "pragma_vectorize": 1}):
+ for ax2_0 in T.serial(T.int64(22),
annotations={"unroll_explicit": 256}):
for ax2_1 in T.thread_binding(T.int64(16),
thread="threadIdx.y"):
for ax2_2 in T.thread_binding(T.int64(32),
thread="threadIdx.x"):
for ax2_3 in T.vectorized(T.int64(1)):
