From 3d7643170fe5fd569cac2ed3d483bf3a44520be8 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 15 Jul 2026 12:59:07 -0700 Subject: [PATCH 1/3] Defer GPU shared/heap allocation fusing to the offload site fuse_gpu_thread_loops still hoists shared, heap, and register allocations and normalizes GPU loop structure, and still computes the lifetime/offset fusion analysis where the barrier structure is intact. But instead of erasing the per-Func allocations and rewriting their loads and stores onto the fused buffer, it now emits a backing allocation plus one aliasing Allocate per Func, whose new_expr is a new offset_pointer(backing, offset) intrinsic. Loads and stores keep their original per-Func names. A new flatten pass in inject_gpu_offload folds each aliasing allocation's offset back into its loads and stores and drops the aliasing nodes, reproducing the flat representation the device backends expect. It runs after the conceptual stmt is captured, so the profiler and the conceptual stmt see per-Func allocation names and stores can be attributed to a specific Func. Supporting fixes: - offset_pointer is a non-pure Intrinsic so CSE/LICM won't lift it out of new_expr and escape the backing allocation's scope. - RemoveDeadAllocations traverses new_expr, so a backing allocation referenced only through aliasing new_exprs isn't dropped as dead. - PartitionLoops no longer moves a let inside an Allocate when the let is used in the allocation's new_expr or condition (evaluated in the outer scope), which previously left a dangling reference for dynamically-sized allocations. - The flatten pass re-derives load/store alignment from the offset, rather than keeping the original alignment across a possibly misaligned shift. Co-Authored-By: Claude Opus 4.8 --- src/FuseGPUThreadLoops.cpp | 131 ++++++++++++++++------------------ src/IR.cpp | 1 + src/IR.h | 12 ++++ src/OffloadGPULoops.cpp | 67 ++++++++++++++++- src/PartitionLoops.cpp | 15 ++-- src/RemoveDeadAllocations.cpp | 12 +++- 6 files changed, 158 insertions(+), 80 deletions(-) diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index ec85e5a383f9..38c776d3169c 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -809,10 +809,19 @@ class ExtractSharedAndHeapAllocations : public IRMutator { 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 distinct backing name plus one aliasing allocation per Func. + const bool simple = number_of_allocs == 1 && memory_type != MemoryType::Heap; + for (const auto &alloc : cluster) { if (name.empty()) { widest_type = alloc.widest_type; - if (number_of_allocs > 1) { + if (!simple) { name = "allocgroup__" + alloc.name; } else { name = alloc.name; @@ -853,7 +862,36 @@ 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 { @@ -861,78 +899,29 @@ class ExtractSharedAndHeapAllocations : public IRMutator { {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()->name, simplify(offset), s); } s = LetStmt::make(total_size_name, total_size, s); } diff --git a/src/IR.cpp b/src/IR.cpp index c5158728f367..a28e0740d6da 100644 --- a/src/IR.cpp +++ b/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", diff --git a/src/IR.h b/src/IR.h index 16016fca819a..9a4534131726 100644 --- a/src/IR.h +++ b/src/IR.h @@ -699,6 +699,18 @@ struct Call : public ExprNode { 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, diff --git a/src/OffloadGPULoops.cpp b/src/OffloadGPULoops.cpp index 7351a038e019..a775d6f4f01f 100644 --- a/src/OffloadGPULoops.cpp +++ b/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 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() : nullptr; + if (c && c->is_intrinsic(Call::offset_pointer)) { + const Variable *base = c->args[0].as(); + internal_assert(base) << "offset_pointer base must be a Variable\n"; + ScopedBinding 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 diff --git a/src/PartitionLoops.cpp b/src/PartitionLoops.cpp index e31b8b9387b2..ae70ea21b54e 100644 --- a/src/PartitionLoops.cpp +++ b/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); diff --git a/src/RemoveDeadAllocations.cpp b/src/RemoveDeadAllocations.cpp index 74a5c32751d2..986a11aa065d 100644 --- a/src/RemoveDeadAllocations.cpp +++ b/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); } } From 7d7e298e4484da50e6ec2ea3dc01f8879a475aef Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 15 Jul 2026 14:33:29 -0700 Subject: [PATCH 2/3] Don't warp-stripe heap allocations in LowerWarpShuffles LowerWarpShuffles treated any non-shared allocation above a lane loop as per-lane register storage, striping its size by the warp size and moving the Allocate inside the lane loop. With GPU allocation fusing now emitting per-Func aliasing allocations, a heap allocation used inside a gpu_lanes loop (e.g. apps/iir_blur) reaches this pass as a real Allocate node and was wrongly striped, which also separated it from its Free and left a dangling reference. Heap is global memory and is never warp-level storage, so skip it alongside GPUShared. On main no heap Allocate nodes reach this pass, so this is a no-op there. Co-Authored-By: Claude Opus 4.8 --- src/LowerWarpShuffles.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/LowerWarpShuffles.cpp b/src/LowerWarpShuffles.cpp index 9f244aad2ce0..ca72d604e4c6 100644 --- a/src/LowerWarpShuffles.cpp +++ b/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. From 9be014915c85c4e75d7a8d28014170edc074cc4d Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 16 Jul 2026 11:18:44 -0700 Subject: [PATCH 3/3] Give fused GPU backing allocations a short unique name The per-Func names now live on the aliasing allocations, so the backing allocation no longer needs a name concatenating every fused Func. Name it unique_name("shared_alloc") or unique_name("global_alloc") instead of the "allocgroup__f1__f2__..." monster name. The single-shared-allocation case still names the backing after its Func, since there are no aliasing wrappers there. Co-Authored-By: Claude Opus 4.8 --- src/FuseGPUThreadLoops.cpp | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index 38c776d3169c..161daac05878 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -802,9 +802,6 @@ 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(); @@ -815,22 +812,25 @@ class ExtractSharedAndHeapAllocations : public IRMutator { // 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 distinct backing name plus one aliasing allocation per Func. + // 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 (!simple) { - 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)