Allow split store_at/compute_at over gpu threads for things stored in registers

src/CheckGPUCrossTalk.cpp

@@ -2,7 +2,9 @@
 
 #include "Bounds.h"
 #include "CanonicalizeGPUVars.h"
+#include "ExprUsesVar.h"
 #include "IR.h"
+#include "IREquality.h"
 #include "IROperator.h"
 #include "IRPrinter.h"
 #include "IRVisitor.h"
@@ -245,6 +247,50 @@ class CheckCrossTalk : public IRVisitor {
             }
         }
 
+        // Which dimensions tell one thread's part from another's. A
+        // dimension whose region is the same whatever thread is asking is one
+        // every thread walks the same way, so a coordinate a load names in it
+        // is one every thread names. Such a dimension cannot be what makes a
+        // load another thread's. It is the region that has to be asked, not
+        // the index: a loop of a thread's own is written the same way by every
+        // thread, and it is its bounds that say which part is whose.
+        // A dimension nothing could be bounded in stays in, so that failing
+        // to work out where an access reaches is still an error rather than a
+        // dimension that gets to sit the comparison out.
+        vector<bool> separates_threads(finder.accesses[0].args.size(), false);
+        for (const auto &region : regions) {
+            for (size_t i = 0; i < region.size(); i++) {
+                if (!region[i].has_lower_bound() || !region[i].has_upper_bound()) {
+                    separates_threads[i] = true;
+                    continue;
+                }
+                for (int t = 0; t < 3; t++) {
+                    const string &n = gpu_thread_name(t);
+                    separates_threads[i] =
+                        separates_threads[i] ||
+                        stmt_or_expr_uses_var(region[i].min, n) ||
+                        stmt_or_expr_uses_var(region[i].max, n);
+                }
+            }
+        }
+
+        // Does store s reach every site load l does, along the dimensions
+        // that tell one thread's part from another's?
+        const auto covers = [&](size_t s, size_t l) {
+            for (size_t i = 0; i < regions[l].size(); i++) {
+                if (!separates_threads[i]) {
+                    continue;
+                }
+                const Interval &want = regions[l][i], &have = regions[s][i];
+                if (!(want.has_lower_bound() && want.has_upper_bound() &&
+                      have.has_lower_bound() && have.has_upper_bound() &&
+                      can_prove(have.min <= want.min && want.max <= have.max))) {
+                    return false;
+                }
+            }
+            return true;
+        };
+
         for (size_t l = 0; l < finder.accesses.size(); l++) {
             const Access &load = finder.accesses[l];
             if (load.is_store) {
@@ -257,20 +303,60 @@ class CheckCrossTalk : public IRVisitor {
             // no other thread stored this one, this thread did. A site this
             // thread never wrote holds a value nothing depends on, like the
             // garbage that pads out a vector.
-            for (size_t s = 0; s < l && !ok; s++) {
+            // Stores earlier in the list have happened. A later one still
+            // counts if it lands somewhere else along a dimension that does
+            // not separate threads, because such a dimension is what carries
+            // an allocation from one run of a loop to the next: the store this
+            // load wants is the one the previous iteration ran, and the thread
+            // that ran it was this one. Where every dimension agrees there is
+            // no such gap, and a later store is simply later.
+            for (size_t s = 0; s < finder.accesses.size() && !ok; s++) {
                 const Access &store = finder.accesses[s];
                 // The store has to be in at least as many loops over threads,
                 // or it is the work of one thread standing in for all of them.
                 if (!store.is_store || store.thread_depth < load.thread_depth) {
                     continue;
                 }
-                ok = true;
-                for (size_t i = 0; i < regions[l].size() && ok; i++) {
-                    const Interval &want = regions[l][i], &have = regions[s][i];
-                    ok = (want.has_lower_bound() && want.has_upper_bound() &&
-                          have.has_lower_bound() && have.has_upper_bound() &&
-                          can_prove(have.min <= want.min && want.max <= have.max));
+                if (s > l) {
+                    bool carried = false;
+                    for (size_t i = 0; i < regions[l].size() && !carried; i++) {
+                        carried = (!separates_threads[i] &&
+                                   !equal(load.canonical_args[i], store.canonical_args[i]));
+                    }
+                    if (!carried) {
+                        continue;
+                    }
+                }
+                ok = covers(s, l);
+            }
+            // Finding one store of this thread's that covers the load is
+            // not enough. What a thread reads is what was written to the site
+            // last, so any store that might land on the same site has to be
+            // this thread's too. A store that cannot reach the site is no
+            // one's business, which is what the overlap test asks.
+            //
+            // Two ways a store that reaches it belongs to someone else: it
+            // runs in fewer loops over threads than the load, so one thread
+            // ran it on everyone's behalf; or it reaches the site from a
+            // different thread, which is what failing to cover the load along
+            // the dimensions that separate threads means.
+            for (size_t s = 0; s < finder.accesses.size() && ok; s++) {
+                const Access &store = finder.accesses[s];
+                if (!store.is_store) {
+                    continue;
+                }
+                if (store.thread_depth >= load.thread_depth && covers(s, l)) {
+                    continue;
+                }
+                bool disjoint = false;
+                for (size_t i = 0; i < regions[l].size() && !disjoint; i++) {
+                    const Interval &a = regions[l][i], &b = regions[s][i];
+                    disjoint = ((a.has_upper_bound() && b.has_lower_bound() &&
+                                 can_prove(a.max < b.min)) ||
+                                (b.has_upper_bound() && a.has_lower_bound() &&
+                                 can_prove(b.max < a.min)));
                 }
+                ok = disjoint;
             }
             if (!ok) {
                 report(op, finder.accesses, load);

src/ScheduleFunctions.cpp

@@ -2017,7 +2017,7 @@ class InjectFunctionRealization : public IRMutator {
 class ComputeLegalSchedules : public IRVisitor {
 public:
     struct Site {
-        bool is_parallel, is_gpu_block;
+        bool is_parallel, is_gpu_block, is_gpu_thread;
         LoopLevel loop_level;
     };
     vector<Site> sites_allowed;
@@ -2054,7 +2054,9 @@ class ComputeLegalSchedules : public IRVisitor {
         // thus any new ones we synthesize we must explicitly lock.
         loop_level.lock();
         const bool is_gpu_block = (f->for_type == ForType::GPUBlock);
-        sites.push_back({f->is_parallel(), is_gpu_block, loop_level});
+        const bool is_gpu_thread = (f->for_type == ForType::GPUThread ||
+                                    f->for_type == ForType::GPULane);
+        sites.push_back({f->is_parallel(), is_gpu_block, is_gpu_thread, loop_level});
 
         f->min.accept(this);
         f->max.accept(this);
@@ -2498,10 +2500,23 @@ bool validate_schedule(Function f, const Stmt &s, const Target &target, bool is_
         return store_idx >= 0 && compute_idx >= 0 && hoist_storage_idx >= 0;
     };
 
+    // Storage private to a GPU thread is not shared by a loop over threads, so
+    // such a loop between where it is stored and where it is computed is not a
+    // race - each thread gets its own copy, which is what the memory type
+    // means. Whether each thread then keeps to its own copy is a different
+    // question, and check_gpu_cross_talk answers it later in lowering, for
+    // these same two memory types.
+    const MemoryType mem = f.schedule().memory_type();
+    const bool thread_private =
+        mem == MemoryType::Register || mem == MemoryType::Stack;
+    const auto races = [&](int i) {
+        return sites[i].is_parallel && !(thread_private && sites[i].is_gpu_thread);
+    };
+
     // Check there isn't a parallel loop between the compute_at and the store_at
     if (all_ok()) {
         for (int i = store_idx + 1; i <= compute_idx; i++) {
-            if (sites[i].is_parallel) {
+            if (races(i)) {
                 err << "Func \"" << f.name()
                     << "\" is stored outside the parallel/vectorized/gpu_block loop over "
                     << sites[i].loop_level.to_string()
@@ -2514,7 +2529,7 @@ bool validate_schedule(Function f, const Stmt &s, const Target &target, bool is_
     // Check there isn't a parallel loop between the compute_at and the hoist_storage_at
     if (all_ok()) {
         for (int i = hoist_storage_idx + 1; i <= compute_idx; i++) {
-            if (sites[i].is_parallel) {
+            if (races(i)) {
                 err << "Func \"" << f.name()
                     << "\" storage is hoisted outside the parallel/vectorized/gpu_block loop over "
                     << sites[i].loop_level.to_string()

test/correctness/CMakeLists.txt

@@ -178,6 +178,7 @@ tests(
     gpu_object_lifetime_3.cpp
     gpu_param_allocation.cpp
     gpu_register_at_block_level.cpp
+    gpu_register_stored_outside_thread_loop.cpp
     gpu_reuse_shared_memory.cpp
     gpu_specialize.cpp
     gpu_store_in_register_with_no_lanes_loop.cpp