Only inject thread barriers where threads disagree about memory

src/FuseGPUThreadLoops.cpp

@@ -424,6 +424,23 @@ class ExtractSharedAndHeapAllocations : public IRMutator {
 public:
     vector<SharedAllocation> 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<std::string, int> storage_groups() {
+        vector<AllocGroup> groups = allocate_funcs(allocations);
+        std::map<std::string, int> 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<string, SharedAllocation *> 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<AllocGroup> allocate_funcs(vector<SharedAllocation> &allocations) {
+    vector<AllocGroup> allocate_funcs(vector<SharedAllocation> 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 &register_allocs;
 
-    std::set<std::string> shared_stores;
-    std::set<std::string> device_stores;
-    std::set<std::string> shared_loads;
-    std::set<std::string> 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<std::string, int> 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<std::string> shared_stores, shared_loads;
+        std::set<std::string> device_stores, device_loads;
+
+        static bool intersects(const std::set<std::string> &a,
+                               const std::set<std::string> &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<StringImm>();
+                    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<Stmt> 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<Stmt> 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()) {
     }
 };