diff --git a/src/transform/thread_storage_sync.cc b/src/transform/thread_storage_sync.cc index 542fa93fef..d76c0d4e0c 100644 --- a/src/transform/thread_storage_sync.cc +++ b/src/transform/thread_storage_sync.cc @@ -1810,15 +1810,16 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { Range::FromMinExtent(loop->min, adjusted_extent)); } - // For WAW (Write-after-Write) and RAR (Read-after-Read), we should use - // the same thread variables because: - // - WAW: doesn't create true data dependency, only need to check if the - // same thread overwrites its own data across iterations - // - RAR: no dependency at all - // For RAW (Read-after-Write) and WAR (Write-after-Read), we need to use - // different thread variables to check cross-thread dependencies. - bool same_access_type = (prev.type == kWrite && curr.type == kWrite) || - (prev.type == kRead && curr.type == kRead); + // A same-iteration WAW must model two distinct threads: program order + // already orders writes from one thread, while writes from different + // threads may target the same address and require a barrier. Keep the + // existing loop-carried WAW model unchanged; correctly extending that + // case also requires shifting its constraints and placing the barrier + // inside the loop. + bool is_same_iteration_waw = + loop == nullptr && prev.type == kWrite && curr.type == kWrite; + bool use_distinct_threads = + prev.type != curr.type || is_same_iteration_waw; PrimExpr thread_condition = Bool(false); Map prev_sub, curr_sub; @@ -1828,14 +1829,7 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { Var old_prev_var = prev.threads[prev.threads.size() + idx - 3]->var; Var old_curr_var = curr.threads[curr.threads.size() + idx - 3]->var; - if (same_access_type) { - // For WAW/RAR: use a single shared Var object for both prev and curr - // This allows the analyzer to see they reference the same thread - Var shared_var(thread_names[idx], old_prev_var.dtype()); - prev_sub.Set(old_prev_var, shared_var); - curr_sub.Set(old_curr_var, shared_var); - } else { - // For RAW/WAR: use different Var objects to model cross-thread access + if (use_distinct_threads) { Var prev_var(std::string(thread_names[idx]) + "1", old_prev_var.dtype()); Var curr_var(std::string(thread_names[idx]) + "2", @@ -1844,17 +1838,20 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { tirx::Or(thread_condition, tirx::NE(prev_var, curr_var)); prev_sub.Set(old_prev_var, prev_var); curr_sub.Set(old_curr_var, curr_var); + } else { + Var shared_var(thread_names[idx], old_prev_var.dtype()); + prev_sub.Set(old_prev_var, shared_var); + curr_sub.Set(old_curr_var, shared_var); } } - if (!same_access_type) { + if (use_distinct_threads) { analyzer.EnterConstraint(thread_condition); } // Two instances in one analyzer, so per-instance binds need their own - // copy; see ConstrSet::RenameFrom. They differ when they are two threads - // (RAW/WAR) or two iterations of one thread (loop carry); a same-type - // pair within one iteration is one execution, where renaming would only - // lose the bind. - if (!same_access_type || loop != nullptr) { + // copy; see ConstrSet::RenameFrom. They differ for cross-thread proofs + // (RAW/WAR/same-iteration WAW) or for two loop iterations. Otherwise the + // pair represents one execution, where renaming would only lose the bind. + if (use_distinct_threads || loop != nullptr) { prev_cset = prev_cset.RenameFrom("", prev_sub, std::nullopt, /*rename_ranges=*/false); curr_cset = curr_cset.RenameFrom("", curr_sub, std::nullopt, diff --git a/testing/python/issue/test_tilelang_issue_thread_sync_cross_thread_waw.py b/testing/python/issue/test_tilelang_issue_thread_sync_cross_thread_waw.py new file mode 100644 index 0000000000..3fcf08ec36 --- /dev/null +++ b/testing/python/issue/test_tilelang_issue_thread_sync_cross_thread_waw.py @@ -0,0 +1,33 @@ +import torch + +import tilelang +import tilelang.language as T +import tilelang.testing + + +@tilelang.jit +def _copy_then_offset_fill( + A: T.Tensor((16,), T.float32), + B: T.Tensor((16,), T.float32), +): + with T.Kernel(1, threads=16): + a_shared = T.alloc_shared((16,), T.float32) + T.copy(A, a_shared) + T.fill(a_shared[10:16], 99.0) + T.copy(a_shared, B) + + +@tilelang.testing.requires_cuda +def test_copy_then_offset_fill_orders_cross_thread_writes(): + a = torch.arange(16, device="cuda", dtype=torch.float32) + expected = a.clone() + expected[10:16] = 99.0 + + for _ in range(10): + b = torch.empty_like(a) + _copy_then_offset_fill(a, b) + torch.testing.assert_close(b, expected, rtol=0, atol=0) + + +if __name__ == "__main__": + tilelang.testing.main() diff --git a/testing/python/transform/test_tilelang_transform_thread_sync.py b/testing/python/transform/test_tilelang_transform_thread_sync.py index aef1e65d60..c223233427 100644 --- a/testing/python/transform/test_tilelang_transform_thread_sync.py +++ b/testing/python/transform/test_tilelang_transform_thread_sync.py @@ -772,6 +772,80 @@ def func(): assert 'T.tvm_storage_sync("shared")' not in s, f"Unexpected sync:\n{s}" +def test_cross_thread_waw_with_shifted_partition_requires_sync(): + """A later writer can target a slot owned by another thread.""" + + @T.prim_func(private=True) + def func(): + s = T.alloc_buffer((128,), dtype="float32", scope="shared") + bx = T.launch_thread("blockIdx.x", 1) + tx = T.launch_thread("threadIdx.x", 64) + ty = T.launch_thread("threadIdx.y", 1) + tz = T.launch_thread("threadIdx.z", 1) + s[tx] = T.cast(tx, "float32") + if tx < 6: + s[tx + 10] = T.float32(99) + + script = run_passes_script(func) + sync = 'T.tvm_storage_sync("shared")' + assert script.count(sync) == 1, f"Expected exactly one barrier:\n{script}" + first_write = script.index('T.Cast("float32", tx)') + assert first_write < script.index(sync) < script.index("if tx < 6"), f"Barrier must separate the two write phases:\n{script}" + + +def test_provably_disjoint_waw_partitions_need_no_sync(): + """Disjoint write partitions should not acquire a conservative barrier.""" + + @T.prim_func(private=True) + def func(): + s = T.alloc_buffer((128,), dtype="float32", scope="shared") + bx = T.launch_thread("blockIdx.x", 1) + tx = T.launch_thread("threadIdx.x", 64) + ty = T.launch_thread("threadIdx.y", 1) + tz = T.launch_thread("threadIdx.z", 1) + s[tx] = T.cast(tx, "float32") + s[tx + 64] = T.float32(99) + + script = run_passes_script(func) + assert 'T.tvm_storage_sync("shared")' not in script, f"Disjoint partitions need no barrier:\n{script}" + + +def test_thread_private_waw_needs_no_sync(): + """Two writes to the same thread-private slot remain program ordered.""" + + @T.prim_func(private=True) + def func(): + s = T.alloc_buffer((64,), dtype="float32", scope="shared") + bx = T.launch_thread("blockIdx.x", 1) + tx = T.launch_thread("threadIdx.x", 64) + ty = T.launch_thread("threadIdx.y", 1) + tz = T.launch_thread("threadIdx.z", 1) + s[tx] = T.cast(tx, "float32") + s[tx] = T.float32(99) + + script = run_passes_script(func) + assert 'T.tvm_storage_sync("shared")' not in script, f"Thread-private writes need no barrier:\n{script}" + + +def test_waw_bind_definition_uses_per_thread_instances(): + """Bind definitions must not collapse the two WAW thread instances.""" + + @T.prim_func(private=True) + def func(): + s = T.alloc_buffer((128,), dtype="float32", scope="shared") + bx = T.launch_thread("blockIdx.x", 1) + tx = T.launch_thread("threadIdx.x", 64) + ty = T.launch_thread("threadIdx.y", 1) + tz = T.launch_thread("threadIdx.z", 1) + idx: T.int32 = tx + 10 + s[idx - 10] = T.cast(tx, "float32") + if tx < 6: + s[idx] = T.float32(99) + + script = run_passes_script(func) + assert 'T.tvm_storage_sync("shared")' in script, f"Bind sharing hid a cross-thread WAW hazard:\n{script}" + + @tilelang.testing.requires_cuda def test_partial_sync_non_warp_multiple_rejected(): """Regression test for issue #2556: a required barrier inside a divergent