diff --git a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp index ba230ee03653..52ed5b9ca193 100644 --- a/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp +++ b/apps/depthwise_separable_conv/depthwise_separable_conv_generator.cpp @@ -98,10 +98,11 @@ class DepthwiseSeparableConvolution : public Generator &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..bbaadb672682 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 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); diff --git a/src/Func.cpp b/src/Func.cpp index 468188530c67..fd851cfe6d7c 100644 --- a/src/Func.cpp +++ b/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(); diff --git a/src/Func.h b/src/Func.h index 4df562e272ca..f149191701f8 100644 --- a/src/Func.h +++ b/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 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..a2126c451870 --- /dev/null +++ b/test/correctness/gpu_max_registers.cpp @@ -0,0 +1,124 @@ +// 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(-1).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; + } + // 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; +} + +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 negative number of registers means nothing, so it is rejected rather + // than quietly treated as zero. + bool ok = true; + 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); + f.gpu_tile(x, y, xi, yi, 8, 8).gpu_max_registers(bad); + }); + } + if (!ok) { + return 1; + } +#endif + + printf("Success!\n"); + return 0; +}