Allow split store_at/compute_at over gpu threads for things stored in registers - #9376
Draft
abadams wants to merge 5 commits into
Draft
Allow split store_at/compute_at over gpu threads for things stored in registers#9376abadams wants to merge 5 commits into
abadams wants to merge 5 commits into
Conversation
validate_schedule rejects any parallel loop between where a Func is stored and where it is computed, because the iterations would write the same storage. A loop over GPU threads does not, when the storage is Register: a thread's registers are its own, so each thread gets a copy rather than sharing one, which is what the memory type means. Exempt that one case. What the exemption does not establish is that each thread then keeps to its own copy, and it does not have to, because check_gpu_cross_talk already answers exactly that question, later in lowering, once storage folding has had its say. Its error message is about this configuration in as many words - storage "scheduled outside the loops over GPU threads, so every thread gets its own copy of it rather than sharing one". Being rejected up front is what stopped such schedules from ever reaching it. Both sides are tested, because the interesting property is the division of labour between the two checks rather than either alone. The correctness test is a schedule that is fine and was refused; the error test is one that is not fine, gets past the validator now, and has to be stopped by the cross-talk check. Without this change both fail at the validator, the second for the wrong reason. The five existing cross-talk error tests all compute at the block level, so none of them has a thread loop between store and compute - this configuration could not be reached before, and so was not covered. Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Storage folding is not why the cross-talk check runs where it does: it refuses to descend into a parallel loop at all, so it cannot rewrite these accesses. The check runs late because the whole loop nest is in place by then, and anticipating what every part of the schedule will do to it would be much harder. Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Uses expect_user_error rather than a test/error file, which is where these are heading. The two cases belong together anyway: what is being tested is the division of labour between the schedule validator and the cross-talk check, and either case alone only shows half of it. Matching on the message matters here rather than just on erroring at all, because before this change the schedule was rejected too - by the validator, for a different reason. Reaching the error only takes compiling, so that half runs on a machine with no GPU, which the test/error version could not do. The two halves skip independently and say which one they skipped: without exceptions there is no error to catch, and without a device there is nothing to run. Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
abadams
marked this pull request as draft
August 21, 2026 21:11
The placement this branch allows is only worth having because a serial loop can sit between the storage and the loop over threads, and a producer can slide over that loop while staying computed inside the threads. The test now schedules that, rather than a placement with nothing in between, which was permitted but pointless - it allocated the whole block's worth per thread so each could write one row of it. Scheduling the real thing shows the check refusing it. A slid producer loads what the previous run of the loop stored, and the check looked only at stores earlier in the list, so the load it complained about was the first access there was. Two things let it through. A dimension whose region is the same whatever thread is asking is one every thread walks identically, so a coordinate a load names in it is one every thread names, and it cannot be what makes a load another thread's - leave those out of the comparison. It is the region that has to be asked and not the index: a loop of a thread's own is written the same way by every thread, and only its bounds say which part is whose. Then a store listed after a load counts when the two differ along such a dimension, because that is what carries an allocation from one run of a loop to the next. Where every dimension agrees there is no gap to carry anything, and a later store is simply later, which is what gpu_register_stages_disagree relies on. Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Skipping the dimensions that don't separate threads was not enough on its own, and made the check accept two schedules it used to reject, both with silently wrong answers. Both are a slice added to an existing error test: one leaves a stage serial so a single thread writes the slice on everyone's behalf, the other maps that stage's threads the other way round so the slice is written by the thread with its coordinates transposed. In both, the load of the second slice was excused by the store to the first, because the dimension that told the slices apart was one of the ones being skipped. What a thread reads is whatever was written to the site last, so finding one store of its own that covers the load says nothing if some other store could have landed there afterwards. A store is only somebody else's business if it cannot reach the site at all. So exonerate a store only when it runs in at least as many loops over threads as the load and covers it along the dimensions that separate threads, and require every other store to be provably disjoint from it. The same question was already being asked too weakly before any of this: a Func with a real pure definition and one stage left serial is accepted today and computes the wrong thing, because the pure definition's store covers everything and satisfies "this thread wrote it at some point". Asking about the last writer instead rejects that too. A dimension nothing could be bounded in now counts as separating, so failing to work out where an access reaches stays an error rather than excusing the dimension from the comparison. Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #9376 +/- ##
==========================================
+ Coverage 70.08% 70.11% +0.03%
==========================================
Files 260 260
Lines 79287 79332 +45
Branches 19327 19351 +24
==========================================
+ Hits 55569 55625 +56
+ Misses 17923 17882 -41
- Partials 5795 5825 +30 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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.
This is a follow-up to #9325
That PR allows for storing things outside a gpu-thread-parallel loop that are computed within it, provided a no-crosstalk check passes. This was to support placing Funcs at the blocks level but keeping their contents in registers. However it also should allow store_at/compute_ats that are split over thread loops. It makes those safe too - it's not like we'll try to slide or fold over a parallel intermediate loop. You might reasonably ask: What's the point of such a schedule if it can't slide or fold? The answer is that it can't slide or fold over just the dimensions spanned by the parallel loops, but there may also be serial loops mixed in that we do want to slide or fold over. This happens when you have a group of warps cooperating on something that need to synchronize themselves within an outer serial reduction loop, and you want to slide a producer to that over the outer serial reduction loop, but still have that producer compute_at inside the group of warps so that things can stay in registers instead of crossing between warps.