Add scheduling directive to cap the max GPU registers

apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp

@@ -98,10 +98,11 @@ class DepthwiseSeparableConvolution : public Generator<DepthwiseSeparableConvolu
             output.dim(2).set_estimate(0, H);
             output.dim(3).set_estimate(0, N);
         } else if (get_target().has_gpu_feature()) {
-            // 0.066ms on a 2060 RTX super. This is about 1.2 TFlops,
-            // which is not a very large fraction of peak. For
-            // comparison though, tensorflow 2.3 achieves 0.13ms via
-            // cudnn 7. So we're twice as fast.
+            // 0.032ms on an RTX 5060 Ti. This is about 2.5 TFlops, which is
+            // not a very large fraction of peak. For comparison, pytorch 2.11
+            // on the same card takes 0.036ms for the same pipeline when given
+            // the same channels-innermost layout, and 0.079ms in its own
+            // default layout.
 
             // This schedule fuses the depthwise conv into the pointwise
             // conv. The results of the depthwise conv are computed inside
@@ -110,16 +111,24 @@ class DepthwiseSeparableConvolution : public Generator<DepthwiseSeparableConvolu
             Var xi, yi, di, dii, xii, yii;
             RVar ro, ri;
 
-            // The pointwise convolution kernel. Produces a 4x4 tile of output.
+            // The pointwise convolution kernel. Produces a 4x2 tile of output.
             Func(output)
-                .tile({d, x, y}, {di, xi, yi}, {16, 4, 4})
+                .tile({d, x, y}, {di, xi, yi}, {16, 4, 2})
                 .tile({di, xi, yi}, {dii, xii, yii}, {1, 2, 2})
                 .gpu_threads(di, xi, yi)
                 .fuse(y, b, b)
                 .gpu_blocks(d, x, b)
                 .unroll(xii)
                 .unroll(yii)
                 .unroll(dii);
+            // ptxas picks 80 registers here, and asking for 64 is worth 0.5%
+            // on an RTX 5060 Ti. The work and the theoretical occupancy are
+            // the same either way - the same loads, stores and FFMAs, nothing
+            // spilled, and a processor holds 24 blocks regardless - but it
+            // keeps more warps in flight, 5.63 active warps per scheduler
+            // against 5.56. The effect is not monotonic in the number, so it
+            // is worth sweeping: 72 is worse than either.
+            output.gpu_max_registers(64);
 
             pointwise_convolved.compute_at(output, di)
                 .reorder(x, y, d)
@@ -154,11 +163,11 @@ class DepthwiseSeparableConvolution : public Generator<DepthwiseSeparableConvolu
                 .unroll(x)
                 .unroll(y);
 
-            // The depthwise convolution kernel. Produces a 4x4 tile
+            // The depthwise convolution kernel. Produces a 4x2 tile
             // of intermediate state, storing the result in shared.
             depthwise_convolved.in()
                 .compute_at(output, d)
-                .tile({d, x, y}, {di, xi, yi}, {32, 4, 4}, TailStrategy::RoundUp)
+                .tile({d, x, y}, {di, xi, yi}, {32, 4, 2}, TailStrategy::RoundUp)
                 .tile({di, xi, yi}, {dii, xii, yii}, {2, 2, 2})
                 .gpu_threads(di, xi, yi)
                 .unroll(xii)

src/CodeGen_GPU_Dev.h

@@ -25,6 +25,12 @@ struct CodeGen_GPU_Dev {
                             const std::string &name,
                             const std::vector<DeviceArgument> &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. */

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<DeviceArgument> &args) override;
 
+    void set_kernel_max_registers(int n) override;
+
     static void test();
 
     std::vector<char> 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<DeviceArgument> &args) {
@@ -286,6 +295,15 @@ void CodeGen_PTX_Dev::add_kernel(Stmt stmt,
                  << "x" << block_size.extent[2] << "\n";
     }
 
+    // A schedule can ask ptxas for a different number of registers per thread
+    // than it would choose for itself, trading how many blocks fit on a
+    // processor against how much it must spill.
+    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);
 

src/Func.cpp

@@ -2615,6 +2615,15 @@ 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 non-negative 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();

src/Func.h

@@ -2728,6 +2728,22 @@ class Func {
      * on MemoryType for more detail. */
     Func &store_in(MemoryType memory_type);
 
+    /** 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. 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
+     * 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

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)

src/Lower.cpp

@@ -496,7 +496,7 @@ void lower_impl(const vector<Function> &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 {

src/OffloadGPULoops.cpp

@@ -96,6 +96,7 @@ class InjectGpuOffload : public IRMutator {
     map<string, bool> state_needed;
 
     const Target &target;
+    const std::map<std::string, Function> &env;
 
     Expr get_state_var(const string &name) {
         // Expr v = Variable::make(type_of<void *>(), 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<std::string, Function> &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<std::string, Function> &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

src/OffloadGPULoops.h

@@ -7,7 +7,11 @@
  * appropriate host runtime module.
  */
 
+#include <map>
+#include <string>
+
 #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<std::string, Function> &env);
 
 }  // namespace Internal
 }  // namespace Halide

src/Schedule.cpp

@@ -239,6 +239,7 @@ struct FuncScheduleContents {
     std::vector<Bound> estimates;
     std::map<std::string, Internal::FunctionPtr> 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;
 }

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,

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