diff --git a/python/tvm/s_tir/dlight/analysis/gemv.py b/python/tvm/s_tir/dlight/analysis/gemv.py index 9c8477eeb679..4e1aa29bf0bf 100644 --- a/python/tvm/s_tir/dlight/analysis/gemv.py +++ b/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: diff --git a/tests/python/s_tir/dlight/test_gpu_gemv.py b/tests/python/s_tir/dlight/test_gpu_gemv.py index 25d3aae3b3ef..00561edbe60a 100644 --- a/tests/python/s_tir/dlight/test_gpu_gemv.py +++ b/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")