[DLight][CUDA] Fix undefined TX in GEMV broadcast epilogue

[Nanmur]

This PR fixes an undefined TX reference in the DLight GPU GEMV inner-reduction schedule.

The broadcast epilogue path splits fused epilogue loops and binds the inner loop to threadIdx.x, but it referenced TX, which is not defined in the surrounding scope or passed into the helper. This PR uses the existing TR tile factor, which is already used for the threadIdx.x direction in this schedule.

A regression test is added to cover the GEMV broadcast epilogue path.

Fixes #19969

Tests:

• python -m pytest tests/python/s_tir/dlight/test_gpu_gemv.py -q
• python -m pytest tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py -q
• python -m pytest tests/python/s_tir/dlight/test_gpu_reduction.py -q

Note: The branch is based on apache/tvm:main. Local main-branch test execution on this Windows machine could not be completed because the available compiled TVM library is from a v0.25 build and does not match the latest Python sources; the same tests passed in the matching local v0.25 environment before rebasing the patch to main.

[gemini-code-assist]

Code Review

This pull request updates the GPU GEMV schedule in python/tvm/s_tir/dlight/gpu/gemv.py by changing the split factor for the epilogue block from TX to TR. It also adds a new unit test test_gemv_broadcast_epilogue in tests/python/s_tir/dlight/test_gpu_gemv.py to validate this scheduling behavior. There are no review comments, so I have no feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[tlopex]

Thanks for fixing this! The same undefined TX reference still exists in the LowBatchGEMV broadcast-epilogue path:

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

Lines 491 to 496 in 9837463

	if is_broadcast_epilogue(sch, block, epilogue):
	sch.reverse_compute_at(epilogue, bx)
	sch.set_scope(block, 0, "shared")
	_, _, _, *s = sch.get_loops(epilogue) # pylint: disable=invalid-name
	_, tx = sch.split(sch.fuse(*s), factors=[None, TX])
	sch.bind(tx, TAG_S)

TX is not defined anywhere in that file either, so this branch will fail with the same runtime error when reached. Since the resulting loop is bound to TAG_S, the corresponding existing tile factor appears to be TS:

_, ts = sch.split(sch.fuse(*s), factors=[None, TS])
sch.bind(ts, TAG_S)

The existing LowBatchGEMV tests do not trigger this broadcast-epilogue branch. Could we fix this occurrence as well and add a regression test for it?

[Nanmur]

Thanks for catching this. I updated the PR to fix the LowBatchGEMV broadcast-epilogue path as well, replacing the undefined TX with TS for the loop bound to TAG_S, and added a regression test that covers this branch.

[Nanmur]

@yongwww The current wasm/pr-head build checked out the latest commit e3738a7 after the brief close/reopen event. Its prepare stage completed successfully, and the build reached the Emscripten (emcc) step before the Jenkins EC2 worker went offline with ChannelClosedException; the job was aborted without a compilation or test failure. Could you please re-run the WASM check when convenient?

The LowBatchGEMV follow-up and regression test requested by @tlopex are included in e3738a7.