Repository navigation
[FIX][S-TIR] Do not reverse-inline a computing consumer into a block with init - #20954
Open
JonSnow1807 wants to merge 1 commit into
Open
JonSnow1807 wants to merge 1 commit into
JonSnow1807 wants to merge 1 commit into
Conversation
…with init DLight's NormalizeFunction removes a unit-extent reduction block var but keeps the init statement. For a matmul with K = 1, the block then has only data-parallel block vars and passes CheckCompleteBlock. reverse_compute_inline then puts the consumer into both the init and the update, so an epilogue f is computed as f(f(0) + a * b) instead of f(a * b). When the producer has an init statement, only allow reverse inline if the consumer copies the value of the inlined tensor to a new location (for example a transpose). This keeps the case that apache#17838 added. Fixes apache#20825
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.
A K = 1 matmul followed by an elementwise op gives wrong values on CUDA. DLight’s
NormalizeFunctionremoves the unit reduction var but the block keeps its init, so it still passes the complete-block check. DLightFallbackthen callsreverse_compute_inlineon the elementwise consumer, which writes it into both the init and the update: an epiloguefbecomesf(f(0) + a * b).#17838 allowed reverse inline into a block with init for a transpose, which is correct because a copy only moves the output. This keeps that case: when the producer has init, the consumer must read the inlined tensor once, as the whole stored value. The same rule rejects a consumer like
D[C[i]] = C[i], which before recursed until the process crashed.The issue’s repro uses
AdjustMatmulOrder, removed in #20712. The same wrong result happens on main witha(10x1)^T @ (b^T @ c)(1x3)and the same exp chain: max abs error 0.1929 before, 0.0 after.Tests: three new tests (two in
test_tir_schedule_compute_inline.py, one intest_gpu_fallback.py). Without the change two fail and the scatter test crashes. s_tir and relax suites show the same failures as main. Run on an RTX PRO 4000 (sm_120), CUDA 12.8.Fixes #20825
Code drafted by both human and AI, and reviewed by human.