diff --git a/Makefile b/Makefile index 5435d670c42c..c7a497ae0502 100644 --- a/Makefile +++ b/Makefile @@ -563,6 +563,7 @@ SOURCE_FILES = \ Prefetch.cpp \ PrintLoopNest.cpp \ Profiling.cpp \ + PromoteGPURegisters.cpp \ PurifyIndexMath.cpp \ PythonExtensionGen.cpp \ Qualify.cpp \ @@ -772,6 +773,7 @@ HEADER_FILES = \ Prefetch.h \ PrefetchDirective.h \ Profiling.h \ + PromoteGPURegisters.h \ PurifyIndexMath.h \ PythonExtensionGen.h \ Qualify.h \ diff --git a/apps/cuda_mat_mul/CMakeLists.txt b/apps/cuda_mat_mul/CMakeLists.txt index 803f28c5ecdb..cce4869a8460 100644 --- a/apps/cuda_mat_mul/CMakeLists.txt +++ b/apps/cuda_mat_mul/CMakeLists.txt @@ -28,7 +28,7 @@ find_package(Halide REQUIRED) add_halide_generator(mat_mul.generator SOURCES mat_mul_generator.cpp) # Filters -add_halide_library(mat_mul FROM mat_mul.generator FEATURES cuda cuda_capability_50 PARAMS size=1024) +add_halide_library(mat_mul FROM mat_mul.generator FEATURES cuda cuda_capability_80 PARAMS size=1024) # Main executable add_executable(runner runner.cpp) diff --git a/apps/cuda_mat_mul/Makefile b/apps/cuda_mat_mul/Makefile index e0dfb78900fe..2733e4e353fe 100644 --- a/apps/cuda_mat_mul/Makefile +++ b/apps/cuda_mat_mul/Makefile @@ -2,7 +2,9 @@ include ../support/Makefile.inc MATRIX_SIZE ?= 1024 -CUDA_SDK ?= /usr/local/cuda-10.0 +CUDA_TARGET ?= host-cuda-cuda_capability_80 + +CUDA_SDK ?= /usr/local/cuda CXXFLAGS += -I $(CUDA_SDK)/include LDFLAGS += -L $(CUDA_SDK)/lib64 -Wl,-rpath,$(CUDA_SDK)/lib64 @@ -15,7 +17,7 @@ $(GENERATOR_BIN)/mat_mul.generator: mat_mul_generator.cpp $(GENERATOR_DEPS) $(BIN)/%/mat_mul.a: $(GENERATOR_BIN)/mat_mul.generator @mkdir -p $(@D) - $^ -g mat_mul -e $(GENERATOR_OUTPUTS) -o $(@D) target=host-cuda-cuda_capability_50 size=$(MATRIX_SIZE) + $^ -g mat_mul -e $(GENERATOR_OUTPUTS) -o $(@D) target=$(CUDA_TARGET) size=$(MATRIX_SIZE) $(BIN)/%/runner: runner.cpp $(BIN)/%/mat_mul.a @mkdir -p $(@D) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index 6f2cb17c8cd6..27c534010674 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -15,14 +15,21 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { class MatMul : public Halide::Generator { public: GeneratorParam size{"size", 1024}; + // The tile of the output one block computes, the piece of it one thread + // holds in registers, and how much of the reduction is staged at a time. + GeneratorParam block_x{"block_x", 64}; + GeneratorParam block_y{"block_y", 64}; + GeneratorParam reg_x{"reg_x", 4}; + GeneratorParam reg_y{"reg_y", 8}; + GeneratorParam chunk{"chunk", 32}; Input> A{"A"}; Input> B{"B"}; Output> out{"out"}; void generate() { - // 688 us on an RTX 2060 - // cublas is 512 us on the same card + // 162 us on an RTX 5060 Ti + // cublas is 150 us on the same card Var x("x"), y("y"), p("p"); @@ -35,24 +42,64 @@ class MatMul : public Halide::Generator { RVar rxo, rxi; if (!using_autoscheduler()) { + const int bx = block_x, by = block_y, rx = reg_x, ry = reg_y, k = chunk; + const int tx = bx / rx, ty = by / ry; + + // A block computes a block_x by block_y tile of the output with + // tx by ty threads, each holding a reg_x by reg_y tile of the + // accumulator in registers. The accumulator lives at block level + // so that the loop over the reduction can sit above the loop over + // threads, which lets one staged panel of each input serve every + // thread in the block. out.bound(x, 0, size) .bound(y, 0, size) - .tile(x, y, xi, yi, 64, 16) - .tile(xi, yi, xii, yii, 4, 8) + .tile(x, y, xi, yi, bx, by) + .tile(xi, yi, xii, yii, rx, ry) .gpu_blocks(x, y) .gpu_threads(xi, yi) + .vectorize(xii) + .unroll(yii); + + prod.compute_at(out, x) + .store_in(MemoryType::Register) + .tile(x, y, xii, yii, rx, ry) + .gpu_threads(x, y) .unroll(xii) .unroll(yii); - prod.compute_at(out, xi) - .vectorize(x) - .unroll(y) - .update() - .reorder(x, y, r) - .vectorize(x) - .unroll(y) - .unroll(r, 8); - A.in().compute_at(prod, r).vectorize(_0).unroll(_1); - B.in().compute_at(prod, r).vectorize(_0).unroll(_1); + + prod.update() + .split(r, rxo, rxi, k) + .tile(x, y, xii, yii, rx, ry) + .reorder(xii, yii, rxi, x, y, rxo) + .gpu_threads(x, y) + .unroll(xii) + .unroll(yii) + .unroll(rxi); + + prod.in().compute_at(out, xi).unroll(x).unroll(y); + + // One panel of each input per block per step of the reduction, + // copied from global to shared by all the threads together. Each + // thread moves four floats at a time, which is the widest + // asynchronous copy the hardware has. Both panels are laid over + // the same grid of threads as the compute, so that no thread sits + // idle in either phase. + Var v("v"), t("t"), ti("ti"), tj("tj"), to("to"); + auto stage = [&](Func f) { + f.compute_at(prod, rxo) + .store_in(MemoryType::GPUSharedAsync) + .split(_0, _0, v, 4) + .fuse(_0, _1, t) + .split(t, t, ti, tx) + .split(t, to, tj, ty) + .gpu_threads(ti, tj) + .reorder(to, ti, tj) + .unroll(to) + .vectorize(v); + }; + stage(A.in()); + stage(B.in()); + A.in().compute_with(B.in(), ti); set_alignment_and_bounds(A, size); set_alignment_and_bounds(B, size); diff --git a/apps/cuda_mat_mul/runner.cpp b/apps/cuda_mat_mul/runner.cpp index 898496632802..fbe5414102eb 100644 --- a/apps/cuda_mat_mul/runner.cpp +++ b/apps/cuda_mat_mul/runner.cpp @@ -10,16 +10,17 @@ using Halide::Runtime::Buffer; using Halide::Tools::benchmark; int main(int argc, char **argv) { - // Our Generator is compiled using cuda_capability_50; if the system running this - // test doesn't have at least that, quietly skip the test. + // Our Generator is compiled using cuda_capability_80, because it stages its + // inputs with asynchronous copies; if the system running this test doesn't + // have at least that, quietly skip the test. const auto *interface = halide_cuda_device_interface(); assert(interface->compute_capability != nullptr); int major, minor; int err = interface->compute_capability(nullptr, &major, &minor); assert(err == 0); int ver = major * 10 + minor; - if (ver < 50) { - printf("[SKIP] This system supports only Cuda compute capability %d.%d, but compute capability 5.0+ is required.\n", major, minor); + if (ver < 80) { + printf("[SKIP] This system supports only Cuda compute capability %d.%d, but compute capability 8.0+ is required.\n", major, minor); return 0; } diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index aa86aa833a49..d835d48e61f4 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -178,6 +178,7 @@ target_sources( Prefetch.h PrefetchDirective.h Profiling.h + PromoteGPURegisters.h PurifyIndexMath.h PythonExtensionGen.h Qualify.h @@ -357,6 +358,7 @@ target_sources( Prefetch.cpp PrintLoopNest.cpp Profiling.cpp + PromoteGPURegisters.cpp PurifyIndexMath.cpp PythonExtensionGen.cpp Qualify.cpp diff --git a/src/Lower.cpp b/src/Lower.cpp index 327382f296eb..753dadb5f6ec 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -52,6 +52,7 @@ #include "PartitionLoops.h" #include "Prefetch.h" #include "Profiling.h" +#include "PromoteGPURegisters.h" #include "PurifyIndexMath.h" #include "Qualify.h" #include "RealizationOrder.h" @@ -364,6 +365,10 @@ void lower_impl(const vector &output_funcs, t.has_feature(Target::Vulkan)) { debug(1) << "Injecting per-block gpu synchronization...\n"; s = fuse_gpu_thread_loops(s); + log("Lowering after fusing GPU thread loops:", s); + + debug(1) << "Promoting GPU register allocations...\n"; + s = promote_gpu_registers(s); log("Lowering after injecting per-block gpu synchronization:", s); } diff --git a/src/PromoteGPURegisters.cpp b/src/PromoteGPURegisters.cpp new file mode 100644 index 000000000000..872e9b881850 --- /dev/null +++ b/src/PromoteGPURegisters.cpp @@ -0,0 +1,164 @@ +#include "PromoteGPURegisters.h" + +#include "IR.h" +#include "IREquality.h" +#include "IRMutator.h" +#include "IROperator.h" +#include "IRVisitor.h" +#include "MultiRamp.h" + +#include + +namespace Halide { +namespace Internal { + +using std::map; +using std::string; +using std::vector; + +namespace { + +// Every access to the allocation, in the order they appear. +vector find_accesses(const Stmt &s, const string &alloc) { + vector indices; + auto note = [&](auto *self, const auto *op) { + if (op->name == alloc) { + indices.push_back(op->index); + } + self->visit_base(op); + }; + visit_with( + s, [&](auto *self, const Store *op) { note(self, op); }, + [&](auto *self, const Load *op) { note(self, op); }); + return indices; +} + +// Which kinds of loop over the threads of a block appear in some IR. +struct LoopKinds { + bool threads = false, lanes = false; +}; + +LoopKinds loop_kinds(const Stmt &s) { + LoopKinds kinds; + visit_with(s, [&](auto *self, const For *op) { + kinds.threads = kinds.threads || op->for_type == ForType::GPUThread; + kinds.lanes = kinds.lanes || op->for_type == ForType::GPULane; + self->visit_base(op); + }); + return kinds; +} + +// Replace each access with the one worked out for it below. +Stmt rewrite_accesses(const Stmt &s, const string &alloc, + const map &rewritten) { + auto index_for = [&](const Expr &index) { + auto it = rewritten.find(index); + internal_assert(it != rewritten.end()); + return it->second; + }; + return mutate_with( + s, + [&](auto *self, const Store *op) { + Stmt s = self->visit_base(op); + if (op->name == alloc) { + const Store *store = s.as(); + s = store->with(store->value, index_for(store->index), store->predicate, + ModulusRemainder()); + } + return s; + }, + [&](auto *self, const Load *op) { + Expr e = self->visit_base(op); + if (op->name == alloc) { + const Load *load = e.as(); + e = load->with(index_for(load->index), load->predicate, ModulusRemainder()); + } + return e; + }); +} + +class PromoteGPURegisters : public IRMutator { +protected: + using IRMutator::visit; + + bool in_threads = false; + vector pending; + + Stmt visit(const Allocate *op) override { + LoopKinds kinds = loop_kinds(op->body); + // An allocation with a loop over lanes inside it is warp-level + // storage, which LowerWarpShuffles stripes across the lanes. Leave it + // alone. Without a loop over threads there is nowhere to put this one, + // and whoever runs it already has it to themselves. + if (!in_threads && op->memory_type == MemoryType::Register && + kinds.threads && !kinds.lanes) { + // Pick it up, and put it back inside the loops over threads. + pending.push_back(op); + return mutate(op->body); + } + return IRMutator::visit(op); + } + + Stmt visit(const For *op) override { + if (op->for_type != ForType::GPUThread || pending.empty()) { + ScopedValue bind(in_threads, + in_threads || op->for_type == ForType::GPUThread || + op->for_type == ForType::GPULane); + return IRMutator::visit(op); + } + + // The outermost loop over threads with allocations to place. Everything + // private to a thread goes inside it. + vector allocs; + allocs.swap(pending); + + Stmt body = op->body; + for (const Allocate *alloc : allocs) { + body = promote(alloc, body); + } + { + ScopedValue bind(in_threads, true); + body = mutate(body); + } + return op->with(op->min, op->max, body); + } + + // Give each site its own registers, and wrap the body in the smaller + // allocation. + Stmt promote(const Allocate *op, Stmt body) { + vector accesses = find_accesses(body, op->name); + + // Each access covers a set of elements, and get_subtile partitions the + // accesses between the distinct sets. Nothing about the layout of a set + // matters here, because the registers it gets are its own, so a dense + // ramp reaches all of them. + vector subtiles; + map rewritten; + string description = "the allocation " + op->name + + ", which is scheduled to live in Register memory outside the " + "loops over GPU threads"; + for (const Expr &index : accesses) { + int subtile = get_subtile(index, description, &subtiles); + // Every subtile has the same shape, and so the same number of + // lanes, because get_subtile rejects accesses that don't. + int lanes = subtiles[subtile].total_lanes(); + Expr base = make_const(index.type().element_of(), subtile * lanes); + rewritten[index] = + lanes == 1 ? base : Ramp::make(base, make_one(base.type()), lanes); + } + + int size = subtiles.empty() ? 0 : (int)subtiles.size() * subtiles[0].total_lanes(); + body = rewrite_accesses(body, op->name, rewritten); + + return op->with({make_const(Int(32), size)}, op->condition, body); + } +}; + +} // namespace + +Stmt promote_gpu_registers(const Stmt &s) { + return PromoteGPURegisters()(s); +} + +} // namespace Internal +} // namespace Halide diff --git a/src/PromoteGPURegisters.h b/src/PromoteGPURegisters.h new file mode 100644 index 000000000000..7ffbbb748369 --- /dev/null +++ b/src/PromoteGPURegisters.h @@ -0,0 +1,36 @@ +#ifndef HALIDE_PROMOTE_GPU_REGISTERS_H +#define HALIDE_PROMOTE_GPU_REGISTERS_H + +/** \file + * + * Defines a lowering pass that turns an allocation in register memory outside + * the loops over GPU threads into one register per thread per site. + */ + +#include "Expr.h" + +namespace Halide { +namespace Internal { + +/** An allocation in MemoryType::Register outside the loops over GPU threads + * describes storage that is private to a thread, so what looks like one + * allocation of many elements is really a handful of registers held by each + * thread. Find the sites it is accessed at, check that a thread's accesses to + * each one always land on the same elements, and that different sites never + * share an element. Then give each site registers of its own and move the + * allocation inside the loops over threads, which is where storage private to + * a thread belongs. + * + * Any two accesses must be provably to the same elements or to none of the same + * elements, which is what MemoryType::Register asks for: "all stores must be at + * constant coordinates". A thread that indexes its own storage dynamically gets + * a user error. + * + * Must run after the loops over threads have been fused, so that there is one + * loop nest for the allocation to move inside of. */ +Stmt promote_gpu_registers(const Stmt &s); + +} // namespace Internal +} // namespace Halide + +#endif diff --git a/test/correctness/gpu_register_at_block_level.cpp b/test/correctness/gpu_register_at_block_level.cpp index 12ed2c467039..84ba03b13305 100644 --- a/test/correctness/gpu_register_at_block_level.cpp +++ b/test/correctness/gpu_register_at_block_level.cpp @@ -63,35 +63,30 @@ int main(int argc, char **argv) { } { - // A tile per thread, walked by serial loops rather than unrolled ones, - // with a tail strategy that lets neighbouring tiles overlap. Threads - // recompute each other's values but each still reads only its own. + // Striped across threads rather than tiled: thread t owns elements t, + // t + 16, t + 32 and t + 48, so it holds four registers. Func f("f"), g("g"); - Var x("x"), y("y"), xo("xo"), yo("yo"), xi("xi"), yi("yi"), xii("xii"), yii("yii"); - Var fxo("fxo"), fyo("fyo"), fxi("fxi"), fyi("fyi"); + Var x("x"), y("y"), xo("xo"), xi("xi"); f(x, y) = x + y * 1000; - g(x, y) = f(x, y) * 2; - g.split(x, xo, xi, 32) - .split(y, yo, yi, 32) - .split(xi, xi, xii, 2) - .split(yi, yi, yii, 2) - .reorder(xii, yii, xi, yi, xo, yo) - .gpu_blocks(xo, yo) - .gpu_threads(xi, yi); + g(x, y) = f(x, y); + g.split(x, xo, xi, 64) + .split(xi, xi, x, 16) + .reorder(xi, x, y) + .unroll(xi) + .gpu_blocks(xo, y) + .gpu_threads(x); f.compute_at(g, xo) .store_in(MemoryType::Register) - .split(x, fxo, fxi, 2) - .split(y, fyo, fyi, 2) - .reorder(fxi, fyi, fxo, fyo) - .gpu_threads(fxo, fyo); - - Buffer result = g.realize({128, 128}, target); - for (int y = 0; y < 128; y++) { - for (int x = 0; x < 128; x++) { - int correct = (x + y * 1000) * 2; - if (result(x, y) != correct) { - printf("tile per thread: result(%d, %d) = %d instead of %d\n", - x, y, result(x, y), correct); + .split(x, xo, xi, 16) + .unroll(xo) + .gpu_threads(xi); + + Buffer result = g.realize({256, 4}, target); + for (int y = 0; y < 4; y++) { + for (int x = 0; x < 256; x++) { + if (result(x, y) != x + y * 1000) { + printf("striped: result(%d, %d) = %d instead of %d\n", + x, y, result(x, y), x + y * 1000); return 1; } } @@ -121,6 +116,39 @@ int main(int argc, char **argv) { } } + { + // Each thread owns four elements accessed as one vector, so the vector + // needs four registers rather than one. + Func f("f"), g("g"); + Var x("x"), y("y"), xo("xo"), yo("yo"), xi("xi"), yi("yi"); + f(x, y) = x + y * 1000; + g(x, y) = f(x, y) * 2; + g.split(x, xo, xi, 32) + .split(y, yo, yi, 8) + .split(xi, xi, x, 4) + .reorder(x, xi, yi, xo, yo) + .gpu_blocks(xo, yo) + .gpu_threads(xi, yi) + .vectorize(x); + f.compute_at(g, xo) + .store_in(MemoryType::Register) + .split(x, xo, xi, 4) + .gpu_threads(xo, y) + .vectorize(xi); + + Buffer result = g.realize({256, 64}, target); + for (int y = 0; y < 64; y++) { + for (int x = 0; x < 256; x++) { + int correct = (x + y * 1000) * 2; + if (result(x, y) != correct) { + printf("vectorized: result(%d, %d) = %d instead of %d\n", + x, y, result(x, y), correct); + return 1; + } + } + } + } + { // A tile per thread, read back through a wrapper computed inside the // loops over threads. The wrapper's own loops start at the thread's diff --git a/test/error/CMakeLists.txt b/test/error/CMakeLists.txt index 5ff9f97d944d..5b949b48a112 100644 --- a/test/error/CMakeLists.txt +++ b/test/error/CMakeLists.txt @@ -74,6 +74,7 @@ tests( fuse_same_var.cpp fuse_vectorized_var_with_rvar.cpp gpu_register_crosstalk.cpp + gpu_register_dynamic_index.cpp gpu_register_shifted_between_threads.cpp gpu_register_stages_disagree.cpp gpu_register_stored_by_one_thread.cpp diff --git a/test/error/gpu_register_dynamic_index.cpp b/test/error/gpu_register_dynamic_index.cpp new file mode 100644 index 000000000000..fee1dbca7a35 --- /dev/null +++ b/test/error/gpu_register_dynamic_index.cpp @@ -0,0 +1,43 @@ +#include "Halide.h" +#include + +using namespace Halide; + +int main(int argc, char **argv) { + Target target = get_jit_target_from_environment(); + if (!target.has_gpu_feature()) { + printf("[SKIP] No GPU target enabled.\n"); + // An error test has to report an error even when it skips. + _halide_user_assert(0); + } + + Func f("f"), g("g"); + Var x("x"), y("y"), xo("xo"), yo("yo"), xi("xi"), yi("yi"), xii("xii"), yii("yii"); + Var fxo("fxo"), fyo("fyo"), fxi("fxi"), fyi("fyi"); + + f(x, y) = x + y * 1000; + g(x, y) = f(x, y) * 2; + + g.split(x, xo, xi, 32) + .split(y, yo, yi, 32) + .split(xi, xi, xii, 2) + .split(yi, yi, yii, 2) + .reorder(xii, yii, xi, yi, xo, yo) + .gpu_blocks(xo, yo) + .gpu_threads(xi, yi); + + // Each thread walks its own 2x2 tile with serial loops rather than + // unrolled ones, so which register to use is only known while running. + // Registers cannot be indexed dynamically. + f.compute_at(g, xo) + .store_in(MemoryType::Register) + .split(x, fxo, fxi, 2) + .split(y, fyo, fyi, 2) + .reorder(fxi, fyi, fxo, fyo) + .gpu_threads(fxo, fyo); + + g.compile_jit(target); + + printf("Success!\n"); + return 0; +}