[DLight][CUDA] Fall back from unsupported GEMV spatial tiles - #20330
Open
Nanmur wants to merge 1 commit into
Open
[DLight][CUDA] Fall back from unsupported GEMV spatial tiles#20330Nanmur wants to merge 1 commit into
Nanmur wants to merge 1 commit into
Conversation
The GEMV rule assumed that each outer spatial tile extent was an IntImm equal to one. Some valid convolution shapes produce a FloorDiv extent greater than one, causing an AttributeError or assertion failure. Use the arithmetic analyzer to prove the unit-extent requirement and decline the GEMV rule when it is not satisfied, allowing the GPU fallback rule to schedule the PrimFunc. Add a regression test for the reported convolution shape.
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
1 participant
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.
The GPU GEMV schedule assumes that the outer spatial loop produced by its thread-tile split has extent one. It checks this by reading
.extent.value, which crashes when the extent is aFloorDiv; for the convolution shapes in #20047, the extent is also genuinely larger than one, so forcing this GEMV schedule is not valid.This change uses
arith.Analyzer.can_prove_equalat the three unit-extent guards. When the selected GEMV configuration cannot prove that the outer extent is one, the rule returnsNone, allowingApplyDefaultScheduleto continue to the GPU Fallback rule.The regression test uses the lowered TIR for the reported
1x1x3x10input and1x1x1x2kernel. It verifies that GEMV declines the unsupported configuration and Fallback produces a scheduled PrimFunc instead of raisingAttributeErrororAssertionError.Fixes #20047
Tests:
python -m pytest tests/python/s_tir/dlight/test_gpu_gemv.py -q(15 passed)python -m ruff check python/tvm/s_tir/dlight/gpu/gemv.py tests/python/s_tir/dlight/test_gpu_gemv.pypython -m ruff format --check python/tvm/s_tir/dlight/gpu/gemv.py tests/python/s_tir/dlight/test_gpu_gemv.py