Don't flag host stack allocations feeding a GPU kernel as cross-talk

src/CheckGPUCrossTalk.cpp

@@ -162,6 +162,7 @@ class CheckCrossTalk : public IRVisitor {
     using IRVisitor::visit;
 
     bool in_threads = false;
+    bool in_kernel = false;
 
     void visit(const For *op) override {
         // An allocation inside a loop over threads or lanes already belongs to
@@ -172,17 +173,20 @@ class CheckCrossTalk : public IRVisitor {
         if (op->for_type == ForType::GPUThread || op->for_type == ForType::GPULane) {
             ScopedValue<bool> bind(in_threads, true);
             IRVisitor::visit(op);
+        } else if (is_gpu(op->for_type)) {
+            ScopedValue<bool> bind(in_kernel, true);
+            IRVisitor::visit(op);
         } else {
             IRVisitor::visit(op);
         }
     }
 
     void visit(const Realize *op) override {
-        // Only memory that is private to a thread, and only when the
-        // allocation is outside the loops over threads. An allocation with an
-        // automatic memory type that lands outside them goes to shared memory,
-        // which the threads of a block really do share.
-        if (!in_threads &&
+        // Only memory that is private to a thread. An allocation with an
+        // automatic memory type that lands outside the loops over threads
+        // goes to shared memory, which the threads of a block really do
+        // share.
+        if (in_kernel && !in_threads &&
             (op->memory_type == MemoryType::Register ||
              op->memory_type == MemoryType::Stack)) {
             check(op);

test/correctness/gpu_register_at_block_level.cpp

@@ -186,6 +186,31 @@ int main(int argc, char **argv) {
         }
     }
 
+    {
+        // An allocation stored on the stack, but computed entirely on the
+        // host and outside any GPU loop at all. This is ordinary host memory
+        // read by every thread of the kernel, not a separate copy per thread,
+        // so there is no cross-talk to check for.
+        Func f("f"), g("g");
+        Var x("x"), y("y"), xi("xi"), yi("yi");
+        f(x) = x * 2;
+        f.compute_root().store_in(MemoryType::Stack);
+        g(x, y) = f(x) + f(x + 1);
+        g.gpu_tile(x, y, x, y, xi, yi, 8, 8);
+
+        Buffer<int> result = g.realize({32, 32}, target);
+        for (int y = 0; y < 32; y++) {
+            for (int x = 0; x < 32; x++) {
+                int correct = x * 2 + (x + 1) * 2;
+                if (result(x, y) != correct) {
+                    printf("host stack input: result(%d, %d) = %d instead of %d\n",
+                           x, y, result(x, y), correct);
+                    return 1;
+                }
+            }
+        }
+    }
+
     printf("Success!\n");
     return 0;
 }