Defer GPU shared/heap allocation fusing to the offload site

src/IR.cpp

@@ -652,6 +652,7 @@ constexpr const char *intrinsic_op_names[] = {
     "mod_round_to_zero",
     "mul_shift_right",
     "mux",
+    "offset_pointer",
     "popcount",
     "prefetch",
     "profiling_enable_instance_marker",

src/IR.h

@@ -606,6 +606,7 @@ struct Call : public ExprNode<Call> {
     // are *not* guaranteed to be stable across time.
     enum IntrinsicOp {
         // keep-sorted start sticky_comments=yes
+
         abs,
         // Absolute difference between two values. absd(a, b) = abs(a - b), but
         // without overflow issues for integer types.
@@ -680,6 +681,16 @@ struct Call : public ExprNode<Call> {
         mod_round_to_zero,
         mul_shift_right,
         mux,
+        // A pointer into a sub-range of another allocation.
+        // offset_pointer(base, offset) returns the base pointer of the
+        // allocation named by the Variable `base`, advanced by `offset`
+        // elements (in units of the aliasing allocation's element type). Used
+        // as the new_expr of an Allocate node to express that it aliases a
+        // sub-range of a larger backing allocation. GPU allocation fusing emits
+        // these so that per-Func names survive lowering (for the profiler and
+        // for a readable conceptual stmt); they are folded into the indices of
+        // loads and stores by inject_gpu_offload before GPU codegen.
+        offset_pointer,
         popcount,
         prefetch,
         // Marks the point where profiling should start counting pipeline instances

src/LowerWarpShuffles.cpp

@@ -658,8 +658,12 @@ class LowerWarpShuffles : public IRMutator {
     }
 
     Stmt visit(const Allocate *op) override {
-        if (this_lane.defined() || op->memory_type == MemoryType::GPUShared) {
-            // Not a warp-level allocation
+        if (this_lane.defined() ||
+            op->memory_type == MemoryType::GPUShared ||
+            op->memory_type == MemoryType::Heap) {
+            // Not a warp-level allocation. Warp-level storage is per-lane
+            // register storage; shared and heap (global) memory are never
+            // striped across lanes.
             return IRMutator::visit(op);
         } else {
             // Pick up this allocation and deposit it inside the loop over lanes at reduced size.

src/OffloadGPULoops.cpp

@@ -14,7 +14,9 @@
 #include "IROperator.h"
 #include "IRPrinter.h"
 #include "InjectHostDevBufferCopies.h"
+#include "ModulusRemainder.h"
 #include "OffloadGPULoops.h"
+#include "Scope.h"
 #include "Simplify.h"
 #include "Util.h"
 
@@ -318,10 +320,73 @@ class InjectGpuOffload : public IRMutator {
     }
 };
 
+// Fold the aliasing allocations left behind by GPU allocation fusing back into
+// direct offset accesses of their backing allocation, and drop the aliasing
+// Allocate nodes. This reproduces the flat representation the device code
+// generators expect, and runs just before them so that everything upstream
+// (the profiler, the conceptual stmt) still sees per-Func allocation names.
+class FlattenAliasedAllocations : public IRMutator {
+    using IRMutator::visit;
+
+    struct Alias {
+        std::string backing;
+        Expr offset;
+    };
+    Scope<Alias> aliases;
+
+    Stmt visit(const Allocate *op) override {
+        // offset_pointer is a (non-pure) Intrinsic specifically so CSE/LICM
+        // won't lift it out of new_expr, so it's still a direct call here.
+        const Call *c = op->new_expr.defined() ? op->new_expr.as<Call>() : nullptr;
+        if (c && c->is_intrinsic(Call::offset_pointer)) {
+            const Variable *base = c->args[0].as<Variable>();
+            internal_assert(base) << "offset_pointer base must be a Variable\n";
+            ScopedBinding<Alias> bind(aliases, op->name, Alias{base->name, c->args[1]});
+            return mutate(op->body);
+        }
+        return IRMutator::visit(op);
+    }
+
+    Stmt visit(const Free *op) override {
+        if (aliases.contains(op->name)) {
+            return Evaluate::make(0);
+        }
+        return IRMutator::visit(op);
+    }
+
+    // Adding the offset into the backing allocation shifts the index, so the
+    // load/store alignment relative to the backing base must be recomputed from
+    // the original alignment and the alignment of the offset.
+    ModulusRemainder shift_alignment(const ModulusRemainder &alignment, const Expr &offset) {
+        return alignment + modulus_remainder(offset);
+    }
+
+    Expr visit(const Load *op) override {
+        if (aliases.contains(op->name)) {
+            const Alias &a = aliases.get(op->name);
+            return Load::make(op->type, a.backing, mutate(op->index) + a.offset,
+                              op->image, op->param, mutate(op->predicate),
+                              shift_alignment(op->alignment, a.offset));
+        }
+        return IRMutator::visit(op);
+    }
+
+    Stmt visit(const Store *op) override {
+        if (aliases.contains(op->name)) {
+            const Alias &a = aliases.get(op->name);
+            return Store::make(a.backing, mutate(op->value), mutate(op->index) + a.offset,
+                               op->param, mutate(op->predicate),
+                               shift_alignment(op->alignment, a.offset));
+        }
+        return IRMutator::visit(op);
+    }
+};
+
 }  // namespace
 
 Stmt inject_gpu_offload(const Stmt &s, const Target &host_target, bool any_strict_float) {
-    return InjectGpuOffload(host_target, any_strict_float).inject(s);
+    Stmt flattened = FlattenAliasedAllocations()(s);
+    return InjectGpuOffload(host_target, any_strict_float).inject(flattened);
 }
 
 }  // namespace Internal

src/PartitionLoops.cpp

@@ -931,12 +931,15 @@ class RenormalizeGPULoops : public IRMutator {
             return mutate(inner);
         } else if (a && in_gpu_loop && !in_thread_loop) {
             internal_assert(a->extents.size() == 1);
-            if (expr_uses_var(a->extents[0], op->name)) {
-                // This var depends on the block index, and is used to
-                // define the size of shared memory. Can't move it
-                // inwards or outwards. Codegen will have to deal with
-                // it when it deduces how much shared or warp-level
-                // memory to allocate.
+            if (expr_uses_var(a->extents[0], op->name) ||
+                expr_uses_var(a->condition, op->name) ||
+                (a->new_expr.defined() && expr_uses_var(a->new_expr, op->name))) {
+                // This var is used in the allocation's extent, condition, or
+                // new_expr, all of which are evaluated in the allocation's
+                // outer scope, so the let can't be moved inside its body. (The
+                // extent case also can't move outwards: it depends on the block
+                // index and codegen deals with it when deducing shared memory
+                // size.)
                 return IRMutator::visit(op);
             } else {
                 Stmt inner = LetStmt::make(op->name, op->value, a->body);

src/RemoveDeadAllocations.cpp

@@ -57,14 +57,22 @@ class RemoveDeadAllocations : public IRMutator {
         allocs.push(op->name, 1);
         Stmt body = mutate(op->body);
 
+        // An aliasing allocation's new_expr may reference the backing
+        // allocation (e.g. offset_pointer(backing, offset)). Mutating it marks
+        // that backing as used, so it isn't mistaken for dead.
+        Expr new_expr = op->new_expr;
+        if (new_expr.defined()) {
+            new_expr = mutate(new_expr);
+        }
+
         if (allocs.contains(op->name) && op->free_function.empty()) {
             allocs.pop(op->name);
             return body;
-        } else if (body.same_as(op->body)) {
+        } else if (body.same_as(op->body) && new_expr.same_as(op->new_expr)) {
             return op;
         } else {
             return Allocate::make(op->name, op->type, op->memory_type, op->extents, op->condition,
-                                  body, op->new_expr, op->free_function, op->padding);
+                                  body, new_expr, op->free_function, op->padding);
         }
     }
 

test/correctness/gpu_reuse_shared_memory.cpp

@@ -165,6 +165,76 @@ int dynamic_shared_test(MemoryType memory_type) {
     return 0;
 }
 
+int repeated_realization_test(MemoryType memory_type) {
+    // A single Func realized several times at the block level, with
+    // non-overlapping lifetimes. Unrolling a serial loop that sits above the
+    // GPU thread loops duplicates the producer's allocation, so several
+    // allocations sharing a name reach the block level of one kernel and get
+    // coalesced into one reused backing allocation. The producer is
+    // tuple-valued so its components carry real names (p.0, p.1) that the
+    // coalescing must preserve.
+    Func p("p"), c("c");
+    Var x("x"), y("y"), xo("xo"), yo("yo"), xi("xi"), yi("yi"), ky("ky");
+
+    p(x, y) = Tuple(x + y, x - y);
+    c(x, y) = p(x, y)[0] + p(x, y + 1)[1];
+
+    c.tile(x, y, xo, yo, xi, yi, 32, 32)
+        .split(yi, ky, yi, 8)
+        .gpu_blocks(xo, yo)
+        .reorder(xi, yi, ky, xo, yo)
+        .gpu_threads(xi, yi)
+        .unroll(ky);
+    p.compute_at(c, ky).store_in(memory_type);
+
+    Buffer<int> out = c.realize({64, 64});
+    for (int y = 0; y < out.height(); y++) {
+        for (int x = 0; x < out.width(); x++) {
+            int correct = (x + y) + (x - (y + 1));
+            if (out(x, y) != correct) {
+                printf("out(%d, %d) = %d instead of %d\n",
+                       x, y, out(x, y), correct);
+                return 1;
+            }
+        }
+    }
+
+    printf("OK\n");
+    return 0;
+}
+
+int repeated_register_realization_test() {
+    // The register-memory analogue of repeated_realization_test: a Func
+    // realized several times inside the thread loops (via an unrolled serial
+    // loop) with disjoint lifetimes. The copies share a name and get coalesced
+    // into one reused register allocation. The producer is tuple-valued so its
+    // components carry real names (p.0, p.1) that the coalescing must preserve.
+    Func p("p"), c("c");
+    Var x("x"), xo("xo"), xi("xi"), u("u");
+
+    p(x) = Tuple(x + 1, x - 1);
+    c(x) = p(x)[0] * 2 + p(x)[1];
+
+    c.split(x, xo, xi, 8)
+        .gpu_blocks(xo)
+        .split(xi, xi, u, 2)
+        .gpu_threads(xi)
+        .unroll(u);
+    p.compute_at(c, u).store_in(MemoryType::Register);
+
+    Buffer<int> out = c.realize({256});
+    for (int x = 0; x < out.width(); x++) {
+        int correct = (x + 1) * 2 + (x - 1);
+        if (out(x) != correct) {
+            printf("out(%d) = %d instead of %d\n", x, out(x), correct);
+            return 1;
+        }
+    }
+
+    printf("OK\n");
+    return 0;
+}
+
 int main(int argc, char **argv) {
     Target t = get_jit_target_from_environment();
     if (!t.has_gpu_feature()) {
@@ -196,6 +266,16 @@ int main(int argc, char **argv) {
                 return 1;
             }
         }
+
+        printf("Running repeated realization test\n");
+        if (repeated_realization_test(memory_type) != 0) {
+            return 1;
+        }
+    }
+
+    printf("Running repeated register realization test\n");
+    if (repeated_register_realization_test() != 0) {
+        return 1;
     }
 
     printf("Success!\n");