Shrink-wrap block-level GPU register allocations

Makefile

@@ -562,6 +562,7 @@ SOURCE_FILES = \
   Prefetch.cpp \
   PrintLoopNest.cpp \
   Profiling.cpp \
+  PromoteGPURegisters.cpp \
   PurifyIndexMath.cpp \
   PythonExtensionGen.cpp \
   Qualify.cpp \
@@ -770,6 +771,7 @@ HEADER_FILES = \
   Prefetch.h \
   PrefetchDirective.h \
   Profiling.h \
+  PromoteGPURegisters.h \
   PurifyIndexMath.h \
   PythonExtensionGen.h \
   Qualify.h \

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)

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);

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;
     }
 

src/CMakeLists.txt

@@ -177,6 +177,7 @@ target_sources(
     Prefetch.h
     PrefetchDirective.h
     Profiling.h
+    PromoteGPURegisters.h
     PurifyIndexMath.h
     PythonExtensionGen.h
     Qualify.h
@@ -356,6 +357,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,164 @@
+#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.
+vector<Expr> find_accesses(const Stmt &s, const string &alloc) {
+    vector<Expr> 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<Expr, Expr, IRDeepCompare> &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<Store>();
+                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<Load>();
+                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<const Allocate *> 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<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) {
+        vector<Expr> 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<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 : 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