Defer GPU shared/heap allocation fusing to the offload site

src/FuseGPUThreadLoops.cpp

@@ -802,26 +802,35 @@ class ExtractSharedAndHeapAllocations : public IRMutator {
             // the cluster (in terms of the alloc_type), and the
             // widest type in the cluster (which may be wider than the
             // alloc_type).
-            string name;
-            Expr total_size = 0;
-            Type widest_type;
             int number_of_allocs = 0;
             for (const auto &alloc : cluster) {
                 number_of_allocs += alloc.group.size();
             }
+
+            // A single shared allocation needs no offset math, so we can name
+            // the backing allocation after the Func directly and skip the
+            // aliasing wrappers entirely, keeping the common case uncluttered.
+            // Anything else (multiple fused allocations, or a heap allocation
+            // that is sliced per-block out of a larger device allocation) gets
+            // a fresh backing name plus one aliasing allocation per Func. The
+            // per-Func names live on the aliasing allocations, so the backing
+            // itself just needs a unique name.
+            const bool simple = number_of_allocs == 1 && memory_type != MemoryType::Heap;
+
+            string name;
+            if (simple) {
+                name = cluster[0].name;
+            } else if (memory_type == MemoryType::Heap) {
+                name = unique_name("global_alloc");
+            } else {
+                name = unique_name("shared_alloc");
+            }
+
+            Expr total_size = 0;
+            Type widest_type = cluster[0].widest_type;
             for (const auto &alloc : cluster) {
-                if (name.empty()) {
+                if (alloc.widest_type.bytes() > widest_type.bytes()) {
                     widest_type = alloc.widest_type;
-                    if (number_of_allocs > 1) {
-                        name = "allocgroup__" + alloc.name;
-                    } else {
-                        name = alloc.name;
-                    }
-                } else {
-                    if (alloc.widest_type.bytes() > widest_type.bytes()) {
-                        widest_type = alloc.widest_type;
-                    }
-                    name += "__" + alloc.name;
                 }
                 int ratio = alloc.widest_type.bytes() / alloc_type.bytes();
                 internal_assert(ratio != 0)
@@ -853,86 +862,66 @@ class ExtractSharedAndHeapAllocations : public IRMutator {
             const string total_size_name = name + ".size";
             Expr total_size_var = Variable::make(Int(32), total_size_name);
 
-            // Make the allocation
+            // Wrap the body in one aliasing allocation per Func, each pointing
+            // at its offset within the backing allocation. Loads and stores
+            // keep their original per-Func names; the offsets get folded in
+            // later by inject_gpu_offload, just before GPU codegen. The offsets
+            // are in units of each allocation's own type; the group offsets
+            // they build on are in units of widest_type across the cluster.
+            if (!simple) {
+                for (int i = (int)(cluster.size()) - 1; i >= 0; i--) {
+                    Expr group_offset = Variable::make(Int(32), name + "." + std::to_string(i) + ".offset");
+                    for (const SharedAllocation &alloc : cluster[i].group) {
+                        Expr offset = group_offset;
+                        internal_assert(alloc.type.bytes() <= widest_type.bytes());
+                        if (alloc.type.bytes() < widest_type.bytes()) {
+                            offset *= (widest_type.bytes() / alloc.type.bytes());
+                        }
+                        offset = simplify(offset);
+                        Expr base = Variable::make(Handle(), name);
+                        // Intrinsic, not PureIntrinsic: this keeps CSE/LICM from
+                        // lifting it out of the aliasing Allocate's new_expr,
+                        // which would move the reference to the backing
+                        // allocation out of the backing allocation's own scope.
+                        Expr aliased = Call::make(Handle(), Call::offset_pointer,
+                                                  {base, offset}, Call::Intrinsic);
+                        s = Allocate::make(alloc.name, alloc.type, alloc.memory_type,
+                                           {alloc.size}, const_true(), s, aliased);
+                    }
+                }
+            }
+
+            // Make the backing allocation.
             if (memory_type == MemoryType::Heap) {
                 global_allocations.push_back(GlobalAllocation{name, total_size, alloc_type});
             } else {
                 s = Allocate::make(name, alloc_type, memory_type,
                                    {total_size_var}, const_true(), s);
             }
 
-            // Define a group offset for each group in the
-            // cluster. The group offsets are in elements of
-            // widest_type across the entire cluster. Using that,
-            // define an individual offset for each allocation in the
-            // group, using units of that allocation's type.
-            for (int i = (int)(cluster.size()) - 1; i >= 0; i--) {
-                Expr group_offset = Variable::make(Int(32), name + "." + std::to_string(i) + ".offset");
-
-                for (const SharedAllocation &alloc : cluster[i].group) {
-                    // Change units, as described above.
-                    Expr offset = group_offset;
-                    internal_assert(alloc.type.bytes() <= widest_type.bytes());
-                    if (alloc.type.bytes() < widest_type.bytes()) {
-                        offset *= (widest_type.bytes() / alloc.type.bytes());
-                    }
-                    offset = simplify(offset);
-
-                    // Rewrite all loads and stores to point to the allocation
-                    // cluster they belong to with the appropriate offset into it.
-                    class RewriteGroupAccess : public IRMutator {
-                        using IRMutator::visit;
-                        Expr visit(const Load *op) override {
-                            if (op->name == alloc_name) {
-                                return Load::make(op->type, cluster_name, mutate(op->index) + offset,
-                                                  op->image, op->param, mutate(op->predicate),
-                                                  op->alignment);
-                            } else {
-                                return IRMutator::visit(op);
-                            }
-                        }
-
-                        Stmt visit(const Store *op) override {
-                            if (op->name == alloc_name) {
-                                return Store::make(cluster_name, mutate(op->value), mutate(op->index) + offset,
-                                                   op->param, mutate(op->predicate), op->alignment);
-                            } else {
-                                return IRMutator::visit(op);
-                            }
-                        }
-                        const string &alloc_name;
-                        const string &cluster_name;
-                        const Expr &offset;
-
-                    public:
-                        RewriteGroupAccess(const string &alloc_name,
-                                           const string &cluster_name,
-                                           const Expr &offset)
-                            : alloc_name(alloc_name), cluster_name(cluster_name), offset(offset) {
-                        }
-                    } rewriter{alloc.name, name, offset};
-                    s = rewriter(s);
-                }
-
-                // Define the group offset in terms of the previous group in the cluster
-                Expr offset;
-                if (i > 0) {
-                    // Build off the last offset
-                    offset = Variable::make(Int(32), name + "." + std::to_string(i - 1) + ".offset");
-                    int ratio = (widest_type.bytes() / cluster[i - 1].widest_type.bytes());
-                    internal_assert(ratio != 0);
-                    offset += simplify((cluster[i - 1].max_size + ratio - 1) / ratio);
-                } else {
-                    if (memory_type == MemoryType::Heap) {
-                        // One slice of a larger global allocation
-                        offset = get_block_id(bs) * total_size_var;
+            // Define the group offsets, each in terms of the previous group in
+            // the cluster.
+            if (!simple) {
+                for (int i = (int)(cluster.size()) - 1; i >= 0; i--) {
+                    string group_offset_name = name + "." + std::to_string(i) + ".offset";
+                    Expr offset;
+                    if (i > 0) {
+                        // Build off the last offset
+                        offset = Variable::make(Int(32), name + "." + std::to_string(i - 1) + ".offset");
+                        int ratio = (widest_type.bytes() / cluster[i - 1].widest_type.bytes());
+                        internal_assert(ratio != 0);
+                        offset += simplify((cluster[i - 1].max_size + ratio - 1) / ratio);
                     } else {
-                        // Base address for shared memory is zero
-                        offset = 0;
+                        if (memory_type == MemoryType::Heap) {
+                            // One slice of a larger global allocation
+                            offset = get_block_id(bs) * total_size_var;
+                        } else {
+                            // Base address for shared memory is zero
+                            offset = 0;
+                        }
                     }
+                    s = LetStmt::make(group_offset_name, simplify(offset), s);
                 }
-
-                s = LetStmt::make(group_offset.as<Variable>()->name, simplify(offset), s);
             }
             s = LetStmt::make(total_size_name, total_size, s);
         }

src/IR.cpp

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

src/IR.h

@@ -699,6 +699,18 @@ 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,
 

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