[Fix][DLight] Decline GEMV reductions with only batch spatial axes

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

@@ -158,8 +158,10 @@ def normalize(
         else:
             s_loops.append(loop)
 
-    assert s_loops
-    assert r_loops
+    # Broadcast operands can make a general reduction look like GEMV.
+    # GEMV needs a non-batch spatial axis as well as a reduction axis.
+    if not s_loops or not r_loops:
+        return None
     if not c_loops:
         c_loops = [sch.add_unit_loop(block_info.block_rv)]
     if not batch_loops:

tests/python/s_tir/dlight/test_gpu_gemv.py

@@ -52,6 +52,41 @@ def before(
     tvm_tirx.ir.assert_structural_equal(scheduled["main"], before)
 
 
+def test_gemv_rejects_all_batch_spatial_axes():
+    frames = T.dynamic("frames", "int64")
+
+    @Ts.function(private=True)
+    def before(
+        data: T.Tensor((T.int64(768), frames), "float32"),
+        mean: T.Tensor((T.int64(768),), "float32"),
+        weights: T.Tensor((T.int64(768), frames), "float32"),
+        result: T.Tensor((T.int64(768),), "float32"),
+    ):
+        for c, t in T.grid(T.int64(768), frames):
+            with Ts.sblock("weighted_variance"):
+                vc, vt = Ts.axis.remap("SR", [c, t])
+                with Ts.init():
+                    result[vc] = T.float32(0)
+                result[vc] += weights[vc, vt] * (
+                    (data[vc, vt] - mean[vc]) * (data[vc, vt] - mean[vc])
+                )
+
+    mod = tvm_tirx.IRModule({"main": before})
+    target = Target(
+        {
+            "kind": "webgpu",
+            "keys": ["webgpu", "gpu"],
+            "max_num_threads": 256,
+            "thread_warp_size": 1,
+            "max_shared_memory_per_block": 32768,
+        }
+    )
+    rejected = dl.ApplyDefaultSchedule(dl.gpu.GEMV(), target=target)(mod)
+    tvm_tirx.ir.assert_structural_equal(rejected["main"], before)
+    scheduled = dl.ApplyDefaultSchedule(dl.gpu.GEMV(), dl.gpu.Reduction(), target=target)(mod)
+    assert scheduled["main"].attrs["tirx.is_scheduled"]
+
+
 def test_gemv_basic():
     # fmt: off
     n = T.dynamic("n", "int32")