Repository navigation
Conversation
DLight's matmul rule can abort default CUDA compilation when a consumer block contains a conditional statement. Consumer-chain inlining assumes that all remaining consumers can be fused, preventing later schedule rules from handling otherwise valid functions. Report unsuccessful consumer-chain inlining and have both Matmul and MetalMatmul discard the partially transformed schedule and return None. Subsequent rules can then schedule the original function. Add regression tests for redundant and nontrivial conditional epilogues, preservation of the original function after a rejected schedule, and CUDA execution through the default Relax compilation pipeline. Validate on H200 with CUDA 12.2 and NVCC: all four new regression tests and 16 existing tests pass. The issue reproducer, adapted to the current APIs, passes exact NumPy comparisons without StmtSimplify. Ruff checks pass. One existing fallback test fails because the apple/m1-gpu target is unregistered; the same failure occurs on the unchanged base revision. Fixes apache#20884
This branch has not been deployed
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 matmul rule can abort default CUDA compilation when a consumer block contains a conditional statement. In #20884, a valid FP32 matmul followed by a conditional epilogue raises a
ScheduleErrorduringcompute_inline, preventing subsequent schedule rules from handling the function.This change makes consumer-chain inlining report failure when a required inline operation is rejected or consumers remain unfused. Both
MatmulandMetalMatmuldiscard the partially transformed schedule and returnNone, allowing subsequent rules to schedule the original function.Regression tests cover redundant conditionals, conditionals with different branch results, preservation of the original function after a rejected schedule, and end-to-end CUDA execution through the default Relax compilation pipeline.
Fixes #20884.
Validation
Tested on an NVIDIA H200 with CUDA 12.2 and NVCC:
tests/python/s_tir/dlight/test_gpu_matmul_conditional.py: 4 passed. Both scheduling regression cases fail at the reportedcompute_inlinecall on the baseline before the fix. The CUDA execution cases check three inputs each against NumPy using exact comparisons.test_gpu_matmul.py,test_gpu_matmul_tensorize.py, andtest_gpu_fallback.py: 16 passed, 1 failed.StmtSimplifyfirst and passes exact NumPy comparisons for all three inputs..pre-commit-config.yamlon both changed files: passed. Hooks with no matching files were skipped.git diff --checkalso passed.The existing failure is
test_gpu_fallback.py::test_fallback_reduction: this build does not recognizeTarget("apple/m1-gpu"), so the test fails during target creation. The unchanged baseline sources reproduce the same failure.