|
| 1 | +#include "PromoteGPURegisters.h" |
| 2 | + |
| 3 | +#include "IR.h" |
| 4 | +#include "IREquality.h" |
| 5 | +#include "IRMutator.h" |
| 6 | +#include "IROperator.h" |
| 7 | +#include "IRVisitor.h" |
| 8 | +#include "MultiRamp.h" |
| 9 | + |
| 10 | +#include <map> |
| 11 | + |
| 12 | +namespace Halide { |
| 13 | +namespace Internal { |
| 14 | + |
| 15 | +using std::map; |
| 16 | +using std::string; |
| 17 | +using std::vector; |
| 18 | + |
| 19 | +namespace { |
| 20 | + |
| 21 | +// Every access to the allocation, in the order they appear. |
| 22 | +vector<Expr> find_accesses(const Stmt &s, const string &alloc) { |
| 23 | + vector<Expr> indices; |
| 24 | + auto note = [&](auto *self, const auto *op) { |
| 25 | + if (op->name == alloc) { |
| 26 | + indices.push_back(op->index); |
| 27 | + } |
| 28 | + self->visit_base(op); |
| 29 | + }; |
| 30 | + visit_with( |
| 31 | + s, [&](auto *self, const Store *op) { note(self, op); }, |
| 32 | + [&](auto *self, const Load *op) { note(self, op); }); |
| 33 | + return indices; |
| 34 | +} |
| 35 | + |
| 36 | +// Which kinds of loop over the threads of a block appear in some IR. |
| 37 | +struct LoopKinds { |
| 38 | + bool threads = false, lanes = false; |
| 39 | +}; |
| 40 | + |
| 41 | +LoopKinds loop_kinds(const Stmt &s) { |
| 42 | + LoopKinds kinds; |
| 43 | + visit_with(s, [&](auto *self, const For *op) { |
| 44 | + kinds.threads = kinds.threads || op->for_type == ForType::GPUThread; |
| 45 | + kinds.lanes = kinds.lanes || op->for_type == ForType::GPULane; |
| 46 | + self->visit_base(op); |
| 47 | + }); |
| 48 | + return kinds; |
| 49 | +} |
| 50 | + |
| 51 | +// Replace each access with the one worked out for it below. |
| 52 | +Stmt rewrite_accesses(const Stmt &s, const string &alloc, |
| 53 | + const map<Expr, Expr, IRDeepCompare> &rewritten) { |
| 54 | + auto index_for = [&](const Expr &index) { |
| 55 | + auto it = rewritten.find(index); |
| 56 | + internal_assert(it != rewritten.end()); |
| 57 | + return it->second; |
| 58 | + }; |
| 59 | + return mutate_with( |
| 60 | + s, |
| 61 | + [&](auto *self, const Store *op) { |
| 62 | + Stmt s = self->visit_base(op); |
| 63 | + if (op->name == alloc) { |
| 64 | + const Store *store = s.as<Store>(); |
| 65 | + s = store->with(store->value, index_for(store->index), store->predicate, |
| 66 | + ModulusRemainder()); |
| 67 | + } |
| 68 | + return s; |
| 69 | + }, |
| 70 | + [&](auto *self, const Load *op) { |
| 71 | + Expr e = self->visit_base(op); |
| 72 | + if (op->name == alloc) { |
| 73 | + const Load *load = e.as<Load>(); |
| 74 | + e = load->with(index_for(load->index), load->predicate, ModulusRemainder()); |
| 75 | + } |
| 76 | + return e; |
| 77 | + }); |
| 78 | +} |
| 79 | + |
| 80 | +class PromoteGPURegisters : public IRMutator { |
| 81 | +protected: |
| 82 | + using IRMutator::visit; |
| 83 | + |
| 84 | + bool in_threads = false; |
| 85 | + vector<const Allocate *> pending; |
| 86 | + |
| 87 | + Stmt visit(const Allocate *op) override { |
| 88 | + LoopKinds kinds = loop_kinds(op->body); |
| 89 | + // An allocation with a loop over lanes inside it is warp-level |
| 90 | + // storage, which LowerWarpShuffles stripes across the lanes. Leave it |
| 91 | + // alone. Without a loop over threads there is nowhere to put this one, |
| 92 | + // and whoever runs it already has it to themselves. |
| 93 | + if (!in_threads && op->memory_type == MemoryType::Register && |
| 94 | + kinds.threads && !kinds.lanes) { |
| 95 | + // Pick it up, and put it back inside the loops over threads. |
| 96 | + pending.push_back(op); |
| 97 | + return mutate(op->body); |
| 98 | + } |
| 99 | + return IRMutator::visit(op); |
| 100 | + } |
| 101 | + |
| 102 | + Stmt visit(const For *op) override { |
| 103 | + if (op->for_type != ForType::GPUThread || pending.empty()) { |
| 104 | + ScopedValue<bool> bind(in_threads, |
| 105 | + in_threads || op->for_type == ForType::GPUThread || |
| 106 | + op->for_type == ForType::GPULane); |
| 107 | + return IRMutator::visit(op); |
| 108 | + } |
| 109 | + |
| 110 | + // The outermost loop over threads with allocations to place. Everything |
| 111 | + // private to a thread goes inside it. |
| 112 | + vector<const Allocate *> allocs; |
| 113 | + allocs.swap(pending); |
| 114 | + |
| 115 | + Stmt body = op->body; |
| 116 | + for (const Allocate *alloc : allocs) { |
| 117 | + body = promote(alloc, body); |
| 118 | + } |
| 119 | + { |
| 120 | + ScopedValue<bool> bind(in_threads, true); |
| 121 | + body = mutate(body); |
| 122 | + } |
| 123 | + return op->with(op->min, op->max, body); |
| 124 | + } |
| 125 | + |
| 126 | + // Give each site its own registers, and wrap the body in the smaller |
| 127 | + // allocation. |
| 128 | + Stmt promote(const Allocate *op, Stmt body) { |
| 129 | + vector<Expr> accesses = find_accesses(body, op->name); |
| 130 | + |
| 131 | + // Each access covers a set of elements, and get_subtile partitions the |
| 132 | + // accesses between the distinct sets. Nothing about the layout of a set |
| 133 | + // matters here, because the registers it gets are its own, so a dense |
| 134 | + // ramp reaches all of them. |
| 135 | + vector<MultiRamp> subtiles; |
| 136 | + map<Expr, Expr, IRDeepCompare> rewritten; |
| 137 | + string description = "the allocation " + op->name + |
| 138 | + ", which is scheduled to live in Register memory outside the " |
| 139 | + "loops over GPU threads"; |
| 140 | + for (const Expr &index : accesses) { |
| 141 | + int subtile = get_subtile(index, description, &subtiles); |
| 142 | + // Every subtile has the same shape, and so the same number of |
| 143 | + // lanes, because get_subtile rejects accesses that don't. |
| 144 | + int lanes = subtiles[subtile].total_lanes(); |
| 145 | + Expr base = make_const(index.type().element_of(), subtile * lanes); |
| 146 | + rewritten[index] = |
| 147 | + lanes == 1 ? base : Ramp::make(base, make_one(base.type()), lanes); |
| 148 | + } |
| 149 | + |
| 150 | + int size = subtiles.empty() ? 0 : (int)subtiles.size() * subtiles[0].total_lanes(); |
| 151 | + body = rewrite_accesses(body, op->name, rewritten); |
| 152 | + |
| 153 | + return op->with({make_const(Int(32), size)}, op->condition, body); |
| 154 | + } |
| 155 | +}; |
| 156 | + |
| 157 | +} // namespace |
| 158 | + |
| 159 | +Stmt promote_gpu_registers(const Stmt &s) { |
| 160 | + return PromoteGPURegisters()(s); |
| 161 | +} |
| 162 | + |
| 163 | +} // namespace Internal |
| 164 | +} // namespace Halide |
![halide-ci[bot]](https://avatars.githubusercontent.com/u/1548551?v=4&size=40){kind=avatar}
0 commit comments