[BugFix][Transform] Detect cross-thread WAW hazards

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<Var, PrimExpr> 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>", prev_sub, std::nullopt,
                                          /*rename_ranges=*/false);
         curr_cset = curr_cset.RenameFrom("<CURR>", curr_sub, std::nullopt,

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()

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