[Bug] DLight GPU GEMV fails with undefined TX in broadcast epilogue path
Problem
Building a Relax module for CUDA can fail inside the DLight GPU GEMV schedule with:
RuntimeError: name 'TX' is not defined
The failure comes from:
# python/tvm/s_tir/dlight/gpu/gemv.py
_, tx = sch.split(sch.fuse(*s), factors=[None, TX])
TX does not appear to be defined in the surrounding scope or passed into the inner apply(...) function. The same function already has TS and TR schedule parameters, and TR is used as the threadIdx.x tile size in nearby scheduling logic.
Source Location
Current main branch appears to still contain this reference:
The same issue is present in v0.25.0:
Environment
TVM version: v0.25.0 based local source tree
Target: cuda
Frontend: Relax
CUDA device visible: tvm.cuda(0).exist == True
OS: Windows
Python: 3.11
Reproduction Context
I hit this while compiling a Relax model for CUDA:
ex = relax.build(mod, target=tvm.target.Target("cuda"))The same Relax module builds and runs correctly on LLVM. The failure happens before CUDA VM execution, during the Relax CUDA build pipeline when DLight applies the GPU GEMV schedule.
The traceback ends at:
File "python/tvm/s_tir/dlight/gpu/gemv.py", line 291, in apply
_, tx = sch.split(sch.fuse(*s), factors=[None, TX])
RuntimeError: name 'TX' is not defined
Expected Behavior
relax.build(mod, target="cuda") should not fail because of an undefined Python variable in the DLight schedule rule.
Actual Behavior
The CUDA build fails during DLight GPU GEMV scheduling with TX undefined.
Suspected Cause
In GEMV.sch_inner_reduction, the broadcast epilogue branch uses TX:
_, tx = sch.split(sch.fuse(*s), factors=[None, TX])
sch.bind(tx, "threadIdx.x")
However, TX is not defined. The schedule configuration defines TS and TR, and TR is bound to TAG_R, which is "threadIdx.x" for CUDA in this path.
A likely minimal fix may be to replace TX with TR:
_, tx = sch.split(sch.fuse(*s), factors=[None, TR])
sch.bind(tx, "threadIdx.x")
However, this should be validated with a regression test for the broadcast epilogue GEMV path.
Suggested Fix
Add a regression test that triggers the DLight GPU GEMV inner-reduction schedule with a broadcast epilogue, then update the undefined TX reference to the correct existing tile variable.
The patch should be separate from any frontend-specific changes, because this is in the DLight GPU schedule layer.
[Bug] DLight GPU GEMV fails with undefined
TXin broadcast epilogue pathProblem
Building a Relax module for CUDA can fail inside the DLight GPU GEMV schedule with:
The failure comes from:
TXdoes not appear to be defined in the surrounding scope or passed into the innerapply(...)function. The same function already hasTSandTRschedule parameters, andTRis used as thethreadIdx.xtile size in nearby scheduling logic.Source Location
Current main branch appears to still contain this reference:
The same issue is present in v0.25.0:
Environment
Reproduction Context
I hit this while compiling a Relax model for CUDA:
The same Relax module builds and runs correctly on LLVM. The failure happens before CUDA VM execution, during the Relax CUDA build pipeline when DLight applies the GPU GEMV schedule.
The traceback ends at:
Expected Behavior
relax.build(mod, target="cuda")should not fail because of an undefined Python variable in the DLight schedule rule.Actual Behavior
The CUDA build fails during DLight GPU GEMV scheduling with
TXundefined.Suspected Cause
In
GEMV.sch_inner_reduction, the broadcast epilogue branch usesTX:However,
TXis not defined. The schedule configuration definesTSandTR, andTRis bound toTAG_R, which is"threadIdx.x"for CUDA in this path.A likely minimal fix may be to replace
TXwithTR:However, this should be validated with a regression test for the broadcast epilogue GEMV path.
Suggested Fix
Add a regression test that triggers the DLight GPU GEMV inner-reduction schedule with a broadcast epilogue, then update the undefined
TXreference to the correct existing tile variable.The patch should be separate from any frontend-specific changes, because this is in the DLight GPU schedule layer.