From 6956ea57b73e8dbf7195b036e1260644da2588d0 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 25 Aug 2026 09:58:09 -0700 Subject: [PATCH 1/2] Only inject thread barriers where threads disagree about memory InjectThreadBarriers visited Blocks pairwise, so it put a barrier at every join in a sequence of statements regardless of whether the two halves touched any of the same memory. Take the whole sequence at once and compare each statement's footprint against everything since the last barrier, emitting one only for a genuine read-after-write, write-after-read, or write-after-write. A barrier only fences the memory spaces named in its mask, so it only clears the footprint of those spaces. Accesses in a space it did not fence are still waiting for a barrier of their own. The footprint is keyed on the storage a Func ends up in rather than its name. ExtractSharedAndHeapAllocations pools allocations with disjoint lifetimes onto the same memory, so two differently named Funcs can be the same bytes at different times, and the hazard between them is real. allocate_funcs takes its argument by value, because the grouping is now asked for twice and sorting the caller's vector would make the second answer depend on the first. Textures are counted too. They are reached through image_load and image_store rather than Load and Store, so walking for the latter alone would miss a Func stored in MemoryType::GPUTexture entirely. Across the correctness tests this drops the number of barriers emitted from 259 to 181. Co-Authored-By: Claude Opus 5 --- src/FuseGPUThreadLoops.cpp | 209 ++++++++++++++++++++++++------------- 1 file changed, 135 insertions(+), 74 deletions(-) diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index a444273c9a46..a15bad37a6ea 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,22 @@ 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 +1508,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()) { } }; From 1067c093299224c7565bac7e7a18cdb2d7b18ef7 Mon Sep 17 00:00:00 2001 From: "halide-ci[bot]" <266445882+halide-ci[bot]@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:49:30 +0000 Subject: [PATCH 2/2] Apply pre-commit auto-fixes --- src/FuseGPUThreadLoops.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index a15bad37a6ea..c4b19f5a4ba4 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -1461,7 +1461,6 @@ class InjectThreadBarriers : public IRMutator { 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) { if (x.name == name) {