[FEAT] Race Detector

examples/race_detector/inplace_neighbor.py

@@ -0,0 +1,36 @@
+"""RW race: Block N reads x[N+1 region] while Block N+1 writes x[N+1 region]."""
+
+import torch
+import triton
+import triton.language as tl
+
+import triton_viz
+from triton_viz.clients import RaceDetector
+
+
+@triton_viz.trace(RaceDetector())
+@triton.jit
+def inplace_neighbor_kernel(x_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+    pid = tl.program_id(0)
+    own = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    own_mask = own < n_elements
+    own_data = tl.load(x_ptr + own, mask=own_mask, other=0.0)
+
+    neighbor = (pid + 1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    neighbor_mask = neighbor < n_elements
+    neighbor_data = tl.load(x_ptr + neighbor, mask=neighbor_mask, other=0.0)
+
+    tl.store(x_ptr + own, own_data + neighbor_data * 0.5, mask=own_mask)
+
+
+if __name__ == "__main__":
+    from triton_viz.core.trace import launches
+
+    n, bs = 32, 8
+    x = torch.randn(n, dtype=torch.float32)
+    inplace_neighbor_kernel[(triton.cdiv(n, bs),)](x, n, bs)
+
+    races = launches[-1].records
+    print(f"Detected {len(races)} race(s)")
+    for r in races:
+        print(f"  {r.race_type.name} at address offset {r.address_offset}")

examples/race_detector/reduction.py

@@ -0,0 +1,33 @@
+"""WAW race: All blocks write to the same scalar address without atomics."""
+
+import torch
+import triton
+import triton.language as tl
+
+import triton_viz
+from triton_viz.clients import RaceDetector
+
+
+@triton_viz.trace(RaceDetector())
+@triton.jit
+def reduction_kernel(input_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    values = tl.load(input_ptr + offsets, mask=mask, other=0.0)
+    block_sum = tl.sum(values)
+    tl.store(output_ptr, block_sum)  # BUG: all blocks -> same addr
+
+
+if __name__ == "__main__":
+    from triton_viz.core.trace import launches
+
+    n, bs = 64, 8
+    inp = torch.randn(n, dtype=torch.float32)
+    out = torch.zeros(1, dtype=torch.float32)
+    reduction_kernel[(triton.cdiv(n, bs),)](inp, out, n, bs)
+
+    races = launches[-1].records
+    print(f"Detected {len(races)} race(s)")
+    for r in races:
+        print(f"  {r.race_type.name} at address offset {r.address_offset}")

examples/race_detector/scatter.py

@@ -0,0 +1,35 @@
+"""WAW race: Different blocks scatter to overlapping destination indices."""
+
+import torch
+import triton
+import triton.language as tl
+
+import triton_viz
+from triton_viz.clients import RaceDetector
+
+
+@triton_viz.trace(RaceDetector())
+@triton.jit
+def scatter_kernel(src_ptr, idx_ptr, dst_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    values = tl.load(src_ptr + offsets, mask=mask, other=0.0)
+    indices = tl.load(idx_ptr + offsets, mask=mask, other=0)
+    tl.store(dst_ptr + indices, values, mask=mask)
+
+
+if __name__ == "__main__":
+    from triton_viz.core.trace import launches
+
+    n, bs = 32, 8
+    src = torch.randn(n, dtype=torch.float32)
+    # Create conflicts: all indices point to [0, 3]
+    idx = torch.randint(0, 4, (n,), dtype=torch.int32)
+    dst = torch.zeros(4, dtype=torch.float32)
+    scatter_kernel[(triton.cdiv(n, bs),)](src, idx, dst, n, bs)
+
+    races = launches[-1].records
+    print(f"Detected {len(races)} race(s)")
+    for r in races:
+        print(f"  {r.race_type.name} at address offset {r.address_offset}")

examples/race_detector/transpose.py

@@ -0,0 +1,36 @@
+"""RW+WAW race: In-place matrix transpose where blocks read/write overlapping addresses."""
+
+import torch
+import triton
+import triton.language as tl
+
+import triton_viz
+from triton_viz.clients import RaceDetector
+
+
+@triton_viz.trace(RaceDetector())
+@triton.jit
+def transpose_kernel(matrix_ptr, N, BLOCK: tl.constexpr):
+    # Each block handles one row: reads row pid, writes column pid
+    pid = tl.program_id(0)
+    cols = tl.arange(0, BLOCK)
+    read_mask = cols < N
+    # Read row pid: matrix[pid, :]
+    read_off = pid * N + cols
+    vals = tl.load(matrix_ptr + read_off, mask=read_mask, other=0.0)
+    # Write column pid: matrix[:, pid]  (i.e. matrix[col, pid] for each col)
+    write_off = cols * N + pid
+    tl.store(matrix_ptr + write_off, vals, mask=read_mask)
+
+
+if __name__ == "__main__":
+    from triton_viz.core.trace import launches
+
+    N, block = 8, 8
+    mat = torch.randn(N, N, dtype=torch.float32)
+    transpose_kernel[(N,)](mat, N, block)
+
+    races = launches[-1].records
+    print(f"Detected {len(races)} race(s)")
+    for r in races:
+        print(f"  {r.race_type.name} at address offset {r.address_offset}")

pyproject.toml

@@ -50,6 +50,7 @@ homepage = "https://github.com/Deep-Learning-Profiling-Tools/triton-viz"
 [project.scripts]
 triton-sanitizer = "triton_viz.wrapper:apply_sanitizer"
 triton-profiler = "triton_viz.wrapper:apply_profiler"
+triton-race-detector = "triton_viz.wrapper:apply_race_detector"
 triton-visualizer = "triton_viz.visualizer_cli:main"
 
 [project.optional-dependencies]