Repository navigation
[FIX][S-TIR] Bind merged shared memory sizes before the merged allocation - #20955
Open
JonSnow1807 wants to merge 1 commit into
Open
JonSnow1807 wants to merge 1 commit into
JonSnow1807 wants to merge 1 commit into
Conversation
…tion MergeSharedMemoryAllocations puts the merged allocation and the tensor views at the start of the launch body. When a shared.dyn shape uses a scalar that is bound later in the body, the merged size and the view type use the variable before its binding. check_well_formed fails, and tvm.compile for CUDA fails in MakePackedAPI. An access offset can have the same problem, when it uses the shape of a tensor that is allocated after the access. Collect the pure scalar bindings of the launch body that always run (not under a condition, a loop or an assert). The merged size, the offsets and the view shapes use copies of these bindings. The copies are placed before the merged allocation, once each and in dependency order. The original bindings do not change. Fixes apache#20870
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.
MergeSharedMemoryAllocationsputs the merged allocation and the tensor views at the start of the launch body, but ashared.dynshape can use a scalar bound later in the body. The merged size then uses the variable before its binding:check_well_formedfails, and the CUDA compile fails inMakePackedAPIwith “variables (cse_v1,) are used, but are not passed in as API arguments”. An access offset that uses a later tensor’s size has the same problem.This collects the pure scalar bindings of the launch body that always run (not under a condition, loop or assert), and places a copy of each one the merged size, offsets and view shapes need before the merged allocation, in dependency order. The original bindings are unchanged, and copying instead of substituting keeps long binding chains from growing the expressions.
Tests: three new tests, including a CUDA compile-and-run check against numpy. All three fail without the change, and the file passes 10 of 10 with it. s_tir and relax show the same failures as main. Run on an RTX PRO 4000 (sm_120), CUDA 12.8.
Fixes #20870
Code drafted by both human and AI, and reviewed by human.