diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index ec85e5a383f9..161daac05878 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/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,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/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); } }