Add scheduling directive to cap the max GPU registers - #9389
Open
abadams wants to merge 13 commits into
Open
Conversation
Caps the registers a thread may use in the kernel a Func's loop over GPU blocks becomes. Fewer registers per thread lets more blocks be resident on one of the GPU's processors, and stops the backend compiler covering the latency of a load by issuing it far ahead of its use. More registers buys the opposite. Which way is better depends on the pipeline, so it is a schedule decision rather than something to infer. Only CUDA does anything with it. The directive becomes an nvvm.maxnreg function attribute, which ptxas turns into .maxnreg. The other GPU APIs offer no equivalent and ignore it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Worth 3% on an RTX 5060 Ti. The interesting part is the direction: the cap is higher than the register count ptxas picks for itself, so the app trades occupancy for keeping more of the accumulator in registers. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
It is one of the methods Generator forwards from an output buffer to the Func behind it, so the wrapper was redundant. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The wording led with lowering the number to fit more blocks, but on the apps measured so far the setting that won was higher than the one ptxas picks for itself, trading resident blocks for keeping more in registers. Neither direction is the default reading. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The comment asserted that a smaller budget stops the shader compiler issuing loads far ahead of their uses, which is not something we have established. Say what the directive does: it constrains instruction scheduling, may cause spilling, and allows more occupancy. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Both comments explained a measurement by a mechanism we have not established: that a smaller register budget stops ptxas hoisting loads, and that the depthwise app gains from keeping its accumulator in registers. What was measured is the register count, the number of blocks that fit, and the runtime. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Zero is already how every layer below spells "no cap": it is the default in the schedule, and add_kernel only attaches the attribute for a positive number. Rejecting it at the API meant a caller passing a value through had to branch around the call to express the default. Only negative numbers are errors now. The test pins the behaviour rather than the guard: zero has to produce a kernel with no .maxnreg, the same as never calling it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The schedule now has a register budget picked on an RTX 5060 Ti, so the headline number should come from the same card. The old figures are kept as the comparison against cudnn they were making, attributed to the 2060 they were measured on. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The comparison was against tensorflow 2.3 on cudnn 7 on a 2060, which nothing here can reproduce. Measure pytorch on the current card instead. The claim of being twice as fast does not survive: pytorch is 0.036ms to our 0.034ms when given the same channels-innermost layout. Its default layout is where the old factor of two came from. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Sweeping the tile sizes by coordinate descent with the register budget moves the output tile from 4x4 to 4x2, and with it the depthwise tile that has to agree with it on how many threads a block has. 0.035ms to 0.032ms. Most of what the register budget was worth is now in the tile size. It was 3% at the old tile and is 0.4% at this one, which is a sign the 80 was compensating for a tile that no longer suited the card. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Measured by running the two schedules alternately, twenty pairs each: 0.5% either way is small, but it is there (paired t of 3.7, and the faster one wins sixteen pairs of twenty). ncu cannot resolve it, since profiling stretches the kernel from 32us to 36us and adds more spread than the effect has size. The disassembly says where it does not come from. Both versions issue the same 284 FFMAs, 76 shared loads, 95 global loads and 2 barriers, neither spills, and a processor holds 24 blocks either way, so occupancy is unchanged. The capped version is even eight instructions longer. What is left is the register allocation and the order of the instructions. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Asking for 64 registers holds more warps on a scheduler than the 80 ptxas picks by itself, 5.63 against 5.56, even though the theoretical occupancy is 50% either way. Say so, and warn that the number is not monotonic: 72 achieves 5.62 and is slower than both, so the setting has to be swept rather than reasoned about. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
alexreinking
requested changes
Aug 25, 2026
Member
There was a problem hiding this comment.
What's here looks fine as-is, but I think this is missing some kind of communication to the user that it does nothing on non-CUDA backends. A user_warning would be warranted here. A few other options:
- Change the spelling to
.cuda_max_registers - Change the parameter format to
.gpu_max_registers(DeviceAPI::CUDA, 64). Then warn if anything besidesCUDAis passed.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #9389 +/- ##
==========================================
+ Coverage 70.03% 70.08% +0.04%
==========================================
Files 261 261
Lines 79223 79390 +167
Branches 19312 19355 +43
==========================================
+ Hits 55487 55643 +156
+ Misses 17923 17910 -13
- Partials 5813 5837 +24 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
I like two. I don't want warnings everywhere. Nor is cuda_... future proof. |
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.
This used to be controlled by HL_CUDA_MAX_REGISTERS. Tuning it is still sadly necessary for some apps (not in main) despite the recent changes. The original comment on that env var said it should be a scheduling directive, so this PR makes it a scheduling directive. For the apps on main, it helps depthwise conv slightly, but has no significant effect elsewhere.
Unfortunately only cuda seems to offer this level of control over the generated code, so it does nothing in other GPU APIs. Rocm also supports it, but we have no rocm backend.