Repository navigation
[Fix][DLight] Handle rank-one GEMV cache loads - #20158
Merged
Merged
Conversation
The DLight GPU GEMV rule assumed that the final two loops around a local cache read were both owned by the cache stage. Rank-one vector inputs create only one cache loop, so the rule included the shared placement loop and attempted to fuse an imperfect nest. Record the cache loop count before compute_at and recover that innermost suffix after placement. Use the sole cache loop directly for rank-one inputs while preserving the existing final-two-loop fusion for higher-rank inputs. Add a focused regression that verifies rank-one GEMV scheduling and vectorized local loading.
SamJSui
marked this pull request as ready for review
August 18, 2026 14:23
wodylm
pushed a commit
to wodylm/mcTVM
that referenced
this pull request
Sep 15, 2026
DLight's GPU GEMV rule assumes that the final two loops around a local cache read are both owned by the cache stage. Rank-one vector inputs create only one cache loop, causing the rule to include the shared placement loop and attempt to fuse an imperfect loop nest. This change records the cache loop count before `compute_at` and recovers that innermost suffix afterward. A single cache loop is used directly, while the existing final-two-loop fusion is preserved for higher-rank inputs. A focused regression verifies that rank-one GEMV scheduling completes and that the local vector load remains vectorized. Fixes apache#20118 Testing: - `python -m pytest tests/python/s_tir/dlight/test_gpu_gemv.py -q` — 14 passed - `python -m pytest tests/python/relax/test_pipeline.py -q` — 24 passed - Exact issue reproducer — passed - RTX 4070 Ti SUPER (SM89): FP32/FP16 widths 2, 3, 32, 33, and 128 — NumPy parity - `pre-commit run --all-files` — passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
DLight's GPU GEMV rule assumes that the final two loops around a local cache
read are both owned by the cache stage. Rank-one vector inputs create only one
cache loop, causing the rule to include the shared placement loop and attempt
to fuse an imperfect loop nest.
This change records the cache loop count before
compute_atand recovers thatinnermost suffix afterward. A single cache loop is used directly, while the
existing final-two-loop fusion is preserved for higher-rank inputs.
A focused regression verifies that rank-one GEMV scheduling completes and that
the local vector load remains vectorized.
Fixes #20118
Testing:
python -m pytest tests/python/s_tir/dlight/test_gpu_gemv.py -q— 14 passedpython -m pytest tests/python/relax/test_pipeline.py -q— 24 passedpre-commit run --all-files— passed