Add scheduling directive to cap the max GPU registers

[abadams]

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.

[alexreinking]

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 Report

❌ Patch coverage is 80.00000% with 7 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.07%. Comparing base (c049b16) to head (e134b2a).
⚠️ Report is 45 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 ⚠️
src/Func.cpp	90.00%	0 Missing and 1 partial ⚠️
src/Lower.cpp	0.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #9389      +/-   ##
==========================================
+ Coverage   70.03%   70.07%   +0.03%     
==========================================
  Files         261      261              
  Lines       79223    79850     +627     
  Branches    19312    19453     +141     
==========================================
+ Hits        55487    55957     +470     
- Misses      17923    18020      +97     
- Partials     5813     5873      +60     

☔ 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]

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.

[abadams]

Actually I'm not sure I get the suggested API. What is the user supposed to say in the world where multiple GPU APIs support this feature and you want to limit the max registers on whatever target we're compiling to, if it supports that? Do they have to awkwardly get a device api from the target? In today's world, to avoid the warning, they're going to have to say:

if (target.has_feature(Target::Cuda)) {
  foo.gpu_max_registers(DeviceAPI::Cuda, 128);
}

which mentions Cuda twice for no good reason. The following seems more natural:

if (target.has_feature(Target::Cuda)) {
  foo.gpu_max_registers(128);
}

and in the multiple-api-support world:

if (target.feature_any_of({Target::Cuda, Target::Metal}) {
  foo.gpu_max_registers(target.get_required_device_api(), 128);
}

vs

if (target.feature_any_of({Target::Cuda, Target::Metal}) {
  foo.gpu_max_registers(128);
}

[mcourteaux]

The suggestion I'd go with takes in the DeviceAPI argument to let the user be explicit about which API it is they want to impose the register limit. Compiling to any other compute API ignores the max_register directive, intentionally silently as the user made clear they want to apply this limit to this specific API.

If instead they do: foo.gpu_max_registers(Metal, 128);, which doesn't support this directive, that warrants a user_warning, as they are scheduling something that today is not supported (but could be potentially in the future).

[abadams]

I understand that rationale in isolation when applied to just the call, but the required surrounding code makes it difficult to justify. Real code will require guarding the call under a target flag in either case. I think this is very ugly:

if (target.feature_any_of({Target::Cuda, Target::Metal}) {
  foo.gpu_max_registers(target.get_required_device_api(), 128);
}

[mcourteaux]

Why would you need the surrounding code? I don't understand. What's the problem with just:

if (target.has_gpu_target()) {
  foo.gpu_tile(...)
     .gpu_max_registers(DeviceAPI::CUDA, 128)
     .gpu_max_registers(DeviceAPI::Metal, 98)
}

When the target has CUDA, it will use 128, when it's Metal it'll pick 98, when it's something else, it will ignore these.

[alexreinking]

Martijn's example is the intended API usage pattern I had in mind

[abadams]

Ah, I understood the suggestion differently, which is why it made no sense to me. The implementation holds a single int, so multiple calls to gpu_max_registers would just clobber each other in my mental model. You're proposing the implementation holds a map from DeviceAPI to int.

[abadams]

The fact that it was surprising to me does give me pause though. Do we have other scheduling directives of the form "if this then that?" Most of them do something unconditionally, right?

[abadams]

My concern more concretely: This adds a scheduling directive which sometimes silently does nothing (if the target's device api does not match the argument). My misunderstanding was because I was assuming it would be an error in this case.

[abadams]

Discussed this with Alex offline and my concerns have been sufficiently allayed. I will merge this once it's green.