From cceca7b9043de3af1e3c91a4e017e3d0bd06b2a0 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 11 Aug 2026 11:52:26 -0700 Subject: [PATCH 1/5] Give each site of a GPU register allocation its own registers An allocation in MemoryType::Register outside the loops over GPU threads is storage private to a thread, so what looks like one allocation of many elements is really a handful of registers held by each thread. The cross-talk check established that each thread keeps to its own part; this shrinks the allocation to just that part. Two accesses by one thread are to the same elements when the distance between them is the same whatever thread it is, because the thread cancels when only comparing accesses made by the same one. That is the question get_subtile already answers for tile memory, so ask it: group the accesses into sets that are each identical or disjoint, reject a partial overlap, and give each set registers of its own. Nothing about how an access covers its elements matters, because the registers a set gets are its own, so a dense ramp reaches them all. A thread that indexes its own storage dynamically has no fixed register to use and gets a user error saying so. Co-Authored-By: Claude Opus 5 --- Makefile | 2 + src/CMakeLists.txt | 2 + src/Lower.cpp | 5 + src/PromoteGPURegisters.cpp | 193 ++++++++++++++++++ src/PromoteGPURegisters.h | 36 ++++ .../gpu_register_at_block_level.cpp | 78 ++++--- test/error/CMakeLists.txt | 1 + test/error/gpu_register_dynamic_index.cpp | 42 ++++ 8 files changed, 334 insertions(+), 25 deletions(-) create mode 100644 src/PromoteGPURegisters.cpp create mode 100644 src/PromoteGPURegisters.h create mode 100644 test/error/gpu_register_dynamic_index.cpp diff --git a/Makefile b/Makefile index 5435d670c42c..bc6f8a104121 100644 --- a/Makefile +++ b/Makefile @@ -560,6 +560,7 @@ SOURCE_FILES = \ Parameter.cpp \ PartitionLoops.cpp \ Pipeline.cpp \ + PromoteGPURegisters.cpp \ Prefetch.cpp \ PrintLoopNest.cpp \ Profiling.cpp \ @@ -769,6 +770,7 @@ HEADER_FILES = \ Parameter.h \ PartitionLoops.h \ Pipeline.h \ + PromoteGPURegisters.h \ Prefetch.h \ PrefetchDirective.h \ Profiling.h \ diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index aa86aa833a49..551c8ab708db 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -175,6 +175,7 @@ target_sources( Parameter.h PartitionLoops.h Pipeline.h + PromoteGPURegisters.h Prefetch.h PrefetchDirective.h Profiling.h @@ -354,6 +355,7 @@ target_sources( Parameter.cpp PartitionLoops.cpp Pipeline.cpp + PromoteGPURegisters.cpp Prefetch.cpp PrintLoopNest.cpp Profiling.cpp diff --git a/src/Lower.cpp b/src/Lower.cpp index 327382f296eb..ec84cf024daa 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -51,6 +51,7 @@ #include "OffloadGPULoops.h" #include "PartitionLoops.h" #include "Prefetch.h" +#include "PromoteGPURegisters.h" #include "Profiling.h" #include "PurifyIndexMath.h" #include "Qualify.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..b95ab488e20a --- /dev/null +++ b/src/PromoteGPURegisters.cpp @@ -0,0 +1,193 @@ +#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. +class FindAccesses : public IRVisitor { + using IRVisitor::visit; + + void visit(const Store *op) override { + if (op->name == alloc) { + indices.push_back(op->index); + } + IRVisitor::visit(op); + } + + void visit(const Load *op) override { + if (op->name == alloc) { + indices.push_back(op->index); + } + IRVisitor::visit(op); + } + + const string &alloc; + +public: + vector indices; + + FindAccesses(const string &alloc) + : alloc(alloc) { + } +}; + +// Which kinds of loop over the threads of a block appear in some IR. +class LoopKinds : public IRVisitor { + using IRVisitor::visit; + + void visit(const For *op) override { + threads = threads || op->for_type == ForType::GPUThread; + lanes = lanes || op->for_type == ForType::GPULane; + IRVisitor::visit(op); + } + +public: + bool threads = false, lanes = false; +}; + +// Replace each access with the one worked out for it below. +class RewriteAccesses : public IRMutator { +public: + using IRMutator::mutate; + +private: + using IRMutator::visit; + + Expr index_for(const Expr &index) const { + auto it = rewritten.find(index); + internal_assert(it != rewritten.end()); + return it->second; + } + + Stmt visit(const Store *op) override { + Stmt s = IRMutator::visit(op); + if (op->name == alloc) { + op = s.as(); + s = op->with(op->value, index_for(op->index), op->predicate, ModulusRemainder()); + } + return s; + } + + Expr visit(const Load *op) override { + Expr e = IRMutator::visit(op); + if (op->name == alloc) { + op = e.as(); + e = op->with(index_for(op->index), op->predicate, ModulusRemainder()); + } + return e; + } + + const string &alloc; + const map &rewritten; + +public: + RewriteAccesses(const string &alloc, const map &rewritten) + : alloc(alloc), rewritten(rewritten) { + } +}; + +class PromoteGPURegisters : public IRMutator { +public: + using IRMutator::mutate; + +private: + using IRMutator::visit; + + bool in_threads = false; + vector pending; + + Stmt visit(const Allocate *op) override { + LoopKinds kinds; + op->body.accept(&kinds); + // 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) { + FindAccesses finder(op->name); + body.accept(&finder); + + // 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 : finder.indices) { + 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 = RewriteAccesses(op->name, rewritten).mutate(body); + + return op->with({make_const(Int(32), size)}, op->condition, body); + } +}; + +} // namespace + +Stmt promote_gpu_registers(const Stmt &s) { + return PromoteGPURegisters().mutate(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..d2056f9291da --- /dev/null +++ b/test/error/gpu_register_dynamic_index.cpp @@ -0,0 +1,42 @@ +#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"); + return 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; +} From ebebd8747ffd5ade415ceff1ff30018e0dac70b4 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 11 Aug 2026 11:52:27 -0700 Subject: [PATCH 2/5] Keep the SGEMM accumulator in registers, and stage asynchronously The accumulator now lives at block level in registers, which puts the loop over the reduction above the loop over threads, so one staged panel of each input serves every thread in the block. The panels are copied from global to shared asynchronously, laid over the same grid of threads as the compute so that no thread sits idle in either phase. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/Makefile | 6 +- apps/cuda_mat_mul/mat_mul_generator.cpp | 75 ++++++++++++++++++++----- 2 files changed, 65 insertions(+), 16 deletions(-) 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); From 8861d3bd4ddef3d42d0350f15c421c29869ba9d9 Mon Sep 17 00:00:00 2001 From: "halide-ci[bot]" <266445882+halide-ci[bot]@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:28:42 +0000 Subject: [PATCH 3/5] Apply pre-commit auto-fixes --- Makefile | 4 ++-- src/CMakeLists.txt | 4 ++-- src/Lower.cpp | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/Makefile b/Makefile index bc6f8a104121..c7a497ae0502 100644 --- a/Makefile +++ b/Makefile @@ -560,10 +560,10 @@ SOURCE_FILES = \ Parameter.cpp \ PartitionLoops.cpp \ Pipeline.cpp \ - PromoteGPURegisters.cpp \ Prefetch.cpp \ PrintLoopNest.cpp \ Profiling.cpp \ + PromoteGPURegisters.cpp \ PurifyIndexMath.cpp \ PythonExtensionGen.cpp \ Qualify.cpp \ @@ -770,10 +770,10 @@ HEADER_FILES = \ Parameter.h \ PartitionLoops.h \ Pipeline.h \ - PromoteGPURegisters.h \ Prefetch.h \ PrefetchDirective.h \ Profiling.h \ + PromoteGPURegisters.h \ PurifyIndexMath.h \ PythonExtensionGen.h \ Qualify.h \ diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 551c8ab708db..d835d48e61f4 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -175,10 +175,10 @@ target_sources( Parameter.h PartitionLoops.h Pipeline.h - PromoteGPURegisters.h Prefetch.h PrefetchDirective.h Profiling.h + PromoteGPURegisters.h PurifyIndexMath.h PythonExtensionGen.h Qualify.h @@ -355,10 +355,10 @@ target_sources( Parameter.cpp PartitionLoops.cpp Pipeline.cpp - PromoteGPURegisters.cpp 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 ec84cf024daa..753dadb5f6ec 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -51,8 +51,8 @@ #include "OffloadGPULoops.h" #include "PartitionLoops.h" #include "Prefetch.h" -#include "PromoteGPURegisters.h" #include "Profiling.h" +#include "PromoteGPURegisters.h" #include "PurifyIndexMath.h" #include "Qualify.h" #include "RealizationOrder.h" From 30b1089f612f4ebfbab843aaffd12f4838412842 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 17 Aug 2026 12:55:07 -0700 Subject: [PATCH 4/5] Build the cuda mat mul app for the capability it needs Staging the inputs with asynchronous copies needs compute capability 8.0, which the app's Makefile asks for but its CMake build did not, so the generator refused to compile it. Raise the guard in the runner to match, so that the test skips on an older GPU rather than failing on one. Also report an error when the dynamic index test skips, since an error test that returns cleanly reads as a failure to the Makefile's harness. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/CMakeLists.txt | 2 +- apps/cuda_mat_mul/runner.cpp | 9 +++++---- test/error/gpu_register_dynamic_index.cpp | 3 ++- 3 files changed, 8 insertions(+), 6 deletions(-) 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/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/test/error/gpu_register_dynamic_index.cpp b/test/error/gpu_register_dynamic_index.cpp index d2056f9291da..fee1dbca7a35 100644 --- a/test/error/gpu_register_dynamic_index.cpp +++ b/test/error/gpu_register_dynamic_index.cpp @@ -7,7 +7,8 @@ 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"); - return 0; + // An error test has to report an error even when it skips. + _halide_user_assert(0); } Func f("f"), g("g"); From d8492ff8f9b3aa7dffa122f2cf7a00ec9401f9a5 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 17 Aug 2026 12:55:12 -0700 Subject: [PATCH 5/5] Use the lambda visitors to find and rewrite the accesses Co-Authored-By: Claude Opus 5 --- src/PromoteGPURegisters.cpp | 131 ++++++++++++++---------------------- 1 file changed, 51 insertions(+), 80 deletions(-) diff --git a/src/PromoteGPURegisters.cpp b/src/PromoteGPURegisters.cpp index b95ab488e20a..872e9b881850 100644 --- a/src/PromoteGPURegisters.cpp +++ b/src/PromoteGPURegisters.cpp @@ -19,101 +19,73 @@ using std::vector; namespace { // Every access to the allocation, in the order they appear. -class FindAccesses : public IRVisitor { - using IRVisitor::visit; - - void visit(const Store *op) override { - if (op->name == alloc) { - indices.push_back(op->index); - } - IRVisitor::visit(op); - } - - void visit(const Load *op) override { +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); } - IRVisitor::visit(op); - } - - const string &alloc; - -public: - vector indices; - - FindAccesses(const string &alloc) - : alloc(alloc) { - } -}; + 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. -class LoopKinds : public IRVisitor { - using IRVisitor::visit; - - void visit(const For *op) override { - threads = threads || op->for_type == ForType::GPUThread; - lanes = lanes || op->for_type == ForType::GPULane; - IRVisitor::visit(op); - } - -public: +struct LoopKinds { bool threads = false, lanes = false; }; -// Replace each access with the one worked out for it below. -class RewriteAccesses : public IRMutator { -public: - using IRMutator::mutate; - -private: - using IRMutator::visit; +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; +} - Expr index_for(const Expr &index) const { +// 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; - } - - Stmt visit(const Store *op) override { - Stmt s = IRMutator::visit(op); - if (op->name == alloc) { - op = s.as(); - s = op->with(op->value, index_for(op->index), op->predicate, ModulusRemainder()); - } - return s; - } - - Expr visit(const Load *op) override { - Expr e = IRMutator::visit(op); - if (op->name == alloc) { - op = e.as(); - e = op->with(index_for(op->index), op->predicate, ModulusRemainder()); - } - return e; - } - - const string &alloc; - const map &rewritten; - -public: - RewriteAccesses(const string &alloc, const map &rewritten) - : alloc(alloc), rewritten(rewritten) { - } -}; + }; + 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 { -public: - using IRMutator::mutate; - -private: +protected: using IRMutator::visit; bool in_threads = false; vector pending; Stmt visit(const Allocate *op) override { - LoopKinds kinds; - op->body.accept(&kinds); + 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, @@ -154,8 +126,7 @@ class PromoteGPURegisters : public IRMutator { // Give each site its own registers, and wrap the body in the smaller // allocation. Stmt promote(const Allocate *op, Stmt body) { - FindAccesses finder(op->name); - body.accept(&finder); + 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 @@ -166,7 +137,7 @@ class PromoteGPURegisters : public IRMutator { string description = "the allocation " + op->name + ", which is scheduled to live in Register memory outside the " "loops over GPU threads"; - for (const Expr &index : finder.indices) { + 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. @@ -177,7 +148,7 @@ class PromoteGPURegisters : public IRMutator { } int size = subtiles.empty() ? 0 : (int)subtiles.size() * subtiles[0].total_lanes(); - body = RewriteAccesses(op->name, rewritten).mutate(body); + body = rewrite_accesses(body, op->name, rewritten); return op->with({make_const(Int(32), size)}, op->condition, body); } @@ -186,7 +157,7 @@ class PromoteGPURegisters : public IRMutator { } // namespace Stmt promote_gpu_registers(const Stmt &s) { - return PromoteGPURegisters().mutate(s); + return PromoteGPURegisters()(s); } } // namespace Internal