diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index ec85e5a383f9..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 @@ -802,26 +841,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].group[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,7 +901,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 +938,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); } @@ -1107,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); @@ -1123,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; @@ -1137,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) { @@ -1190,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/src/IR.cpp b/src/IR.cpp index ae0e38a36251..f343ad83383b 100644 --- a/src/IR.cpp +++ b/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", diff --git a/src/IR.h b/src/IR.h index 000d8a994851..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,6 +681,16 @@ 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, // Marks the point where profiling should start counting pipeline instances 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. 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); } } 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");