Shrink-wrap block-level GPU register allocations

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 \

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)

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<MatMul> {
 public:
     GeneratorParam<int> 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<int> block_x{"block_x", 64};
+    GeneratorParam<int> block_y{"block_y", 64};
+    GeneratorParam<int> reg_x{"reg_x", 4};
+    GeneratorParam<int> reg_y{"reg_y", 8};
+    GeneratorParam<int> chunk{"chunk", 32};
     Input<Buffer<float, 2>> A{"A"};
     Input<Buffer<float, 2>> B{"B"};
 
     Output<Buffer<float, 2>> 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<MatMul> {
         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);

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

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<Function> &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);
     }
 

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 <map>
+
+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<Expr> 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<Store>();
+            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<Load>();
+            e = op->with(index_for(op->index), op->predicate, ModulusRemainder());
+        }
+        return e;
+    }
+
+    const string &alloc;
+    const map<Expr, Expr, IRDeepCompare> &rewritten;
+
+public:
+    RewriteAccesses(const string &alloc, const map<Expr, Expr, IRDeepCompare> &rewritten)
+        : alloc(alloc), rewritten(rewritten) {
+    }
+};
+
+class PromoteGPURegisters : public IRMutator {
+public:
+    using IRMutator::mutate;
+
+private:
+    using IRMutator::visit;
+
+    bool in_threads = false;
+    vector<const Allocate *> 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<bool> 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<const Allocate *> allocs;
+        allocs.swap(pending);
+
+        Stmt body = op->body;
+        for (const Allocate *alloc : allocs) {
+            body = promote(alloc, body);
+        }
+        {
+            ScopedValue<bool> 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<MultiRamp> subtiles;
+        map<Expr, Expr, IRDeepCompare> 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