From 8922f9dd22c7b520d7aa22b795add736009e70bb Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 25 Aug 2026 10:20:23 -0700 Subject: [PATCH 01/13] Add a gpu_max_registers scheduling directive 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 --- src/CodeGen_GPU_Dev.h | 6 ++ src/CodeGen_PTX_Dev.cpp | 18 ++++ src/Func.cpp | 8 ++ src/Func.h | 14 +++ src/Generator.h | 1 + src/Lower.cpp | 2 +- src/OffloadGPULoops.cpp | 29 +++++- src/OffloadGPULoops.h | 7 +- src/Schedule.cpp | 10 +++ src/Schedule.h | 7 ++ test/correctness/CMakeLists.txt | 1 + test/correctness/gpu_max_registers.cpp | 119 +++++++++++++++++++++++++ 12 files changed, 216 insertions(+), 6 deletions(-) create mode 100644 test/correctness/gpu_max_registers.cpp diff --git a/src/CodeGen_GPU_Dev.h b/src/CodeGen_GPU_Dev.h index be56625dac55..2d326e309026 100644 --- a/src/CodeGen_GPU_Dev.h +++ b/src/CodeGen_GPU_Dev.h @@ -25,6 +25,12 @@ struct CodeGen_GPU_Dev { const std::string &name, const std::vector &args) = 0; + /** Cap the registers a thread of the next kernel added may use. Zero, the + * default, leaves it to the backend compiler. Only CUDA does anything with + * this; the other APIs offer no equivalent and ignore it. */ + virtual void set_kernel_max_registers(int n) { + } + /** (Re)initialize the GPU kernel module. This is separate from compile, * since a GPU device module will often have many kernels compiled into it * for a single pipeline. */ diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 399795cc3f56..26e55d36945a 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -47,6 +47,8 @@ class CodeGen_PTX_Dev : public CodeGen_LLVM, public CodeGen_GPU_Dev { const std::string &name, const std::vector &args) override; + void set_kernel_max_registers(int n) override; + static void test(); std::vector compile_to_src() override; @@ -63,6 +65,9 @@ class CodeGen_PTX_Dev : public CodeGen_LLVM, public CodeGen_GPU_Dev { protected: using CodeGen_LLVM::visit; + /** What the schedule asked for, if anything. Zero leaves it to ptxas. */ + int kernel_max_registers = 0; + /** (Re)initialize the PTX module. This is separate from compile, since * a PTX device module will often have many kernels compiled into it for * a single pipeline. */ @@ -194,6 +199,10 @@ class BlockSize : public IRVisitor { bool known = true; }; +void CodeGen_PTX_Dev::set_kernel_max_registers(int n) { + kernel_max_registers = n; +} + void CodeGen_PTX_Dev::add_kernel(Stmt stmt, const std::string &name, const std::vector &args) { @@ -286,6 +295,15 @@ void CodeGen_PTX_Dev::add_kernel(Stmt stmt, << "x" << block_size.extent[2] << "\n"; } + // A schedule can ask for fewer registers per thread than ptxas would + // choose. That lets more blocks be resident at once, and stops it covering + // the latency of a load by issuing it far ahead of its use. + if (kernel_max_registers > 0) { + function->addFnAttr("nvvm.maxnreg", std::to_string(kernel_max_registers)); + debug(2) << "Kernel " << name << " is capped at " + << kernel_max_registers << " registers per thread\n"; + } + // Now verify the function is ok verifyFunction(*function); diff --git a/src/Func.cpp b/src/Func.cpp index 468188530c67..43c8e5f55ee2 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -2615,6 +2615,14 @@ Func &Func::store_in(MemoryType t) { return *this; } +Func &Func::gpu_max_registers(int n) { + invalidate_cache(); + user_assert(n > 0) << "gpu_max_registers must be given a positive number of " + << "registers, but " << name() << " was given " << n << ".\n"; + func.schedule().gpu_max_registers() = n; + return *this; +} + Func &Func::stream_loads() { invalidate_cache(); Stage(func, func.definition(), 0).stream_loads(); diff --git a/src/Func.h b/src/Func.h index 4df562e272ca..dbc00a63802f 100644 --- a/src/Func.h +++ b/src/Func.h @@ -2728,6 +2728,20 @@ class Func { * on MemoryType for more detail. */ Func &store_in(MemoryType memory_type); + /** Cap the registers a thread may use in the kernel this 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 and + * holding the result in a register until then. It costs whatever that + * scheduling freedom was buying, so it is worth measuring rather than + * guessing: the fastest setting is not always the smallest, and a cap that + * doesn't change how many blocks fit is all cost and no benefit. + * + * Only has an effect when compiling for CUDA, and only when the PTX + * version in use has the .maxnreg directive. Other GPU APIs offer no + * equivalent, and ignore this. */ + Func &gpu_max_registers(int n); + /** Use non-temporal (streaming) loads for every direct read this Func's * pure (initial) definition makes of another Func. Equivalent to calling * stream_loads() on Stage 0; see \ref Stage::stream_loads. To stream the diff --git a/src/Generator.h b/src/Generator.h index 7d7f2d1a7d5c..447391c364f7 100644 --- a/src/Generator.h +++ b/src/Generator.h @@ -2319,6 +2319,7 @@ class GeneratorOutputBase : public GIOBase { HALIDE_FORWARD_METHOD(Func, fuse) HALIDE_FORWARD_METHOD(Func, gpu) HALIDE_FORWARD_METHOD(Func, gpu_blocks) + HALIDE_FORWARD_METHOD(Func, gpu_max_registers) HALIDE_FORWARD_METHOD(Func, gpu_single_thread) HALIDE_FORWARD_METHOD(Func, gpu_threads) HALIDE_FORWARD_METHOD(Func, gpu_tile) diff --git a/src/Lower.cpp b/src/Lower.cpp index 753dadb5f6ec..ef1752074412 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -496,7 +496,7 @@ void lower_impl(const vector &output_funcs, if (t.has_gpu_feature()) { debug(1) << "Offloading GPU loops...\n"; - s = inject_gpu_offload(s, t, any_strict_float); + s = inject_gpu_offload(s, t, any_strict_float, env); debug(2) << "Lowering after splitting off GPU loops:\n" << s << "\n\n"; } else { diff --git a/src/OffloadGPULoops.cpp b/src/OffloadGPULoops.cpp index e67f6fad3f4c..c3ba271d5a4e 100644 --- a/src/OffloadGPULoops.cpp +++ b/src/OffloadGPULoops.cpp @@ -96,6 +96,7 @@ class InjectGpuOffload : public IRMutator { map state_needed; const Target ⌖ + const std::map &env; Expr get_state_var(const string &name) { // Expr v = Variable::make(type_of(), name); @@ -172,10 +173,28 @@ class InjectGpuOffload : public IRMutator { // compile the kernel string kernel_name = c_print_name(unique_name("kernel_" + loop->name)); + // The loop the kernel is made from is named after the Func it came + // from, so the schedule that asked for a register cap can be found + // again here, where the kernel is handed to the backend. + int max_registers = 0; + { + const std::string &n = loop->name; + for (size_t i = 0; i + 2 < n.size(); i++) { + if (n[i] == '.' && n[i + 1] == 's' && isdigit(n[i + 2])) { + auto it = env.find(n.substr(0, i)); + if (it != env.end()) { + max_registers = it->second.schedule().gpu_max_registers(); + } + break; + } + } + } + CodeGen_GPU_Dev *gpu_codegen = cgdev[loop->device_api].get(); user_assert(gpu_codegen != nullptr) << "Loop is scheduled on device " << loop->device_api << " which does not appear in target " << target.to_string() << "\n"; + gpu_codegen->set_kernel_max_registers(max_registers); gpu_codegen->add_kernel(loop, kernel_name, closure_args); // get the actual name of the generated kernel for this loop @@ -247,8 +266,9 @@ class InjectGpuOffload : public IRMutator { } public: - InjectGpuOffload(const Target &target, bool any_strict_float) - : target(target) { + InjectGpuOffload(const Target &target, bool any_strict_float, + const std::map &env) + : target(target), env(env) { Target device_target = target; // For the GPU target we just want to pass the flags, to avoid the // generated kernel code unintentionally having any dependence on the @@ -383,9 +403,10 @@ class FlattenAliasedAllocations : public IRMutator { } // namespace -Stmt inject_gpu_offload(const Stmt &s, const Target &host_target, bool any_strict_float) { +Stmt inject_gpu_offload(const Stmt &s, const Target &host_target, bool any_strict_float, + const std::map &env) { Stmt flattened = FlattenAliasedAllocations()(s); - return InjectGpuOffload(host_target, any_strict_float).inject(flattened); + return InjectGpuOffload(host_target, any_strict_float, env).inject(flattened); } } // namespace Internal diff --git a/src/OffloadGPULoops.h b/src/OffloadGPULoops.h index 97cd7737271f..bdba3d4e3027 100644 --- a/src/OffloadGPULoops.h +++ b/src/OffloadGPULoops.h @@ -7,7 +7,11 @@ * appropriate host runtime module. */ +#include +#include + #include "Expr.h" +#include "Function.h" namespace Halide { @@ -17,7 +21,8 @@ namespace Internal { /** Pull loops marked with GPU device APIs to a separate * module, and call them through the appropriate host runtime module. */ -Stmt inject_gpu_offload(const Stmt &s, const Target &host_target, bool any_strict_float); +Stmt inject_gpu_offload(const Stmt &s, const Target &host_target, bool any_strict_float, + const std::map &env); } // namespace Internal } // namespace Halide diff --git a/src/Schedule.cpp b/src/Schedule.cpp index 948233112b7c..45ba4661c9e0 100644 --- a/src/Schedule.cpp +++ b/src/Schedule.cpp @@ -239,6 +239,7 @@ struct FuncScheduleContents { std::vector estimates; std::map wrappers; MemoryType memory_type = MemoryType::Auto; + int gpu_max_registers = 0; bool memoized = false; bool async = false; // This is an extent of the ring buffer and expected to be a positive integer. @@ -375,6 +376,7 @@ FuncSchedule FuncSchedule::deep_copy( copy.contents->bounds = contents->bounds; copy.contents->estimates = contents->estimates; copy.contents->memory_type = contents->memory_type; + copy.contents->gpu_max_registers = contents->gpu_max_registers; copy.contents->memoized = contents->memoized; copy.contents->memoize_eviction_key = contents->memoize_eviction_key; copy.contents->async = contents->async; @@ -402,6 +404,14 @@ MemoryType &FuncSchedule::memory_type() { return contents->memory_type; } +int FuncSchedule::gpu_max_registers() const { + return contents->gpu_max_registers; +} + +int &FuncSchedule::gpu_max_registers() { + return contents->gpu_max_registers; +} + bool &FuncSchedule::memoized() { return contents->memoized; } diff --git a/src/Schedule.h b/src/Schedule.h index ba3d1eea5ca3..62dda656ccde 100644 --- a/src/Schedule.h +++ b/src/Schedule.h @@ -631,6 +631,13 @@ class FuncSchedule { // @{ MemoryType memory_type() const; MemoryType &memory_type(); + + /** The most registers a thread of the kernel this Func's loop over gpu + * blocks becomes may use. Zero means let the backend decide. */ + // @{ + int gpu_max_registers() const; + int &gpu_max_registers(); + // @} // @} /** You may explicitly bound some of the dimensions of a function, diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index b2aafd0f789d..f7b70900a571 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -167,6 +167,7 @@ tests( gpu_jit_explicit_copy_to_device.cpp gpu_large_alloc.cpp gpu_many_kernels.cpp + gpu_max_registers.cpp gpu_metal_completion_handler_error_check.cpp gpu_mixed_dimensionality.cpp gpu_mixed_shared_mem_types.cpp diff --git a/test/correctness/gpu_max_registers.cpp b/test/correctness/gpu_max_registers.cpp new file mode 100644 index 000000000000..014bb7653c5d --- /dev/null +++ b/test/correctness/gpu_max_registers.cpp @@ -0,0 +1,119 @@ +// Exercises Func::gpu_max_registers, which caps the registers a thread of the +// kernel may use. The cap is a hint to the backend compiler about a tradeoff it +// would otherwise make on its own, so what is checked here is that it reaches +// the generated code, that it does not change the answer, and that a +// nonsensical cap is rejected. + +#include "Halide.h" +#include "expect_user_error.h" +#include "halide_test_dirs.h" +#include +#include +#include + +using namespace Halide; + +namespace { + +// A kernel with enough live values to have an opinion about registers. +Func make_pipeline(Var x, Var y) { + Func f("f"); + Expr e = cast(x + y); + for (int i = 0; i < 8; i++) { + e = e * e + cast(x - i) - cast(y + i); + } + f(x, y) = e; + return f; +} + +float expected(int x, int y) { + float e = (float)(x + y); + for (int i = 0; i < 8; i++) { + e = e * e + (float)(x - i) - (float)(y + i); + } + return e; +} + +// Compiling for CUDA embeds the PTX in the host assembly, so the .maxnreg +// directive is visible there. No device is needed to check this. +std::string assembly_for(int max_registers) { + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func f = make_pipeline(x, y); + f.gpu_tile(x, y, xi, yi, 8, 8); + if (max_registers > 0) { + f.gpu_max_registers(max_registers); + } + Target t = get_host_target() + .with_feature(Target::CUDA) + .with_feature(Target::CUDACapability80); + std::string path = Internal::get_test_tmp_dir() + "gpu_max_registers.s"; + f.compile_to_assembly(path, {}, "f", t); + std::ifstream in(path); + std::stringstream ss; + ss << in.rdbuf(); + return ss.str(); +} + +bool check_directive_reaches_ptx() { + if (assembly_for(0).find("maxnreg") != std::string::npos) { + printf("FAIL: .maxnreg appeared without being asked for\n"); + return false; + } + if (assembly_for(40).find("maxnreg 40") == std::string::npos) { + printf("FAIL: .maxnreg 40 did not reach the generated PTX\n"); + return false; + } + return true; +} + +bool check_answer() { + Target t = get_jit_target_from_environment(); + if (!t.has_feature(Target::CUDA)) { + printf("[SKIP] Not running the pipeline: target has no CUDA feature.\n"); + return true; + } + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func f = make_pipeline(x, y); + f.gpu_tile(x, y, xi, yi, 8, 8).gpu_max_registers(32); + Buffer out = f.realize({64, 64}, t); + for (int y = 0; y < out.height(); y++) { + for (int x = 0; x < out.width(); x++) { + float want = expected(x, y); + if (out(x, y) != want) { + printf("FAIL: out(%d, %d) = %f, expected %f\n", x, y, out(x, y), want); + return false; + } + } + } + return true; +} + +} // namespace + +int main(int argc, char **argv) { + if (!check_directive_reaches_ptx()) { + return 1; + } + if (!check_answer()) { + return 1; + } + +#if HALIDE_WITH_EXCEPTIONS + // A cap has to be a number of registers, so zero and negatives are + // rejected rather than silently meaning "no cap". + bool ok = true; + for (int bad : {0, -8}) { + ok &= expect_user_error("gpu_max_registers", "gpu_max_registers", [&]() { + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func f = make_pipeline(x, y); + f.gpu_tile(x, y, xi, yi, 8, 8).gpu_max_registers(bad); + }); + } + if (!ok) { + return 1; + } +#endif + + printf("Success!\n"); + return 0; +} From a3bd9bdac1a3c4b5fa750a8e431cd88a44284abe Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 25 Aug 2026 10:20:23 -0700 Subject: [PATCH 02/13] Use gpu_max_registers in the depthwise separable conv app 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 --- .../depthwise_separable_conv_generator.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp index ba230ee03653..19c7e2da08e2 100644 --- a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp +++ b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp @@ -120,6 +120,11 @@ class DepthwiseSeparableConvolution : public Generator Date: Tue, 25 Aug 2026 11:07:30 -0700 Subject: [PATCH 03/13] Call gpu_max_registers on the output buffer directly 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 --- .../depthwise_separable_conv_generator.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp index 19c7e2da08e2..27534a9afb0f 100644 --- a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp +++ b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp @@ -124,7 +124,7 @@ class DepthwiseSeparableConvolution : public Generator Date: Tue, 25 Aug 2026 11:37:26 -0700 Subject: [PATCH 04/13] Say that gpu_max_registers works in both directions 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 --- src/Func.h | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/src/Func.h b/src/Func.h index dbc00a63802f..daa6c97c2929 100644 --- a/src/Func.h +++ b/src/Func.h @@ -2728,14 +2728,18 @@ class Func { * on MemoryType for more detail. */ Func &store_in(MemoryType memory_type); - /** Cap the registers a thread may use in the kernel this 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 + /** Set how many registers a thread may use in the kernel this Func's loop + * over gpu blocks becomes. This is a budget in both directions, not just a + * cap: left alone, the backend compiler picks a number that fits a certain + * number of blocks on one of the GPU's processors, and this overrides that + * choice in whichever direction you ask for. + * + * Fewer registers per thread fits more blocks, and stops the compiler * covering the latency of a load by issuing it far ahead of its use and - * holding the result in a register until then. It costs whatever that - * scheduling freedom was buying, so it is worth measuring rather than - * guessing: the fastest setting is not always the smallest, and a cap that - * doesn't change how many blocks fit is all cost and no benefit. + * holding the result in a register until then. More registers fits fewer + * blocks, but keeps more of the working set in registers. Either can win, + * so measure rather than guess. A number that doesn't change how many + * blocks fit is all cost and no benefit. * * Only has an effect when compiling for CUDA, and only when the PTX * version in use has the .maxnreg directive. Other GPU APIs offer no From ff76877859b23c0e291f7ed13204ba06acd2d0e1 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 25 Aug 2026 11:41:48 -0700 Subject: [PATCH 05/13] Describe gpu_max_registers without claiming a mechanism 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 --- src/Func.h | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/src/Func.h b/src/Func.h index daa6c97c2929..1a706c9ecae3 100644 --- a/src/Func.h +++ b/src/Func.h @@ -2728,18 +2728,15 @@ class Func { * on MemoryType for more detail. */ Func &store_in(MemoryType memory_type); - /** Set how many registers a thread may use in the kernel this Func's loop - * over gpu blocks becomes. This is a budget in both directions, not just a - * cap: left alone, the backend compiler picks a number that fits a certain - * number of blocks on one of the GPU's processors, and this overrides that - * choice in whichever direction you ask for. - * - * Fewer registers per thread fits more blocks, and stops the compiler - * covering the latency of a load by issuing it far ahead of its use and - * holding the result in a register until then. More registers fits fewer - * blocks, but keeps more of the working set in registers. Either can win, - * so measure rather than guess. A number that doesn't change how many - * blocks fit is all cost and no benefit. + /** Tell the GPU shader compiler to fit the kernel this Func's loop over gpu + * blocks becomes under a given number of registers per thread. A smaller + * budget allows more blocks to be resident on one of the GPU's processors + * at once, but constrains the compiler's instruction scheduling, and may + * make it spill values to memory. + * + * Leaving this unset does not mean no limit. It means the GPU driver picks + * a value automatically, so asking for more registers than it would have + * chosen is also a meaningful thing to do. * * Only has an effect when compiling for CUDA, and only when the PTX * version in use has the .maxnreg directive. Other GPU APIs offer no From 59258c255825d2fd29945d95328edb133ddafc3b Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 25 Aug 2026 11:42:51 -0700 Subject: [PATCH 06/13] Drop the remaining unfounded claims about gpu_max_registers 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 --- .../depthwise_separable_conv_generator.cpp | 8 ++++---- src/CodeGen_PTX_Dev.cpp | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp index 27534a9afb0f..95dada167f2e 100644 --- a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp +++ b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp @@ -120,10 +120,10 @@ class DepthwiseSeparableConvolution : public Generator 0) { function->addFnAttr("nvvm.maxnreg", std::to_string(kernel_max_registers)); debug(2) << "Kernel " << name << " is capped at " From b02352233e88a8ad76b6a13adeaf7b1db0390f34 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 25 Aug 2026 11:47:08 -0700 Subject: [PATCH 07/13] Let gpu_max_registers(0) ask for the automatic choice 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 --- src/Func.cpp | 5 +++-- src/Func.h | 3 ++- test/correctness/gpu_max_registers.cpp | 15 ++++++++++----- 3 files changed, 15 insertions(+), 8 deletions(-) diff --git a/src/Func.cpp b/src/Func.cpp index 43c8e5f55ee2..fd851cfe6d7c 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -2617,8 +2617,9 @@ Func &Func::store_in(MemoryType t) { Func &Func::gpu_max_registers(int n) { invalidate_cache(); - user_assert(n > 0) << "gpu_max_registers must be given a positive number of " - << "registers, but " << name() << " was given " << n << ".\n"; + user_assert(n >= 0) << "gpu_max_registers must be given a non-negative number " + << "of registers, but " << name() << " was given " << n + << ".\n"; func.schedule().gpu_max_registers() = n; return *this; } diff --git a/src/Func.h b/src/Func.h index 1a706c9ecae3..f149191701f8 100644 --- a/src/Func.h +++ b/src/Func.h @@ -2736,7 +2736,8 @@ class Func { * * Leaving this unset does not mean no limit. It means the GPU driver picks * a value automatically, so asking for more registers than it would have - * chosen is also a meaningful thing to do. + * chosen is also a meaningful thing to do. Zero asks for that automatic + * choice, which is what an unscheduled Func gets. * * Only has an effect when compiling for CUDA, and only when the PTX * version in use has the .maxnreg directive. Other GPU APIs offer no diff --git a/test/correctness/gpu_max_registers.cpp b/test/correctness/gpu_max_registers.cpp index 014bb7653c5d..dc57d24986eb 100644 --- a/test/correctness/gpu_max_registers.cpp +++ b/test/correctness/gpu_max_registers.cpp @@ -40,7 +40,7 @@ std::string assembly_for(int max_registers) { Var x("x"), y("y"), xi("xi"), yi("yi"); Func f = make_pipeline(x, y); f.gpu_tile(x, y, xi, yi, 8, 8); - if (max_registers > 0) { + if (max_registers >= 0) { f.gpu_max_registers(max_registers); } Target t = get_host_target() @@ -55,7 +55,7 @@ std::string assembly_for(int max_registers) { } bool check_directive_reaches_ptx() { - if (assembly_for(0).find("maxnreg") != std::string::npos) { + if (assembly_for(-1).find("maxnreg") != std::string::npos) { printf("FAIL: .maxnreg appeared without being asked for\n"); return false; } @@ -63,6 +63,11 @@ bool check_directive_reaches_ptx() { printf("FAIL: .maxnreg 40 did not reach the generated PTX\n"); return false; } + // Zero asks for the automatic choice, which is the same as not asking. + if (assembly_for(0).find("maxnreg") != std::string::npos) { + printf("FAIL: gpu_max_registers(0) still capped the registers\n"); + return false; + } return true; } @@ -99,10 +104,10 @@ int main(int argc, char **argv) { } #if HALIDE_WITH_EXCEPTIONS - // A cap has to be a number of registers, so zero and negatives are - // rejected rather than silently meaning "no cap". + // A negative number of registers means nothing, so it is rejected rather + // than quietly treated as zero. bool ok = true; - for (int bad : {0, -8}) { + for (int bad : {-1, -8}) { ok &= expect_user_error("gpu_max_registers", "gpu_max_registers", [&]() { Var x("x"), y("y"), xi("xi"), yi("yi"); Func f = make_pipeline(x, y); From 53e8943c68e1be88edf42c884b6227a25e6a83eb Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 25 Aug 2026 11:52:00 -0700 Subject: [PATCH 08/13] Report the depthwise conv GPU runtime on the card it was tuned on 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 --- .../depthwise_separable_conv_generator.cpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp index 95dada167f2e..c8d025baae5d 100644 --- a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp +++ b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp @@ -98,10 +98,10 @@ class DepthwiseSeparableConvolution : public Generator Date: Tue, 25 Aug 2026 11:54:11 -0700 Subject: [PATCH 09/13] Compare the depthwise conv app against a baseline we can rerun 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 --- .../depthwise_separable_conv_generator.cpp | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp index c8d025baae5d..e93e1e6b1950 100644 --- a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp +++ b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp @@ -99,9 +99,10 @@ class DepthwiseSeparableConvolution : public Generator Date: Tue, 25 Aug 2026 12:08:11 -0700 Subject: [PATCH 10/13] Retune the depthwise conv GPU schedule for the current card 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 --- .../depthwise_separable_conv_generator.cpp | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp index e93e1e6b1950..89ecdc141080 100644 --- a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp +++ b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp @@ -98,7 +98,7 @@ class DepthwiseSeparableConvolution : public Generator Date: Tue, 25 Aug 2026 12:17:46 -0700 Subject: [PATCH 11/13] Say what gpu_max_registers actually buys the depthwise conv app 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 --- .../depthwise_separable_conv_generator.cpp | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp index 89ecdc141080..9965ca70e9af 100644 --- a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp +++ b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp @@ -121,9 +121,12 @@ class DepthwiseSeparableConvolution : public Generator Date: Tue, 25 Aug 2026 12:27:01 -0700 Subject: [PATCH 12/13] Note the occupancy gpu_max_registers actually achieves here 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 --- .../depthwise_separable_conv_generator.cpp | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp index 9965ca70e9af..52ed5b9ca193 100644 --- a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp +++ b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp @@ -122,11 +122,12 @@ class DepthwiseSeparableConvolution : public Generator Date: Tue, 25 Aug 2026 19:38:47 +0000 Subject: [PATCH 13/13] Apply pre-commit auto-fixes --- test/correctness/gpu_max_registers.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/correctness/gpu_max_registers.cpp b/test/correctness/gpu_max_registers.cpp index dc57d24986eb..a2126c451870 100644 --- a/test/correctness/gpu_max_registers.cpp +++ b/test/correctness/gpu_max_registers.cpp @@ -7,9 +7,9 @@ #include "Halide.h" #include "expect_user_error.h" #include "halide_test_dirs.h" -#include #include #include +#include using namespace Halide;