Shrink-wrap block-level GPU register allocations - #9343
Open
abadams wants to merge 3 commits into
Open
Conversation
alexreinking
requested changes
Aug 14, 2026
alexreinking
left a comment
alexreinking
left a comment
Member
There was a problem hiding this comment.
Let's see if we can get Claude to lean more on the lambda mutators
Comment on lines
+22
to
+47
| class FindAccesses : public IRVisitor { | ||
| using IRVisitor::visit; | ||
|
|
||
| void visit(const Store *op) override { | ||
| if (op->name == alloc) { | ||
| indices.push_back(op->index); | ||
| } | ||
| IRVisitor::visit(op); | ||
| } | ||
|
|
||
| void visit(const Load *op) override { | ||
| if (op->name == alloc) { | ||
| indices.push_back(op->index); | ||
| } | ||
| IRVisitor::visit(op); | ||
| } | ||
|
|
||
| const string &alloc; | ||
|
|
||
| public: | ||
| vector<Expr> indices; | ||
|
|
||
| FindAccesses(const string &alloc) | ||
| : alloc(alloc) { | ||
| } | ||
| }; |
Member
There was a problem hiding this comment.
This would be a ~5 line visit_with
Comment on lines
+105
to
+110
| public: | ||
| using IRMutator::mutate; | ||
|
|
||
| private: | ||
| using IRMutator::visit; | ||
|
|
Member
There was a problem hiding this comment.
Why? Shouldn't these both be protected, anyway?
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## abadams/gpu_register_crosstalk_check #9343 +/- ##
========================================================================
- Coverage 70.13% 70.06% -0.08%
========================================================================
Files 260 261 +1
Lines 79264 79338 +74
Branches 19320 19337 +17
========================================================================
- Hits 55594 55586 -8
- Misses 17898 17920 +22
- Partials 5772 5832 +60 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
An allocation in MemoryType::Register outside the loops over GPU threads is storage private to a thread, so what looks like one allocation of many elements is really a handful of registers held by each thread. The cross-talk check established that each thread keeps to its own part; this shrinks the allocation to just that part. Two accesses by one thread are to the same elements when the distance between them is the same whatever thread it is, because the thread cancels when only comparing accesses made by the same one. That is the question get_subtile already answers for tile memory, so ask it: group the accesses into sets that are each identical or disjoint, reject a partial overlap, and give each set registers of its own. Nothing about how an access covers its elements matters, because the registers a set gets are its own, so a dense ramp reaches them all. A thread that indexes its own storage dynamically has no fixed register to use and gets a user error saying so. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The accumulator now lives at block level in registers, which puts the loop over the reduction above the loop over threads, so one staged panel of each input serves every thread in the block. The panels are copied from global to shared asynchronously, laid over the same grid of threads as the compute so that no thread sits idle in either phase. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
abadams
force-pushed
the
abadams/gpu_register_sgemm
branch
from
August 16, 2026 22:21
0ba4a36 to
8861d3b
Compare
Comment on lines
+7
to
+11
| Target target = get_jit_target_from_environment(); | ||
| if (!target.has_gpu_feature()) { | ||
| printf("[SKIP] No GPU target enabled.\n"); | ||
| return 0; | ||
| } |
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
2 participants
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.
#9325 checks that block-level register storage is safe, but it still duplicates that storage per-thread. This new pass, which runs later, decomposes that into a set of actual ptx registers for each thread to use by identifying which subtiles of the block-level allocation are actually used by each thread. Without this, each thread writes to a dynamic thread-id-dependent index into the allocation, resulting in it spilling to local instead of truly being a register. There doesn't seem to be a good way to merge this with the earlier PR's pass, unfortunately. That one checks legality, and then this one is an optimization that exploits the legality. The need to run at different times.
This is the missing ingredient required to hit >90% of cublas performance in the cuda mat mul app for float32 (i.e. cuda cores, not tensor cores). #9283 hits similar performance levels for narrower float mat muls.