[Fix][DLight] Handle rank-one GEMV cache loads

python/tvm/s_tir/dlight/gpu/gemv.py

@@ -164,13 +164,14 @@ def apply(
                 and shared_mem_usage.value <= max_smem
             )
 
-            # vectorize load A
-            # (TODO) this is now actually problematic since the number of loops is dependent on the
-            # number of dimensions of A_q
             Aq_local = sch.cache_read(rf, read_buffer_index=1, storage_scope="local")
+            num_cache_loops = len(sch.get_loops(block=Aq_local))
             sch.compute_at(Aq_local, r, preserve_unit_loops=True)
-            s_local, r_local = sch.get_loops(block=Aq_local)[-2:]
-            fused_load = sch.fuse(s_local, r_local)
+            cache_loops = sch.get_loops(block=Aq_local)[-num_cache_loops:]
+            if len(cache_loops) == 1:
+                fused_load = cache_loops[0]
+            else:
+                fused_load = sch.fuse(*cache_loops[-2:])
             aq_vec_len = max(1, VEC_LOAD // get_bytes(sch.get(Aq_local).reads[0].buffer.dtype))
             fused_load, vec_load = sch.split(
                 fused_load, factors=[None, aq_vec_len], preserve_unit_iters=True

tests/python/s_tir/dlight/test_gpu_gemv.py

@@ -1114,6 +1114,34 @@ def before(
     assert mod["main"].attrs["tirx.is_scheduled"] == 1
 
 
+def test_gemv_rank_one_vector_input():
+    @T.prim_func(private=True, s_tir=True)
+    def before(
+        matrix: T.Buffer((2, 2), "float32"),
+        vector: T.Buffer((2,), "float32"),
+        output: T.Buffer((2,), "float32"),
+    ):
+        T.func_attr({"tirx.noalias": True})
+        for i, k in T.grid(2, 2):
+            with T.sblock("gemv"):
+                vi, vk = T.axis.remap("SR", [i, k])
+                T.reads(matrix[vi, vk], vector[vk])
+                T.writes(output[vi])
+                with T.init():
+                    output[vi] = T.float32(0)
+                output[vi] += matrix[vi, vk] * vector[vk]
+
+    mod = tvm.IRModule({"main": before})
+    with Target("nvidia/geforce-rtx-3090-ti"):
+        mod = dl.ApplyDefaultSchedule(dl.gpu.GEMV())(mod)
+
+    assert mod["main"].attrs["tirx.is_scheduled"] == 1
+    sch = tvm.s_tir.Schedule(mod)
+    vector_local = sch.get_sblock("vector_local")
+    vector_load_loop = sch.get(sch.get_loops(vector_local)[-1])
+    assert vector_load_loop.kind == tvm.tirx.ForKind.VECTORIZED
+
+
 def test_gemv_broadcast_epilogue():
     # fmt: off
     @T.prim_func(private=True, s_tir=True)