From 3d7643170fe5fd569cac2ed3d483bf3a44520be8 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 15 Jul 2026 12:59:07 -0700 Subject: [PATCH 1/5] 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/5] 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/5] 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) From b27b5ccaeb7b943e3739b43a7ec3274b46899198 Mon Sep 17 00:00:00 2001 From: "halide-ci[bot]" <266445882+halide-ci[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 16:00:38 +0000 Subject: [PATCH 4/5] Apply pre-commit auto-fixes --- src/IR.h | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/IR.h b/src/IR.h index c59cb4ddddab..3a149a8c2715 100644 --- a/src/IR.h +++ b/src/IR.h @@ -606,6 +606,7 @@ struct Call : public ExprNode { // 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,7 +681,6 @@ 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` @@ -691,7 +691,6 @@ struct Call : public ExprNode { // 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 From b9b4179469b330b6109d077d72260d82f6212de2 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Fri, 17 Jul 2026 11:23:09 -0700 Subject: [PATCH 5/5] Give fused GPU allocations meaningful names via merging The block-level shared/heap and register allocation extractors in fuse_gpu_thread_loops previously disambiguated same-named allocations (e.g. from unrolling a loop that holds an allocation) by suffixing each with a counter. That suffix looks like a tuple index and obscures the Func name the profiler needs for attribution. Instead, keep each Func's own name and coalesce same-named allocations with disjoint lifetimes into one allocation sized to the largest, reusing the storage. Since loads and stores keep their names, the shared/heap extractor's load/store visitors now only track liveness (plus the per-thread index striping for allocations inside thread loops), and the register extractor's load/store visitors are gone entirely. Both extractors coalesce via an operator() override that runs immediately after the mutation. Also drops now-dead code: the AllocGroup::name accumulation, the write-only register_allocations scope, and a redundant local init. Adds correctness coverage for the repeated-realization case in both shared/heap and register memory using tuple-valued Funcs. Co-Authored-By: Claude Opus 4.8 --- src/FuseGPUThreadLoops.cpp | 151 ++++++++++++------- test/correctness/gpu_reuse_shared_memory.cpp | 80 ++++++++++ 2 files changed, 177 insertions(+), 54 deletions(-) diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index 161daac05878..61c9990109d5 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -253,8 +253,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator { struct AllocGroup { AllocGroup() = default; AllocGroup(const SharedAllocation &alloc) - : name(alloc.name), - widest_type(alloc.type), + : widest_type(alloc.type), max_size(alloc.size), memory_type(alloc.memory_type) { group.push_back(alloc); @@ -274,7 +273,6 @@ class ExtractSharedAndHeapAllocations : public IRMutator { max_size = max(max_size, alloc.size / size_ratio); } group.push_back(alloc); - name += "_" + alloc.name; } // Only need to check the back of the vector since we always insert @@ -283,7 +281,6 @@ class ExtractSharedAndHeapAllocations : public IRMutator { return group.back().liveness.max < stage; } - string name; Type widest_type; Expr max_size; // In units of the widest type vector group; // Groups of allocs that should be coalesced together @@ -457,8 +454,6 @@ class ExtractSharedAndHeapAllocations : public IRMutator { return IfThenElse::make(condition, then_case, else_case); } - int alloc_node_counter = 0; - Stmt visit(const Allocate *op) override { user_assert(!op->new_expr.defined()) << "Allocate node inside GPU kernel has custom new expression.\n" @@ -484,7 +479,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator { << "but is scheduled to live in " << op->memory_type << " memory.\n"; SharedAllocation alloc; - alloc.name = op->name + "." + std::to_string(alloc_node_counter++); + alloc.name = op->name; alloc.type = op->type; alloc.liveness = IntInterval(barrier_stage, barrier_stage); alloc.size = 1; @@ -533,13 +528,14 @@ class ExtractSharedAndHeapAllocations : public IRMutator { if (it != shared.end()) { SharedAllocation *alloc = it->second; alloc->liveness.max = barrier_stage; - Expr predicate = mutate(op->predicate); - Expr index = mutate_index(alloc, op->index); - return Load::make(op->type, alloc->name, - index, op->image, op->param, predicate, op->alignment); - } else { - return IRMutator::visit(op); + // The allocation keeps its name, so we only need to rewrite the + // node when the storage is striped across threads. + if (alloc->striped_over_threads) { + return Load::make(op->type, op->name, mutate_index(alloc, op->index), + op->image, op->param, mutate(op->predicate), op->alignment); + } } + return IRMutator::visit(op); } Stmt visit(const Store *op) override { @@ -547,14 +543,12 @@ class ExtractSharedAndHeapAllocations : public IRMutator { if (it != shared.end()) { SharedAllocation *alloc = it->second; alloc->liveness.max = barrier_stage; - Expr predicate = mutate(op->predicate); - Expr index = mutate_index(alloc, op->index); - Expr value = mutate(op->value); - return Store::make(alloc->name, value, index, - op->param, predicate, op->alignment); - } else { - return IRMutator::visit(op); + if (alloc->striped_over_threads) { + return Store::make(op->name, mutate(op->value), mutate_index(alloc, op->index), + op->param, mutate(op->predicate), op->alignment); + } } + return IRMutator::visit(op); } Stmt visit(const LetStmt *op) override { @@ -758,6 +752,51 @@ class ExtractSharedAndHeapAllocations : public IRMutator { vector global_allocations; public: + // A single Func can be realized at several sites at the block level with + // disjoint lifetimes (e.g. when a serial loop above the GPU threads is + // unrolled), producing multiple allocations that share a name. Collapse + // each such set into one allocation spanning the union of their lifetimes + // and sized to the largest, so the Func keeps its own name instead of one + // suffixed per copy. The loads and stores already all use that name, and + // the host-side size computation already accumulates a max across the + // copies, so nothing else needs rewriting. + void merge_repeated_allocations() { + vector merged; + map index; + for (SharedAllocation &s : allocations) { + auto it = index.find(s.name); + if (it == index.end()) { + index[s.name] = merged.size(); + merged.push_back(s); + continue; + } + SharedAllocation &m = merged[it->second]; + internal_assert(m.type == s.type && + m.memory_type == s.memory_type && + m.striped_over_threads == s.striped_over_threads && + m.size_computed_on_host == s.size_computed_on_host) + << "Mismatched allocations share the name " << s.name << "\n"; + internal_assert(m.liveness.max < s.liveness.min || + s.liveness.max < m.liveness.min) + << "Allocations sharing the name " << s.name + << " have overlapping lifetimes and cannot be coalesced\n"; + m.liveness.min = std::min(m.liveness.min, s.liveness.min); + m.liveness.max = std::max(m.liveness.max, s.liveness.max); + if (!m.size_computed_on_host) { + m.size = simplify(max(m.size, s.size)); + } + } + allocations.swap(merged); + } + + // Run the mutator, then coalesce repeated allocations so callers always + // see a finalized allocation list. + Stmt operator()(const Stmt &s) { + Stmt result = mutate(s); + merge_repeated_allocations(); + return result; + } + Stmt rewrap_block(Stmt s, const ExtractBlockSize &bs) { // Combine the allocations into groups that have disjoint @@ -819,7 +858,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator { string name; if (simple) { - name = cluster[0].name; + name = cluster[0].group[0].name; } else if (memory_type == MemoryType::Heap) { name = unique_name("global_alloc"); } else { @@ -1096,9 +1135,6 @@ class ExtractRegisterAllocations : public IRMutator { } } - int alloc_node_counter = 0; - Scope alloc_renaming; - Stmt visit(const Allocate *op) override { if (in_lane_loop) { return IRMutator::visit(op); @@ -1112,10 +1148,8 @@ class ExtractRegisterAllocations : public IRMutator { << "it must live in stack memory, heap memory, or registers. " << "Shared allocations at this loop level are not yet supported.\n"; - ScopedBinding p(register_allocations, op->name, 0); - RegisterAllocation alloc; - alloc.name = op->name + "." + std::to_string(alloc_node_counter++); + alloc.name = op->name; alloc.type = op->type; alloc.size = 1; alloc.loop_var = loop_var; @@ -1126,36 +1160,12 @@ class ExtractRegisterAllocations : public IRMutator { alloc.memory_type = op->memory_type; allocations.push_back(alloc); - { - ScopedBinding bind(alloc_renaming, op->name, alloc.name); - return mutate(op->body); - } - } - - Expr visit(const Load *op) override { - const string *new_name = alloc_renaming.find(op->name); - if (!new_name) { - new_name = &(op->name); - } - return Load::make(op->type, *new_name, mutate(op->index), - op->image, op->param, mutate(op->predicate), - op->alignment); - } - - Stmt visit(const Store *op) override { - const string *new_name = alloc_renaming.find(op->name); - if (!new_name) { - new_name = &(op->name); - } - return Store::make(*new_name, mutate(op->value), mutate(op->index), - op->param, mutate(op->predicate), op->alignment); + return mutate(op->body); } template auto visit_let(const LetOrLetStmt *op) -> decltype(op->body) { - auto body = op->body; - - body = mutate(op->body); + auto body = mutate(op->body); Expr value = mutate(op->value); for (RegisterAllocation &s : allocations) { @@ -1179,12 +1189,45 @@ class ExtractRegisterAllocations : public IRMutator { return visit_let(op); } - Scope register_allocations; string loop_var; public: vector allocations; + // Multiple realizations of the same Func inside the thread loops (e.g. from + // unrolling a loop that holds a register allocation) share a name with + // disjoint, sequential lifetimes. Coalesce them into one allocation sized + // to the largest, so the Func keeps its own name (needed for profiler + // attribution) and reuses the scarce register storage instead of getting + // one array per copy. The loads and stores already all use that name. + void merge_repeated_allocations() { + vector merged; + map index; + for (RegisterAllocation &s : allocations) { + auto it = index.find(s.name); + if (it == index.end()) { + index[s.name] = merged.size(); + merged.push_back(s); + continue; + } + RegisterAllocation &m = merged[it->second]; + internal_assert(m.type == s.type && + m.memory_type == s.memory_type && + m.loop_var == s.loop_var) + << "Mismatched register allocations share the name " << s.name << "\n"; + m.size = simplify(max(m.size, s.size)); + } + allocations.swap(merged); + } + + // Run the mutator, then coalesce repeated allocations so callers always + // see a finalized allocation list. + Stmt operator()(const Stmt &s) { + Stmt result = mutate(s); + merge_repeated_allocations(); + return result; + } + Stmt rewrap(Stmt body, const string &loop_var) { for (RegisterAllocation &alloc : allocations) { if ((!loop_var.empty() && ends_with(alloc.loop_var, loop_var)) || diff --git a/test/correctness/gpu_reuse_shared_memory.cpp b/test/correctness/gpu_reuse_shared_memory.cpp index 37e932d78273..8b0e8b7f9b90 100644 --- a/test/correctness/gpu_reuse_shared_memory.cpp +++ b/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 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 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");