diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index a444273c9a46..c4b19f5a4ba4 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -424,6 +424,23 @@ class ExtractSharedAndHeapAllocations : public IRMutator { public: vector allocations; + /** Which allocations end up sharing one piece of memory. Two Funcs in the + * same group are the same bytes at different times, so a hazard between + * them is a hazard on that memory even though the names differ, and + * whoever places thread barriers has to see them as one thing. The groups + * are decided by liveness measured here, before any barriers are placed, + * so asking now gives the same answer as rewrapping later will. */ + std::map storage_groups() { + vector groups = allocate_funcs(allocations); + std::map result; + for (int i = 0; i < (int)groups.size(); i++) { + for (const SharedAllocation &a : groups[i].group) { + result[a.name] = i; + } + } + return result; + } + protected: map shared; @@ -796,7 +813,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator { // Given some allocations, return a vector of allocation group where each group // consists of a number of allocations which should be coalesced together // in the shared memory. - vector allocate_funcs(vector &allocations) { + vector allocate_funcs(vector allocations) { // Sort based on the ascending order of the min liveness stage, // then sort based on the ascending order of the max liveness stage. sort(allocations.begin(), allocations.end(), @@ -1428,17 +1445,21 @@ class ExtractRegisterAllocations : public IRMutator { class InjectThreadBarriers : public IRMutator { protected: - bool in_threads = false, injected_barrier; + bool in_threads = false, injected_barrier = false; using IRMutator::visit; const ExtractSharedAndHeapAllocations &block_allocs; const ExtractRegisterAllocations ®ister_allocs; - std::set shared_stores; - std::set device_stores; - std::set shared_loads; - std::set device_loads; + // Names that share memory answer to the same key, so a hazard between two + // Funcs the allocator coalesced is not missed. + std::map storage_group; + + std::string storage_key(const std::string &name) { + auto it = storage_group.find(name); + return it == storage_group.end() ? name : "group " + std::to_string(it->second); + } MemoryType memory_type_for_name(const std::string &name) { for (const auto &x : register_allocs.allocations) { @@ -1486,99 +1507,138 @@ class InjectThreadBarriers : public IRMutator { } } - Stmt visit(const Store *op) override { - debug(4) << "Encountered store to " << op->name << "\n"; - auto mem_type = memory_type_for_name(op->name); - switch (mem_type) { - case MemoryType::GPUSharedAsync: - case MemoryType::GPUShared: - debug(4) << " memory type is shared\n"; - shared_stores.insert(op->name); - break; - case MemoryType::Auto: - case MemoryType::Heap: - case MemoryType::GPUTexture: - debug(4) << " memory type is heap or auto\n"; - device_stores.insert(op->name); - break; - case MemoryType::Stack: - case MemoryType::Register: - case MemoryType::LockedCache: - case MemoryType::VTCM: - case MemoryType::AMXTile: - break; + // What a statement touches in the memory the threads of a block share. + struct Footprint { + std::set shared_stores, shared_loads; + std::set device_stores, device_loads; + + static bool intersects(const std::set &a, + const std::set &b) { + for (const auto &x : a) { + if (b.count(x)) { + return true; + } + } + return false; } - return IRMutator::visit(op); - } + // Which memory spaces this statement and everything before it since + // the last barrier disagree about: it reads what was written, writes + // what was read, or writes what was written. Any of those needs the + // threads brought back together first. + int conflict(const Footprint &earlier) const { + int mask = 0; + if (intersects(shared_loads, earlier.shared_stores) || + intersects(shared_stores, earlier.shared_loads) || + intersects(shared_stores, earlier.shared_stores)) { + mask |= CodeGen_GPU_Dev::MemoryFenceType::Shared; + } + if (intersects(device_loads, earlier.device_stores) || + intersects(device_stores, earlier.device_loads) || + intersects(device_stores, earlier.device_stores)) { + mask |= CodeGen_GPU_Dev::MemoryFenceType::Device; + } + return mask; + } - Expr visit(const Load *op) override { - debug(4) << "Encountered load from " << op->name << "\n"; - auto mem_type = memory_type_for_name(op->name); - switch (mem_type) { + // Forget the accesses a barrier has now ordered. A barrier only + // fences the spaces named in its mask, so accesses in a space it did + // not fence are still waiting for one. + void clear(int mask) { + if (mask & CodeGen_GPU_Dev::MemoryFenceType::Shared) { + shared_stores.clear(); + shared_loads.clear(); + } + if (mask & CodeGen_GPU_Dev::MemoryFenceType::Device) { + device_stores.clear(); + device_loads.clear(); + } + } + + void add(const Footprint &other) { + shared_stores.insert(other.shared_stores.begin(), other.shared_stores.end()); + shared_loads.insert(other.shared_loads.begin(), other.shared_loads.end()); + device_stores.insert(other.device_stores.begin(), other.device_stores.end()); + device_loads.insert(other.device_loads.begin(), other.device_loads.end()); + } + }; + + void record(Footprint &f, const std::string &raw_name, bool is_store) { + const std::string name = storage_key(raw_name); + switch (memory_type_for_name(raw_name)) { case MemoryType::GPUSharedAsync: case MemoryType::GPUShared: - debug(4) << " memory type is shared\n"; - shared_loads.insert(op->name); + (is_store ? f.shared_stores : f.shared_loads).insert(name); break; case MemoryType::Auto: case MemoryType::Heap: case MemoryType::GPUTexture: - debug(4) << " memory type is heap or auto\n"; - device_loads.insert(op->name); + (is_store ? f.device_stores : f.device_loads).insert(name); break; - case MemoryType::Stack: - case MemoryType::Register: - case MemoryType::LockedCache: - case MemoryType::VTCM: - case MemoryType::AMXTile: + default: break; } + } - return IRMutator::visit(op); + Footprint footprint_of(const Stmt &s) { + Footprint f; + visit_with( + s, + [&](auto *self, const Store *op) { + record(f, op->name, true); + self->visit_base(op); + }, + [&](auto *self, const Load *op) { + record(f, op->name, false); + self->visit_base(op); + }, + [&](auto *self, const Call *op) { + // A texture is not loaded and stored but passed through these, + // which name it in their first argument. + if (op->is_intrinsic(Call::image_load) || + op->is_intrinsic(Call::image_store)) { + const StringImm *name = op->args[0].as(); + internal_assert(name) << "Malformed " << op->name << "\n"; + record(f, name->value, op->is_intrinsic(Call::image_store)); + } + self->visit_base(op); + }); + return f; } + // Taking the statements of a block a pair at a time puts a barrier at every + // join, whether or not the two halves disagree about any memory. Take the + // whole sequence at once and put barriers only where a statement meets + // something an earlier one left. Stmt visit(const Block *op) override { - if (!in_threads && op->rest.defined()) { - // First, we record which loads from shared/device memory occur - // in the rest block - Stmt rest = mutate(op->rest); + if (in_threads || !op->rest.defined()) { + return IRMutator::visit(op); + } - // Now, record which stores occur in the first stmt - // of this block - shared_stores.clear(); - device_stores.clear(); - Stmt first = mutate(op->first); + std::vector stmts = Block::to_vector(op); - // If there are any loads in the rest part that - // load from something stored in first, insert the appropriate - // fence type - int mask = 0; - for (const auto &st : shared_stores) { - auto elem = shared_loads.find(st); - if (elem != shared_loads.end()) { - mask |= CodeGen_GPU_Dev::MemoryFenceType::Shared; - break; - } - } - for (const auto &st : device_stores) { - auto elem = device_loads.find(st); - if (elem != device_loads.end()) { - mask |= CodeGen_GPU_Dev::MemoryFenceType::Device; - break; - } + std::vector result; + Footprint pending; + for (const Stmt &s : stmts) { + Stmt mutated = mutate(s); + Footprint here = footprint_of(mutated); + int mask = here.conflict(pending); + if (mask) { + result.push_back(make_barrier(mask)); + injected_barrier = true; + pending.clear(mask); } - injected_barrier = true; - return Block::make({first, make_barrier(mask), rest}); - } else { - return IRMutator::visit(op); + pending.add(here); + result.push_back(mutated); } + return Block::make(result); } public: InjectThreadBarriers(ExtractSharedAndHeapAllocations &sha, ExtractRegisterAllocations &ra) : block_allocs(sha), - register_allocs(ra) { + register_allocs(ra), + storage_group(sha.storage_groups()) { } };