Nanmur opened a new issue, #19969: URL: https://github.com/apache/tvm/issues/19969
# [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: ```text RuntimeError: name 'TX' is not defined ``` The failure comes from: ```python # 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: - https://raw.githubusercontent.com/apache/tvm/main/python/tvm/s_tir/dlight/gpu/gemv.py The same issue is present in v0.25.0: - https://raw.githubusercontent.com/apache/tvm/v0.25.0/python/tvm/s_tir/dlight/gpu/gemv.py ## Environment ```text 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: ```python 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: ```text 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`: ```python _, 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`: ```python _, 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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
