diff --git a/src/CheckGPUCrossTalk.cpp b/src/CheckGPUCrossTalk.cpp index 73f2a4ecfb6f..0ae58b660713 100644 --- a/src/CheckGPUCrossTalk.cpp +++ b/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 separates_threads(finder.accesses[0].args.size(), false); + for (const auto ®ion : 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); diff --git a/src/ScheduleFunctions.cpp b/src/ScheduleFunctions.cpp index f2a31b1b6562..869148aaf86a 100644 --- a/src/ScheduleFunctions.cpp +++ b/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 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() diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index d88c13fa177f..328937c11c39 100644 --- a/test/correctness/CMakeLists.txt +++ b/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 diff --git a/test/correctness/gpu_register_stored_outside_thread_loop.cpp b/test/correctness/gpu_register_stored_outside_thread_loop.cpp new file mode 100644 index 000000000000..b83d626bb7f1 --- /dev/null +++ b/test/correctness/gpu_register_stored_outside_thread_loop.cpp @@ -0,0 +1,214 @@ +// Register memory is private to a GPU thread, so a loop over threads between +// where a Func is stored and where it is computed is not a race: each thread +// gets its own copy rather than sharing one. +// +// The point of such a placement is not the loop over threads, which nothing +// can slide or fold over. It is that a serial loop can sit in there too. A +// group of warps cooperates on a walk over a serial reduction, and a producer +// feeding it slides over that walk while staying computed inside the warps, so +// its values never have to leave registers to cross between them. That needs +// the storage above the walk, which puts it above the loop over warps as well. +// +// What makes it safe is that each warp only ever touches the rows of its own +// copy that it wrote, and that is check_gpu_cross_talk's question rather than +// the schedule validator's. So the two halves below are a pair: a placement +// that is fine and used to be refused up front, and one that is not fine, gets +// past the validator, and has to be caught by the cross-talk check instead. + +#include "Halide.h" +#include "expect_user_error.h" +#include + +using namespace Halide; + +namespace { + +// The width of the accumulator, how many rows a warp takes, how many warps +// share a block, how many blocks there are, and how many steps the walk takes. +const int W = 8, rows_per_warp = 16, WARPS = 4, BLOCKS = 2, N = 16; +const int Y = rows_per_warp * WARPS * BLOCKS; + +// A producer stored above a serial walk, computed within the loop over warps +// inside it, and slid over the walk. Each warp writes and reads only its own +// rows, and only the two steps of the walk that are live. +int slid_over_a_walk_inside_the_warps(const Target &target, MemoryType mem) { + Func p("p"), acc("acc"), out("out"); + Var x("x"), y("y"), t("t"), yo("yo"), yw("yw"), yi("yi"); + RDom rt(0, N, "rt"); + + p(y, t) = cast((y + 1) * (t + 1)); + + // Each step needs two consecutive values of p, so p slides over the walk. + acc(x, y) = 0.f; + acc(x, y) += (p(y, rt) - p(y, rt - 1)) * cast(x + 1); + + out(x, y) = acc(x, y); + + out.bound(x, 0, W).bound(y, 0, Y).compute_root(); + // A block holds the warps that share the walk, and each of them a stripe + // of rows, so that the loop over warps below has all of them in it. + out.split(y, yo, yw, rows_per_warp * WARPS) + .split(yw, yw, yi, rows_per_warp) + .gpu_blocks(yo) + .gpu_threads(yw); + + acc.compute_at(out, yo) + .split(y, yw, yi, rows_per_warp) + .gpu_threads(yw) + .unroll(yi); + // The walk is outside the loop over warps, so the warps run it together. + // The rows a warp owns are unrolled, because storage a thread indexes with + // a loop variable cannot be registers. + acc.update() + .split(y, yw, yi, rows_per_warp) + .reorder(x, yi, yw, rt) + .gpu_threads(yw) + .unroll(yi); + + // Stored above the walk - and so above the loop over warps - computed + // within the warps, and folded down to the two steps that are live. + // + // Everything that indexes it is unrolled, here and above, because a + // register is only a register if which one it is, is known: the rows a + // warp owns on both sides, and the two live steps it rotates between. + p.store_at(out, yo) + .compute_at(acc, yw) + .store_in(mem) + .unroll(y) + .unroll(t); + + Buffer result = out.realize({W, Y}, target); + result.copy_to_host(); + + // p(y, rt) - p(y, rt - 1) is y + 1 at every step, so each row accumulates + // that times the column, once per step. A warp that read another warp's + // rows would get a different row's answer. + for (int y = 0; y < Y; y++) { + for (int x = 0; x < W; x++) { + float correct = (float)N * (float)(y + 1) * (float)(x + 1); + if (result(x, y) != correct) { + printf("result(%d, %d) = %f instead of %f (memory type %d)\n", + x, y, (double)result(x, y), (double)correct, (int)mem); + return 1; + } + } + } + return 0; +} + +#if HALIDE_WITH_EXCEPTIONS + +// Reaching the error only needs compiling, so this half runs without a device. +Target compile_only_target() { + return get_host_target() + .with_feature(Target::CUDA) + .with_feature(Target::CUDACapability80); +} + +// Stored at the block level, computed within the loop over threads in y, and +// spread across the threads in x by its own gpu_threads. No thread holds the +// whole of f, and the value each one wants from its neighbour is not in the +// copy it has. +void reads_another_threads_copy() { + Func f("f"), g("g"); + Var x("x"), y("y"), xo("xo"), yo("yo"), xi("xi"), yi("yi"); + + f(x, y) = x + y * 1000; + g(x, y) = f(x, y) + f(x - 1, y); + + g.gpu_tile(x, y, xo, yo, xi, yi, 16, 16); + f.store_at(g, xo) + .compute_at(g, yi) + .store_in(MemoryType::Register) + .gpu_threads(x); + + g.compile_jit(compile_only_target()); +} + +// Finding one store of this thread's that covers a load is not enough: what a +// thread reads is whatever was written to the site last. Here the c == 1 slice +// is written by a stage with no loops over threads of its own, so one thread +// writes it on everyone's behalf, while the c == 0 slice is written by all of +// them. c is a constant in every access, so it cannot be what tells one +// thread's part from another's, and the load of c == 1 must not be excused by +// the store to c == 0. +void one_thread_stands_in_for_all_in_a_slice() { + Func f("f"), g("g"); + Var x("x"), y("y"), c("c"), xi("xi"), yi("yi"); + + f(x, y, c) = undef(); + f(x, y, 0) = x + y * 1000; + f(x, y, 1) = x + y * 1000 + 7; + g(x, y) = f(x, y, 0) + f(x, y, 1); + + g.gpu_tile(x, y, x, y, xi, yi, 16, 16); + f.compute_at(g, x).store_in(MemoryType::Register).bound(c, 0, 2); + f.update(0).gpu_threads(x, y); + // f.update(1) is left serial, so one thread runs all of it. + + g.compile_jit(compile_only_target()); +} + +// The same, with the stage that disagrees running in as many loops over +// threads as the load but mapping them the other way round, so the site a +// thread reads in the c == 1 slice was written by the thread with its +// coordinates transposed. +void stages_disagree_within_a_slice() { + Func f("f"), g("g"); + Var x("x"), y("y"), c("c"), xi("xi"), yi("yi"); + + f(x, y, c) = undef(); + f(x, y, 0) = x + y * 1000; + f(x, y, 1) = x + y * 1000 + 7; + g(x, y) = f(x, y, 0) + f(x, y, 1); + + g.gpu_tile(x, y, x, y, xi, yi, 16, 16); + f.compute_at(g, x).store_in(MemoryType::Register).bound(c, 0, 2); + f.update(0).gpu_threads(x, y); + f.update(1).reorder(y, x).gpu_threads(y, x); + + g.compile_jit(compile_only_target()); +} + +#endif // HALIDE_WITH_EXCEPTIONS + +} // namespace + +int main(int argc, char **argv) { +#if HALIDE_WITH_EXCEPTIONS + int failures = 0; + failures += !expect_user_error("reads_another_threads_copy", + "keeps to its own part", + reads_another_threads_copy); + failures += !expect_user_error("one_thread_stands_in_for_all_in_a_slice", + "keeps to its own part", + one_thread_stands_in_for_all_in_a_slice); + failures += !expect_user_error("stages_disagree_within_a_slice", + "keeps to its own part", + stages_disagree_within_a_slice); + if (failures != 0) { + printf("%d schedule(s) did not produce the expected user error\n", failures); + return 1; + } +#else + printf("[SKIP] Halide was compiled without exceptions, so the schedule " + "that must be rejected is not exercised.\n"); +#endif + + Target target = get_jit_target_from_environment(); + if (!target.has_gpu_feature()) { + printf("[SKIP] No GPU target enabled, so the schedule that must be " + "accepted is not exercised.\n"); + return 0; + } + // Both memory types are private to a thread, and check_gpu_cross_talk + // checks both, so the schedule is allowed for both. + for (MemoryType mem : {MemoryType::Register, MemoryType::Stack}) { + if (slid_over_a_walk_inside_the_warps(target, mem) != 0) { + return 1; + } + } + + printf("Success!\n"); + return 0; +}