diff --git a/src/CheckGPUCrossTalk.cpp b/src/CheckGPUCrossTalk.cpp index 73f2a4ecfb6f..6d2b09c25d34 100644 --- a/src/CheckGPUCrossTalk.cpp +++ b/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 bind(in_threads, true); IRVisitor::visit(op); + } else if (is_gpu(op->for_type)) { + ScopedValue 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); diff --git a/test/correctness/gpu_register_at_block_level.cpp b/test/correctness/gpu_register_at_block_level.cpp index 84ba03b13305..19db9322fc81 100644 --- a/test/correctness/gpu_register_at_block_level.cpp +++ b/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 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; }