Skip to content

Add scheduling directive to cap the max GPU registers - #9389

Open
abadams wants to merge 13 commits into
mainfrom
abadams/gpu_max_registers
Open

Add scheduling directive to cap the max GPU registers#9389
abadams wants to merge 13 commits into
mainfrom
abadams/gpu_max_registers

Conversation

@abadams

@abadams abadams commented Aug 25, 2026

Copy link
Copy Markdown
Member

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.

abadams and others added 13 commits August 25, 2026 10:20
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 alexreinking left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. Change the spelling to .cuda_max_registers
  2. Change the parameter format to .gpu_max_registers(DeviceAPI::CUDA, 64). Then warn if anything besides CUDA is passed.

@codecov

codecov Bot commented Aug 26, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 84.37500% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.08%. Comparing base (c049b16) to head (6d0e76e).
⚠️ Report is 3 commits behind head on main.

Files with missing lines Patch % Lines
src/OffloadGPULoops.cpp 75.00% 0 Missing and 3 partials ⚠️
src/CodeGen_PTX_Dev.cpp 66.66% 1 Missing and 1 partial ⚠️
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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@mcourteaux

Copy link
Copy Markdown
Contributor

Change the parameter format to .gpu_max_registers(DeviceAPI::CUDA, 64). Then warn if anything besides CUDA is passed.

I like two. I don't want warnings everywhere. Nor is cuda_... future proof.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants