diff --git a/README.md b/README.md index 9aff6f85..721a02d1 100644 --- a/README.md +++ b/README.md @@ -282,6 +282,36 @@ def kernel(x_ptr, n, BLOCK: tl.constexpr): - Requires Triton 3.8; on another release every launch is `unsupported` (`host-compile-unavailable`), nothing compiled, and the program goes on. +### Compiled race detector + +`RaceDetector(compile=True)` (or `tile-race --compile my_script.py`) decides +whether two program instances of a launch can race on global memory, from the +host-compiled TTIR, like the compiled sanitizer: no GPU is needed, kernels are +not run, every autotune config is analyzed, and a config that fails to compile +for the target is `unsupported` (`compile-failed`) while the program goes on. + +- Each launch gets a `RaceVerdict` in its records, after a `CompiledRaceRecord` + per race (printed too). `ok` is a proof; its `scope` says what it holds for: + any scalar arguments and any grid (`any-params-any-grid`), this launch's + arguments and any grid, or this launch's arguments and grid. Its + `qualifiers` say what it assumes or read: `assumes no i32 overflow` (address + or mask arithmetic that could wrap somewhere in that scope, though not at + this launch: one that can wrap here is `unsupported`, `integer-overflow`), + `assumes termination` (the kernel spins on a flag), `reads tensor contents` + (the proof read an index table, so it holds for those contents). "Any grid" + is any extent along the grid axes the kernel reads. +- `unsupported` names what was not modeled (e.g. a `while` loop that is not a + spin on a flag, `spin-shape`; a CAS operand the integer model cannot follow, + `cas-value`), a Z3 timeout, or races found only under over-approximation + (`potential-race`). +- A loaded integer index is read from the tensor's contents at the launch + (integer, contiguous, at most 16384 elements); since kernels are not run, a + later launch sees the contents the script left, not an earlier kernel's + outputs (nor what an autotune `pre_hook` would have written). +- Races are always printed; the parts of a launch that were not decided, and + a race-free launch's scope, are printed under `TILELENS_VERBOSE` or + `tile-race`. + ### Save and load traces ```py diff --git a/tests/end_to_end/test_compiled_race_detector.py b/tests/end_to_end/test_compiled_race_detector.py new file mode 100644 index 00000000..8c3bd4dd --- /dev/null +++ b/tests/end_to_end/test_compiled_race_detector.py @@ -0,0 +1,579 @@ +"""RaceDetector(compile=True) on real launches: each launch is compiled on the +host for the client's target, its TTIR analyzed for global-memory races and +the kernel not run, on CPU tensors with Triton's driver unreachable. The +kernels are tests/unit/race_detector/compiled_kernels.py's, whose analysis +tests/unit/race_detector checks in detail.""" + +from __future__ import annotations + +import importlib +import os +import subprocess +import sys +import threading +from pathlib import Path +from types import MappingProxyType + +import pytest +import torch +import triton +import triton.language as tl + +import tilelens +from tilelens.clients import ( + CompiledRaceDetector, + CompiledRaceRecord, + RaceDetector, + RaceVerdict, +) +from tilelens.clients.race_detector.compiled import client as client_module +from tilelens.clients.race_detector.compiled.leaves import LaunchTensor +from tilelens.clients.race_detector.hb_common import UnsupportedSymbolicRaceQuery +from tilelens.clients.race_detector.two_copy_symbolic_hb_solver import ( + TwoCopySymbolicHBSolver, +) +from tilelens.core.config import DEFAULT_IR_TARGET +from tilelens.core.config import config as cfg +from tilelens.ir.launch import LaunchBinding, tensor_facts +from tilelens.ir.ttir_reader import parse_ttir + +# tests/ is on sys.path (tests/conftest.py's rootdir under pytest's default +# prepend import mode): the kernels and the TTIR injector are the unit +# tests' own. +from unit.ir.inject import inject_ttir +from unit.race_detector import compiled_kernels as K + +trace_module = importlib.import_module("tilelens.core.trace") + +REPO = Path(__file__).resolve().parents[2] +B = K.B + + +@pytest.fixture(autouse=True) +def _setup(unreachable_driver, monkeypatch): + unreachable_driver("IR mode queried Triton's driver") + for name in ("TILELENS_IR_TARGET", "TRITON_VIZ_IR_TARGET", "TRITON_INTERPRET"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(cfg, "ir_target", DEFAULT_IR_TARGET) + monkeypatch.setattr(cfg, "enable_race_detector", True) + from triton import knobs + + # A traced launch elsewhere may have left interpret mode pinned on. + monkeypatch.setitem(knobs.runtime.__dict__, "interpret", False) + + +def i32(n): + return torch.zeros(n, dtype=torch.int32) + + +def f32(n): + return torch.zeros(n, dtype=torch.float32) + + +def launch(kernel, grid, detector=None, **args): + """Launch ``kernel`` (a compiled_kernels name or a JITFunction) traced by + a compiled race detector with ``args`` (and, for a compiled_kernels + name, its SPECS constexprs); returns the detector.""" + if isinstance(kernel, str): + args = {**K.SPECS[kernel][1], **args} + kernel = getattr(K, kernel) + det = detector or RaceDetector(compile=True) + tilelens.trace(det)(kernel)[grid](**args) + return det + + +def races(det): + return { + (r.race_type, r.first.tensor_name, r.programs[0][0], r.programs[1][0]) + for r in det.last_reports + } + + +@pytest.mark.parametrize( + "name, race", + [ + ("waw", ("WAW", "out")), + ("raw", ("RAW", "x")), + ("war", ("WAR", "x")), + ("atomic_vs_load", ("RAW", "x")), + ], +) +def test_each_race_class_is_recorded_and_printed(name, race, capsys): + tensors = {"x": i32(B), "out": i32(4 * B)} + det = launch( + name, (2,), **{k: v for k, v in tensors.items() if k in K.SPECS[name][0]} + ) + verdict = det.last_verdict + assert isinstance(verdict, RaceVerdict) and verdict.status == "races" + assert races(det) == {(*race, 0, 1)} + records = trace_module.launches[-1].records + assert records[:-1] == det.last_reports and records[-1] is verdict + (config,) = verdict.per_config + assert (config.status, config.n_reports) == ("races", len(det.last_reports)) + record = det.last_reports[0] + assert isinstance(record, CompiledRaceRecord) + assert record.first.source_location[0] == K.__file__ + assert record.first.element is not None + out = capsys.readouterr().out + assert out.count("Data Race Detected") == len(det.last_reports) + assert "(compiled race detector)" in out and f"Race Type: {race[0]}" in out + + +I32 = "assumes no i32 overflow" +TERM = "assumes termination" +CONTENT = "reads tensor contents" + + +def _perm(): + return torch.randperm(4 * B, generator=torch.Generator().manual_seed(0)).to( + torch.int32 + ) + + +# (kernel, grid, args, status, rung, qualifiers, refusal kind) +ROWS = [ + # a proof per rung + ("add", 8, lambda: dict(x=f32(1024), y=f32(1024), out=f32(1024), n=1000), + "ok", "T0", (I32,), None), + ("atomic_only", 4, lambda: dict(x=i32(B)), "ok", "T0", (), None), + ("strided", 4, lambda: dict(x=f32(64 * B), out=f32(64 * B), stride=B), + "ok", "T1", (I32,), None), + ("two_slots", 2, lambda: dict(out=i32(2 * B)), "ok", "T1-launch", (), None), + ("scatter", 4, lambda: dict(idx=_perm(), out=i32(4 * B)), + "ok", "T1", (I32, CONTENT), None), + # fences: a global barrier orders; a conditional one does not; in a loop + # one barrier orders the same iteration only + ("store_fence_load", 1, lambda: dict(hist=i32(B), out=i32(B)), + "ok", "T1-launch", (), None), + ("store_flag_fence_load", 1, lambda: dict(hist=i32(B), out=i32(B), flag=3), + "races", None, (), None), + ("loop_rows_fenced", 1, lambda: dict(t=i32(4 * B), u=i32(4 * B), n=4), + "ok", "T0", (I32,), None), + ("loop_buffer_fenced", 1, lambda: dict(t=i32(4 * B), u=i32(4 * B), n=4), + "races", None, (), None), + # dependency order + ("inplace", 4, lambda: dict(x=i32(4 * B)), "ok", "T0", (I32,), None), + ("inplace_shift", 4, lambda: dict(x=i32(4 * B + 1)), "races", None, (), None), + # awaits + ("comm_comp", 3, lambda: dict(sem=i32(1), payload=f32(B), out=f32(2 * B)), + "ok", "T1", (I32, TERM), None), + ("counted_while", 1, lambda: dict(x=i32(1), n=4), + "unsupported", None, (), "spin-shape"), + ("cas_lock", 2, lambda: dict(lock=i32(1), data=i32(B)), + "ok", "T1", (TERM,), None), + # the CAS operand gate + ("cas_bool_operand", 2, lambda: dict(flag=i32(1), out=i32(4)), + "ok", "T1", (), None), + ("cas_cast_operand", 2, lambda: dict(flag=i32(1), data=i32(1)), + "unsupported", None, (), "cas-value"), + # multipath: an early return, several loops, a loop under a branch whose + # bounds wrap only where it is not reached + ("early_return", 4, lambda: dict(x=f32(4 * B), out=f32(4 * B), T=4 * B), + "ok", "T0", (I32,), None), + ("sequential_loops", 4, lambda: dict(x=f32(B), out=f32(32 * B), n=2), + "ok", "T1", (I32,), None), + ("nested_loops", 4, lambda: dict(x=f32(64 * B), out=f32(64 * B), n=2, m=3), + "ok", "T1", (I32,), None), + ("loop_under_guard", 8, lambda: dict(out=i32(8 * B)), "ok", "T0", (), None), + # width obligations: qualified any-grid, discharged at the launch + ("wrapped_slots", 2, lambda: dict(out=i32(3)), "ok", "T0", (I32,), None), + # program 2 wraps to element 0 at this launch + ("wrapped_slots", 3, lambda: dict(out=i32(3)), + "unsupported", None, (), "integer-overflow"), + ("launch_wrap", 2, lambda: dict(x=i32(2), out=i32(2 * B), big=2**30), + "unsupported", None, (), "integer-overflow"), + ("where_wrap", 3, lambda: dict(out=i32(4)), "races", None, (), None), + # a mask on float flags (no snapshot): a possible race, not a witness + ("table_mask", 4, lambda: dict(flags=f32(4 * B), out=i32(B)), + "unsupported", None, (), "potential-race"), +] # fmt: skip + + +@pytest.mark.parametrize( + "name, grid, args, status, rung, qualifiers, kind", + ROWS, + ids=[f"{row[0]}-{i}" for i, row in enumerate(ROWS)], +) +def test_verdicts(name, grid, args, status, rung, qualifiers, kind): + det = launch(name, (grid,), **args()) + verdict = det.last_verdict + (config,) = verdict.per_config + assert (verdict.status, config.status, config.rung) == (status, status, rung) + assert config.qualifiers == qualifiers == verdict.qualifiers + assert (verdict.refusal and verdict.refusal.kind) == kind + if status == "ok": + scope = { + "T0": "any-params-any-grid", + "T1": "this-params-any-grid", + "T1-launch": "this-params-this-grid", + }[rung] + assert verdict.scope == scope and det.last_reports == [] + if name == "two_slots": + # the races at three programs and more are evidence only + assert all( + max(a[0], b[0]) >= 2 for a, b in (r.programs for r in config.grid_fragile) + ) + + +def test_the_global_barrier_is_a_fence_only_with_global_bits(): + """The reader records the barrier between the store and the reload; with + a barrier of shared memory only (injected), the pair races.""" + args = dict(hist=i32(B), out=i32(B), N=B) + det = RaceDetector(compile=True) + traced = tilelens.trace(det)(K.store_fence_load) + texts = [] + with inject_ttir(traced, lambda t: texts.append(t) or t): + traced[(1,)](**args) + assert det.last_status == "ok" + (fence,) = parse_ttir(texts[0], multipath=True, while_regions=True).fences + assert "ttg.barrier all" in texts[0] + with inject_ttir( + traced, lambda t: t.replace("ttg.barrier all", "ttg.barrier local") + ): + traced[(1,)](**args) + assert det.last_status == "races" + assert races(det) == {("RAW", "hist", 0, 0)} + + +def test_a_possible_race_beside_a_race_is_withheld(): + det = launch("table_mask_and_waw", (2,), flags=f32(2 * B), out=i32(B), both=i32(B)) + assert det.last_status == "races" + assert {r.first.tensor_name for r in det.last_reports} == {"both"} + (note,) = det.last_verdict.notes + assert "possible race(s) under over-approximation were withheld" in note + + +def test_the_snapshot_bounds(monkeypatch): + det = launch("scatter", (4,), idx=_perm(), out=i32(4 * B)) + assert det.last_status == "ok" + monkeypatch.setattr(CompiledRaceDetector, "SNAPSHOT_MAX_ELEMENTS", 4 * B - 1) + det = launch("scatter", (4,), idx=_perm(), out=i32(4 * B)) + assert det.last_verdict.refusal.kind == "snapshot-bound" + monkeypatch.setattr(CompiledRaceDetector, "SNAPSHOT_MAX_ELEMENTS", 4 * B) + monkeypatch.setattr(CompiledRaceDetector, "SNAPSHOT_CAP_BYTES", 4 * B) + det = launch("scatter", (4,), idx=_perm(), out=i32(4 * B)) + assert det.last_verdict.refusal.kind == "indirect-address" + assert "snapshot cap" in det.last_verdict.refusal.message + + +def test_a_timeout_is_unsupported(monkeypatch): + monkeypatch.setattr(CompiledRaceDetector, "T0_TIMEOUT_MS", 1) + monkeypatch.setattr(CompiledRaceDetector, "T1_TIMEOUT_MS", 1) + idx = torch.arange(4 * B, dtype=torch.int32) + idx[5] = 40 + det = launch("scatter", (4,), idx=idx, out=i32(4 * B)) + assert det.last_status == "unsupported" + assert det.last_verdict.refusal.kind == "timeout" + + +def test_an_undecided_launch_grid_requery_records_no_race(monkeypatch): + found = TwoCopySymbolicHBSolver.find_races + + def requery_times_out(self, **kw): + if self.only_pairs is not None: + raise UnsupportedSymbolicRaceQuery("undecided", reason_unknown="timeout") + return found(self, **kw) + + monkeypatch.setattr(TwoCopySymbolicHBSolver, "find_races", requery_times_out) + det = launch("two_slots", (2,), out=i32(2 * B)) + verdict = det.last_verdict + assert (verdict.status, verdict.refusal.kind) == ("unsupported", "timeout") + # the any-grid races are evidence only + assert trace_module.launches[-1].records == [verdict] and not det.last_reports + assert verdict.per_config[0].grid_fragile + + +def test_a_config_that_needs_two_ctas_is_unsupported_and_the_rest_decided(): + @triton.autotune( + configs=[ + triton.Config({"BLOCK": 16}, num_ctas=1), + triton.Config({"BLOCK": 16}, num_ctas=2), + ], + key=[], + ) + @triton.jit + def copy(x, out, BLOCK: tl.constexpr): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + tl.store(out + offs, tl.load(x + offs)) + + det = launch(copy, (4,), x=f32(64), out=f32(64)) + verdict = det.last_verdict + ok, failed = verdict.per_config + assert (ok.config["num_ctas"], ok.status, ok.rung) == (1, "ok", "T0"), ok.refusal + assert (failed.config["num_ctas"], failed.status) == (2, "unsupported") + assert failed.refusal.kind == verdict.refusal.kind == "compile-failed" + assert "RaceDetector(compile=True, target=...)" in failed.refusal.message + assert verdict.status == "unsupported" + + +def test_the_launch_holds_what_every_config_proved(): + @triton.autotune( + configs=[triton.Config({"SLOTS": False}), triton.Config({"SLOTS": True})], + key=[], + ) + @triton.jit + def rows(out, SLOTS: tl.constexpr): + # with SLOTS, race-free for two programs only + pid = tl.program_id(0) + row = pid % 2 if SLOTS else pid + tl.store(out + row * 16 + tl.arange(0, 16), 1) + + det = launch(rows, (2,), out=i32(2 * B)) + verdict = det.last_verdict + assert [c.rung for c in verdict.per_config] == ["T0", "T1-launch"] + assert (verdict.status, verdict.scope) == ("ok", "this-params-this-grid") + + +def test_a_race_outranks_a_config_that_failed_to_compile(): + @triton.autotune( + configs=[ + triton.Config({"BLOCK": 32}, num_ctas=1), + triton.Config({"BLOCK": 32}, num_ctas=2), + ], + key=[], + ) + @triton.jit + def store_rows(out, BLOCK: tl.constexpr): + tl.store(out + tl.program_id(0) * 16 + tl.arange(0, BLOCK), 1) + + det = launch(store_rows, (4,), out=i32(128)) + assert [c.status for c in det.last_verdict.per_config] == ["races", "unsupported"] + assert det.last_status == "races" + + +def test_bindings_that_differ_in_grid_or_contents_are_solved_apart(monkeypatch): + solves = [] + analyze = client_module.analyze_graph + monkeypatch.setattr( + client_module, + "analyze_graph", + lambda *a, **kw: solves.append(a[3]) or analyze(*a, **kw), + ) + det = CompiledRaceDetector() + out = i32(2 * B) + facts = tensor_facts(out) + + def binding(grid): + return LaunchBinding( + params=MappingProxyType({}), + tensors=MappingProxyType({"out": facts}), + constexprs=MappingProxyType({}), + raw_grid=grid, + grid=grid, + config=MappingProxyType({}), + ) + + bindings = [binding((2, 1, 1)), binding((3, 1, 1)), binding((2, 1, 1))] + bindings.append(binding((2, 1, 1))) + for b, value in zip(bindings, (0, 0, 0, 1)): + det._contents[id(b)] = (b, {"out": LaunchTensor(facts, (value,) * 2 * B)}) + _, verdict, _ = det._analyze_config( + None, K.graph("two_slots"), "text", {}, bindings, {} + ) + # the third binding is the first's; the fourth differs in its contents + assert solves == [(2, 1, 1), (3, 1, 1), (2, 1, 1)] + assert verdict.status == "races" + + +def test_launches_on_two_threads_take_turns(): + """Z3's main context is not thread-safe: two host threads' traced + launches (of two kernels: one JITFunction is traced by one launch at a + time) are analyzed one at a time.""" + verdicts, errors = [], [] + + def work(name, **args): + try: + for _ in range(5): + verdicts.append(launch(name, (4,), **args).last_status) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + threads = [ + threading.Thread(target=work, args=("inplace",), kwargs=dict(x=i32(4 * B))), + threading.Thread( + target=work, + args=("add",), + kwargs=dict(x=f32(64), y=f32(64), out=f32(64), n=64), + ), + ] + for t in threads: + t.start() + for t in threads: + t.join() + assert errors == [] and verdicts == ["ok"] * 10 + + +def test_undecided_parts_always_print_once_and_proofs_only_under_verbose( + monkeypatch, capsys +): + monkeypatch.setattr(cfg, "verbose", False) + det = RaceDetector(compile=True) + for _ in range(3): + launch("counted_while", (1,), detector=det, x=i32(1), n=4) + launch("add", (2,), x=f32(64), y=f32(64), out=f32(64), n=64) + out = capsys.readouterr().out + assert out.count("not analyzed: spin-shape") == 1 + assert "race-free" not in out + monkeypatch.setattr(cfg, "verbose", True) + launch("add", (2,), x=f32(64), y=f32(64), out=f32(64), n=64) + assert "race-free" in capsys.readouterr().out + + +def test_autotune_configs_are_folded_and_identical_ones_solved_once(monkeypatch): + @triton.autotune( + configs=[ + triton.Config({"BLOCK": 16}, num_warps=2), + triton.Config({"BLOCK": 16}, num_warps=4), + triton.Config({"BLOCK": 32}, num_warps=4), + ], + key=[], + ) + @triton.jit + def store_rows(out, BLOCK: tl.constexpr): + # BLOCK=32 programs overlap on 16-element rows + tl.store(out + tl.program_id(0) * 16 + tl.arange(0, BLOCK), 1) + + solves = [] + analyze = client_module.analyze_graph + monkeypatch.setattr( + client_module, + "analyze_graph", + lambda *a, **kw: solves.append(a[0].kernel_name) or analyze(*a, **kw), + ) + det = launch(store_rows, (4,), out=i32(128)) + configs = det.last_verdict.per_config + assert [(c.config["BLOCK"], c.config["num_warps"], c.status) for c in configs] == [ + (16, 2, "ok"), + (16, 4, "ok"), + (32, 4, "races"), + ] + assert len(solves) == 2 # the two BLOCK=16 configs share a TTIR + assert det.last_status == "races" + assert {r.config["BLOCK"] for r in det.last_reports} == {32} + + +def test_the_verdict_and_races_survive_save_and_load(tmp_path, monkeypatch): + det = launch("two_slots", (3,), out=i32(2 * B)) + det = launch("two_slots", (2,), detector=det, out=i32(2 * B)) + monkeypatch.setattr(trace_module, "launches", trace_module.launches[-2:]) + saved = [list(launch.records) for launch in trace_module.launches] + tilelens.load(tilelens.save(tmp_path / "trace.zip")) + assert [list(launch.records) for launch in trace_module.launches] == saved + assert saved[0][-1].status == "races" + assert saved[1][-1].per_config[0].grid_fragile + + +# ── fenced CAS publications and the comm/comp pattern ── + + +@K.real_jit +def guarded_acq_rel(flag_ptr, data_ptr, out_ptr): + pid = tl.program_id(0) + tl.store(data_ptr, 1, mask=pid == 0) + tl.debug_barrier() + old = tl.atomic_cas(flag_ptr, tl.where(pid == 0, 0, 1), 1, sem="acq_rel") + tl.debug_barrier() + x = tl.load(data_ptr, mask=(pid == 1) & (old == 1), other=0) + tl.store(out_ptr + pid, x, mask=(pid == 1) & (old == 1)) + + +@K.real_jit +def acquire_on_failure(flag_ptr, data_ptr, out_ptr): + # the consumer's CAS never succeeds (cmp 7), yet reading the release + # acquires it + pid = tl.program_id(0) + tl.store(data_ptr, 1, mask=pid == 0) + tl.debug_barrier() + old = tl.atomic_cas(flag_ptr, tl.where(pid == 0, 0, 7), 1, sem="acq_rel") + tl.debug_barrier() + x = tl.load(data_ptr, mask=(pid == 1) & (old == 1), other=0) + tl.store(out_ptr + pid, x, mask=(pid == 1) & (old == 1)) + + +@K.real_jit +def role_specific_order(flag_ptr, data_ptr, out_ptr): + pid = tl.program_id(0) + if pid == 0: + tl.store(data_ptr, 1) + tl.debug_barrier() + tl.atomic_cas(flag_ptr, 0, 1, sem="release") + else: + old = tl.atomic_cas(flag_ptr, 1, 1, sem="acquire") + tl.debug_barrier() + value = tl.load(data_ptr, mask=old == 1, other=0) + tl.store(out_ptr + pid, value, mask=old == 1) + + +@K.real_jit +def cas_unlock_mutex(flag_ptr, data_ptr, out_ptr): + pid = tl.program_id(0) + while tl.atomic_cas(flag_ptr, 0, 1, sem="acquire") != 0: + pass + tl.debug_barrier() + tl.store(data_ptr, tl.load(data_ptr) + 1) + tl.debug_barrier() + tl.atomic_cas(flag_ptr, 1, 0, sem="release") + tl.store(out_ptr + pid, 1) + + +@K.real_jit +def fenced_tile_handoff(flag_ptr, data_ptr, out_ptr, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = tl.arange(0, BLOCK) + tl.store(data_ptr + offs, offs + 1, mask=pid == 0) + tl.debug_barrier() + old = tl.atomic_cas(flag_ptr, tl.where(pid == 0, 0, 1), 1, sem="acq_rel") + tl.debug_barrier() + value = tl.load(data_ptr + offs, mask=(pid == 1) & (old == 1), other=0) + tl.store(out_ptr + offs, value, mask=(pid == 1) & (old == 1)) + + +@pytest.mark.parametrize( + "kernel", + [ + guarded_acq_rel, + acquire_on_failure, + role_specific_order, + cas_unlock_mutex, + fenced_tile_handoff, + ], +) +def test_fenced_cas_publications_are_proved(kernel): + kw = {"BLOCK": B} if kernel is fenced_tile_handoff else {} + det = launch(kernel, (2,), flag_ptr=i32(1), data_ptr=i32(B), out_ptr=i32(B), **kw) + assert det.last_status == "ok", det.last_verdict + assert det.last_reports == [] + + +def test_comm_comp_is_proved_assuming_termination(): + det = launch("comm_comp", (3,), sem=i32(1), payload=f32(B), out=f32(2 * B)) + assert det.last_status == "ok" + assert TERM in det.last_verdict.qualifiers + + +# ── entry points ── + + +def test_the_cli_compile_flag_runs_the_compiled_race_detector(tmp_path): + script = tmp_path / "waw.py" + script.write_text( + "import torch, triton, triton.language as tl\n" + "@triton.jit\n" + "def waw(out, BLOCK: tl.constexpr):\n" + " tl.store(out + tl.arange(0, BLOCK), tl.program_id(0))\n" + "waw[(2,)](torch.zeros(16, dtype=torch.int32), BLOCK=16)\n" + "print('launch returned')\n" + ) + cli = ( + f"import sys; sys.argv = ['tile-race', '--compile', {str(script)!r}]; " + "from tilelens.wrapper import apply_race_detector; apply_race_detector()" + ) + env = {k: v for k, v in os.environ.items() if k != "TRITON_INTERPRET"} + env.update(PYTHONPATH=str(REPO), CUDA_VISIBLE_DEVICES="") + proc = subprocess.run( + [sys.executable, "-c", cli], capture_output=True, text=True, env=env, cwd=REPO + ) + assert proc.returncode == 0, proc.stderr + assert proc.stdout.count("(compiled race detector)") == 1 + assert f"File: {script}, Line: 4, in waw" in proc.stdout + assert "launch returned" in proc.stdout diff --git a/tests/unit/ir/inject.py b/tests/unit/ir/inject.py new file mode 100644 index 00000000..becbaf05 --- /dev/null +++ b/tests/unit/ir/inject.py @@ -0,0 +1,69 @@ +"""Feed a chosen TTIR text to a traced launch's IR clients. + +``inject_ttir(traced, text_or_fn)`` wraps the trace's host compiler (the +ClientManager's ``compiler``): each call is still bound, specialized, +compiled and delivered by the product (bind errors, dedup, the target and +Launch.grid included); only the compiled kernel's ``asm["ttir"]`` is +replaced, by ``text_or_fn`` or by ``text_or_fn(product_ttir)``. The +specialization becomes a hash of the product's and the injected text, so +two injected texts never share one. + +The injected function must keep every argument the product's has, with the +same types (it may add arguments the launch binds); a mismatch raises +AssertionError when the context exits (an error raised inside the compile +would be delivered as a compile failure instead). + +Not a test module: pytest imports it (python_files = *.py) and finds nothing. +""" + +from __future__ import annotations + +import re +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from dataclasses import replace +from types import MappingProxyType + +from tilelens.ir.capture import content_key + +_FUNC = re.compile(r"tt\.func\s+(?:\w+\s+)?@(\w+)\((.*?)\)\s*(?:->|attributes|\{)") +_ARG = re.compile(r"%([\w.$-]+)\s*:\s*([^\s,)]+)") + + +def func_args(text: str) -> tuple[str, dict[str, str]]: + """The kernel's name and its arguments' types, by name.""" + match = _FUNC.search(text) + assert match is not None, "no tt.func in the TTIR text" + return match.group(1), dict(_ARG.findall(match.group(2))) + + +@contextmanager +def inject_ttir(traced, text_or_fn: str | Callable[[str], str]) -> Iterator[None]: + compiler = traced.client_manager.compiler + compile_ = compiler.compile + errors: list[str] = [] + + def compile(jit_fn, args, kwargs, *, target, **options): + kernel = compile_(jit_fn, args, kwargs, target=target, **options) + product = kernel.asm["ttir"] + text = text_or_fn(product) if callable(text_or_fn) else text_or_fn + name, wanted = func_args(product) + got_name, got = func_args(text) + if got_name != name: + errors.append(f"the injected function is @{got_name}, not @{name}") + for arg, ty in wanted.items(): + if got.get(arg) != ty: + errors.append(f"argument %{arg}: {ty} injected as {got.get(arg)}") + asm = dict(kernel.asm, ttir=text) + return replace( + kernel, + hash=content_key(f"{kernel.hash}\0{text}"), + asm=MappingProxyType(asm), + ) + + compiler.compile = compile + try: + yield + finally: + del compiler.compile + assert not errors, "; ".join(errors) diff --git a/tests/unit/ir/test_lowering.py b/tests/unit/ir/test_lowering.py index 9b4ad1d2..9daa8db2 100644 --- a/tests/unit/ir/test_lowering.py +++ b/tests/unit/ir/test_lowering.py @@ -313,6 +313,18 @@ def test_division_truncates_toward_zero(a, b, quotient, remainder): assert _eval(Bin("%", Const(a), Const(b), 32)) == remainder +def test_a_client_divrem_encodes_quotients_and_remainders(): + seen = [] + + def divrem(t, a, b): + seen.append(t.op) + return a + b + + term = Bin("u//", Param("x"), Bin("%", Param("x"), Const(3))) + lowered = Lowerer(_graph(), _Leaves(), divrem=divrem).lower(term) + assert seen == ["%", "u//"] and _proved(lowered == _v("x") + _v("x") + 3) + + def test_loop_terms_read_the_clients_iteration(): """LoopVar is lower + k * step, IterArgOffset offset0 + k * delta, with k the leaves' iteration of that loop: an IterArgInfo without a loop_ssa diff --git a/tests/unit/race_detector/compiled_kernels.py b/tests/unit/race_detector/compiled_kernels.py new file mode 100644 index 00000000..88bedbe5 --- /dev/null +++ b/tests/unit/race_detector/compiled_kernels.py @@ -0,0 +1,583 @@ +"""Kernels whose TTIR the compiled race detector's unit tests analyze. + +``ttir(name)`` host-compiles kernel ``name`` (with its ``SPECS`` entry: the +signature and constexprs) up to the TTIR stage at test time, through +``tests/unit/ir/ttir_corpus.py``. + +Not a test module: pytest imports it (python_files = *.py) and finds nothing. +""" + +from __future__ import annotations + +import functools + +import triton +import triton.language as tl +from triton import knobs + +from tilelens.ir.ttir_reader import AccessGraph, parse_ttir + +B = 16 + + +def real_jit(fn): + """``triton.jit``, building a JITFunction whatever TRITON_INTERPRET says + (tests/unit/test_multithreading.py sets it when imported, and a module + collected after it would get InterpretedFunctions); the knob is put back + as it was.""" + missing = object() + previous = knobs.runtime.__dict__.get("interpret", missing) + knobs.runtime.__dict__["interpret"] = False + try: + return triton.jit(fn) + finally: + if previous is missing: + knobs.runtime.__dict__.pop("interpret", None) + else: + knobs.runtime.__dict__["interpret"] = previous + + +# ── one kernel per race class ── + + +@real_jit +def add(x, y, out, n, BLOCK: tl.constexpr): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + m = offs < n + tl.store(out + offs, tl.load(x + offs, mask=m) + tl.load(y + offs, mask=m), mask=m) + + +@real_jit +def waw(out, BLOCK: tl.constexpr): + tl.store(out + tl.arange(0, BLOCK), tl.program_id(0)) + + +@real_jit +def raw(x, out, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = tl.arange(0, BLOCK) + tl.store(x + offs, 1, mask=pid == 0) + tl.store(out + pid * BLOCK + offs, tl.load(x + offs, mask=pid == 1), mask=pid == 1) + + +@real_jit +def war(x, out, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = tl.arange(0, BLOCK) + tl.store(out + pid * BLOCK + offs, tl.load(x + offs, mask=pid == 0), mask=pid == 0) + tl.store(x + offs, 2, mask=pid == 1) + + +@real_jit +def atomic_vs_load(x, out, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = tl.arange(0, BLOCK) + tl.atomic_add(x + offs, 1, mask=pid == 0) + tl.store(out + pid * BLOCK + offs, tl.load(x + offs), mask=pid == 1) + + +@real_jit +def atomic_only(x, BLOCK: tl.constexpr): + tl.atomic_add(x + tl.arange(0, BLOCK), 1) + + +@real_jit +def tile2d(out, BLOCK: tl.constexpr): + # one make_range for both dimensions of the tile + r = tl.arange(0, BLOCK) + tl.store(out + r[:, None] * BLOCK + r[None, :], 1) + + +# ── rungs ── + + +@real_jit +def strided(x, out, stride, BLOCK: tl.constexpr): + # pid * stride is nonlinear with the params symbolic: no T0 + offs = tl.program_id(0) * stride + tl.arange(0, BLOCK) + tl.store(out + offs, tl.load(x + offs)) + + +@real_jit +def two_slots(out, BLOCK: tl.constexpr): + # race-free for at most two programs only + tl.store(out + (tl.program_id(0) % 2) * BLOCK + tl.arange(0, BLOCK), 1) + + +@real_jit +def scatter(idx, out, BLOCK: tl.constexpr): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + tl.store(out + tl.load(idx + offs), 1) + + +@real_jit +def launch_wrap(x, out, big, BLOCK: tl.constexpr): + # race-free for two programs; with big = 2**30, pid * big * 2 wraps + # at program 1 + pid = tl.program_id(0) + v = tl.load(x + (pid * big * 2) // big // 2) + tl.store(out + (pid % 2) * BLOCK + tl.arange(0, BLOCK), v) + + +@real_jit +def wrapped_slots(out): + # (pid * 2**30) // 2**30 is pid until the product wraps, at pid 2 + pid = tl.program_id(0) + tl.store(out + ((pid * 1073741824) // 1073741824) % 3, 1) + + +@real_jit +def two_widths(out): + # an i32 product that wraps at pids 2 and 3, an i64 one from pid 4 on, + # on disjoint paths: no one program shows both widths wrapping + pid = tl.program_id(0) + if pid < 4: + tl.store(out + (pid * 1073741824) // 1073741824, 1) + else: + tl.store(out + (pid.to(tl.int64) * 2**61) // 2**61, 1) + + +@real_jit +def where_wrap(out): + # pid * 2**30 wraps from pid 2 on, where the other arm is taken: pids + # >= 2 store out[0], as pid 0 does + pid = tl.program_id(0) + v = pid * 1073741824 + tl.store(out + tl.where(pid >= 2, 0, v // 1073741824), 1) + + +@real_jit +def masked_table(flags, out, stride, BLOCK: tl.constexpr): + # pid * stride: no T0; race-free whatever the table holds + offs = tl.program_id(0) * stride + tl.arange(0, BLOCK) + tl.store(out + offs, 1, mask=tl.load(flags + offs) != 0) + + +@real_jit +def table_mask(flags, out, BLOCK: tl.constexpr): + # every program stores out[0:BLOCK] where its flags are set + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + tl.store(out + tl.arange(0, BLOCK), 1, mask=tl.load(flags + offs) != 0) + + +@real_jit +def table_mask_and_waw(flags, out, both, BLOCK: tl.constexpr): + d = tl.arange(0, BLOCK) + tl.store(out + d, 1, mask=tl.load(flags + tl.program_id(0) * BLOCK + d) != 0) + tl.store(both + d, tl.program_id(0)) + + +@real_jit +def masked_gather_other(idx, out, BLOCK: tl.constexpr): + # lanes >= 8 load nothing and store out[0], in every program + d = tl.arange(0, BLOCK) + i = tl.load(idx + tl.program_id(0) * BLOCK + d, mask=d < 8, other=0) + tl.store(out + i, 1) + + +@real_jit +def handoff_index(flag, idx, out): + # program 0 publishes idx[0] = 5 and stores out[5] after its release; + # program 1 acquires, reads idx[0] and stores out[idx[0]]: a WAW + pid = tl.program_id(0) + if pid == 0: + tl.store(idx, 5) + tl.debug_barrier() + tl.atomic_xchg(flag, 1, sem="release") + tl.store(out + 5, 1) + else: + while tl.atomic_add(flag, 0, sem="acquire") != 1: + pass + tl.debug_barrier() + tl.store(out + tl.load(idx), 2) + + +# ── fences ── + + +@real_jit +def store_then_load(hist, out, N: tl.constexpr): + p = tl.program_id(0) + tl.store(hist + tl.arange(0, N), tl.arange(0, N)) + tl.store(out + p, tl.load(hist + p)) + + +@real_jit +def store_fence_load(hist, out, N: tl.constexpr): + p = tl.program_id(0) + tl.store(hist + tl.arange(0, N), tl.arange(0, N)) + tl.debug_barrier() + tl.store(out + p, tl.load(hist + p)) + + +@real_jit +def store_flag_fence_load(hist, out, flag, N: tl.constexpr): + p = tl.program_id(0) + tl.store(hist + tl.arange(0, N), tl.arange(0, N)) + if flag: + tl.debug_barrier() + tl.store(out + p, tl.load(hist + p)) + + +@real_jit +def loop_rows_fenced(t, u, n, BLOCK: tl.constexpr): + # each iteration writes its own row, then reads it back reversed + d = tl.arange(0, BLOCK) + for i in range(n): + tl.store(t + i * BLOCK + d, d) + tl.debug_barrier() + tl.store(u + i * BLOCK + d, tl.load(t + i * BLOCK + (BLOCK - 1 - d))) + + +@real_jit +def loop_buffer_fenced(t, u, n, BLOCK: tl.constexpr): + # one buffer for every iteration: nothing orders iteration i's reversed + # read before iteration i + 1's write + d = tl.arange(0, BLOCK) + for i in range(n): + tl.store(t + d, d + i) + tl.debug_barrier() + tl.store(u + d, tl.load(t + (BLOCK - 1 - d))) + + +@real_jit +def loop_buffer_two_fences(t, u, n, BLOCK: tl.constexpr): + d = tl.arange(0, BLOCK) + for i in range(n): + tl.store(t + d, d + i) + tl.debug_barrier() + tl.store(u + d, tl.load(t + (BLOCK - 1 - d))) + tl.debug_barrier() + + +# ── dependencies ── + + +@real_jit +def inplace(x, BLOCK: tl.constexpr): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + tl.store(x + offs, tl.load(x + offs) + 1) + + +@real_jit +def inplace_shift(x, BLOCK: tl.constexpr): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + tl.store(x + offs, tl.load(x + offs + 1)) + + +@real_jit +def atomic_chain(x, out, n, BLOCK: tl.constexpr): + # the store's value comes from the load through an unmasked exchange + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + m = offs < n + v = tl.load(x + offs, mask=m) + tl.atomic_add(out + offs, v, mask=m) + tl.store(x + offs, tl.atomic_xchg(out + offs, v), mask=m) + + +@real_jit +def masked_chain(x, y, flag, BLOCK: tl.constexpr): + # where the atomic is masked off, the store's value is not the load's + offs = tl.arange(0, BLOCK) + old = tl.atomic_add(y + offs, tl.load(x + offs), mask=offs < flag) + tl.store(x + offs, old) + + +# ── tile shapes ── + + +@real_jit +def one_address(out, BLOCK: tl.constexpr): + # every lane stores its own value to out[pid] + tl.store( + out + tl.program_id(0) + tl.zeros([BLOCK], dtype=tl.int32), tl.arange(0, BLOCK) + ) + + +@real_jit +def one_address_masked(out, BLOCK: tl.constexpr): + r = tl.arange(0, BLOCK) + tl.store(out + tl.zeros([BLOCK], dtype=tl.int32), r, mask=r == 0) + + +@real_jit +def broadcast_rows(out, BLOCK: tl.constexpr): + # the columns of a row store different values to out[row] + r = tl.arange(0, BLOCK) + p = tl.broadcast_to(out + r[:, None], (BLOCK, BLOCK)) + tl.store(p, r[:, None] * BLOCK + r[None, :]) + + +# ── awaits ── + + +@real_jit +def comm_comp( + sem, + payload, + out, + BLOCK: tl.constexpr, + SEM: tl.constexpr, + EXPECT: tl.constexpr, + ALL_POLL: tl.constexpr, +): + # program 0 publishes a payload and arrives on a semaphore; the others + # await it, then read the payload + pid = tl.program_id(0) + if pid < 1: + offs = tl.arange(0, BLOCK) + tl.store(payload + offs, (offs + 1).to(tl.float32)) + tl.debug_barrier() + tl.atomic_xchg(sem, 1, sem="release") + else: + if ALL_POLL or pid == 1: + while tl.atomic_add(sem, 0, sem=SEM) != EXPECT: + pass + tl.debug_barrier() + offs = tl.arange(0, BLOCK) + tl.store(out + (pid - 1) * BLOCK + offs, tl.load(payload + offs)) + + +@real_jit +def exit_wrap(flag, data, out, big, BLOCK: tl.constexpr): + # program 4's expected flag value wraps to 0 in i32, the initial value: + # it reads data while program 0 writes it + pid = tl.program_id(0) + offs = tl.arange(0, BLOCK) + if pid == 0: + tl.store(data + offs, offs + 1) + tl.debug_barrier() + tl.atomic_xchg(flag, 1, sem="release") + elif (pid == 1) | (pid == 4): + while tl.atomic_add(flag, 0, sem="acquire") != (pid * big) // big: + pass + tl.debug_barrier() + tl.store(out + pid * BLOCK + offs, tl.load(data + offs)) + + +@real_jit +def counted_while(x, n): + i = 0 + while i < n: + i += 1 + tl.store(x, i) + + +@real_jit +def cas_lock(lock, data, BLOCK: tl.constexpr, SEM: tl.constexpr): + while tl.atomic_cas(lock, 0, 1, sem="acquire") != 0: + pass + tl.debug_barrier() + offs = tl.arange(0, BLOCK) + tl.store(data + offs, tl.load(data + offs) + 1) + tl.debug_barrier() + tl.atomic_xchg(lock, 0, sem=SEM) + + +# ── CAS operands ── + + +@real_jit +def cas_handoff( + flag, data, out, PSEM: tl.constexpr, CSEM: tl.constexpr, SCOPE: tl.constexpr +): + # program 0 publishes data with a CAS; the others read it once their + # (failing) CAS observes the publication + pid = tl.program_id(0) + if pid == 0: + tl.store(data, 1) + tl.debug_barrier() + tl.atomic_cas(flag, 0, 1, sem=PSEM, scope=SCOPE) + else: + old = tl.atomic_cas(flag, 0, 0, sem=CSEM, scope=SCOPE) + tl.debug_barrier() + tl.store(out + pid, tl.load(data, mask=old == 1), mask=old == 1) + + +@real_jit +def cas_bool_operand(flag, out): + pid = tl.program_id(0) + old = tl.atomic_cas(flag, (pid == 0).to(tl.int32), 1) + tl.store(out + pid, old) + + +@real_jit +def cas_cast_operand(flag, data): + prior = tl.atomic_cas(data, 0, 1) + tl.atomic_cas(flag, prior.to(tl.int8).to(tl.int32), 1) + + +@real_jit +def cas_arith_operand(flag, data): + prior = tl.atomic_cas(data, 0, 1) + tl.atomic_cas(flag, prior + 1, 1) + + +@real_jit +def cas_in_loop(flag, n): + for _ in range(n): + tl.atomic_cas(flag, 0, 1) + + +# ── multipath ── + + +@real_jit +def early_return(x, out, T, BLOCK: tl.constexpr): + pid = tl.program_id(0) + if pid * BLOCK >= T: + return + offs = pid * BLOCK + tl.arange(0, BLOCK) + tl.store(out + offs, tl.load(x + offs)) + + +@real_jit +def sequential_loops(x, out, n, BLOCK: tl.constexpr, SHIFT: tl.constexpr): + pid = tl.program_id(0) + d = tl.arange(0, BLOCK) + for i in range(n): + tl.store(out + (pid * 2 * n + i) * BLOCK + d, tl.load(x + d)) + for j in range(n): + tl.store(out + ((pid * 2 + SHIFT) * n + j) * BLOCK + d, j) + + +@real_jit +def read_loop_then_inplace(x, y, n, BLOCK: tl.constexpr, IN_LOOP: tl.constexpr): + # a loop that only reads; then x updated in place, in a second loop or + # in no loop + pid = tl.program_id(0) + d = tl.arange(0, BLOCK) + acc = tl.zeros([BLOCK], dtype=tl.float32) + for i in range(n): + acc += tl.load(y + pid * BLOCK + d) + if IN_LOOP: + for j in range(2): + p = x + (pid * 2 + j) * BLOCK + d + tl.store(p, tl.load(p) + acc) + else: + p = x + pid * BLOCK + d + tl.store(p, tl.load(p) + acc) + + +@real_jit +def grouped_swizzle(out, n_m, n_n, GROUP: tl.constexpr): + # the grouped program-id order of a tiled matmul: one store per tile + pid = tl.program_id(0) + width = GROUP * n_n + first_m = (pid // width) * GROUP + group = tl.minimum(n_m - first_m, GROUP) + pid_m = first_m + (pid % width) % group + pid_n = (pid % width) // group + tl.store(out + pid_m * n_n + pid_n, pid) + + +@real_jit +def nested_loops(x, out, n, m, BLOCK: tl.constexpr): + pid = tl.program_id(0) + d = tl.arange(0, BLOCK) + for i in range(n): + for j in range(m): + row = (pid * n + i) * m + j + tl.store(out + row * BLOCK + d, tl.load(x + row * BLOCK + d)) + + +@real_jit +def loop_under_guard(out, BLOCK: tl.constexpr): + # the loop's bounds wrap for pid >= 8 only, where it is not reached + pid = tl.program_id(0) + d = tl.arange(0, BLOCK) + if pid < 4: + for i in range(pid * 268435456, pid * 268435456 + 2): + tl.store(out + (pid * 2 + i - pid * 268435456) * BLOCK + d, 1) + + +SPECS: dict[str, tuple[dict, dict]] = { + "add": ({"x": "*fp32", "y": "*fp32", "out": "*fp32", "n": "i32"}, {"BLOCK": B}), + "waw": ({"out": "*i32"}, {"BLOCK": B}), + "raw": ({"x": "*i32", "out": "*i32"}, {"BLOCK": B}), + "war": ({"x": "*i32", "out": "*i32"}, {"BLOCK": B}), + "atomic_vs_load": ({"x": "*i32", "out": "*i32"}, {"BLOCK": B}), + "atomic_only": ({"x": "*i32"}, {"BLOCK": B}), + "tile2d": ({"out": "*i32"}, {"BLOCK": B}), + "strided": ({"x": "*fp32", "out": "*fp32", "stride": "i32"}, {"BLOCK": B}), + "two_slots": ({"out": "*i32"}, {"BLOCK": B}), + "wrapped_slots": ({"out": "*i32"}, {}), + "where_wrap": ({"out": "*i32"}, {}), + "two_widths": ({"out": "*i32"}, {}), + "masked_table": ({"flags": "*i32", "out": "*i32", "stride": "i32"}, {"BLOCK": B}), + "table_mask": ({"flags": "*fp32", "out": "*i32"}, {"BLOCK": B}), + "table_mask_and_waw": ( + {"flags": "*fp32", "out": "*i32", "both": "*i32"}, + {"BLOCK": B}, + ), + "masked_gather_other": ({"idx": "*i32", "out": "*i32"}, {"BLOCK": B}), + "handoff_index": ({"flag": "*i32", "idx": "*i32", "out": "*i32"}, {}), + "atomic_chain": ({"x": "*fp32", "out": "*fp32", "n": "i32"}, {"BLOCK": B}), + "masked_chain": ({"x": "*i32", "y": "*i32", "flag": "i32"}, {"BLOCK": B}), + "one_address": ({"out": "*i32"}, {"BLOCK": B}), + "one_address_masked": ({"out": "*i32"}, {"BLOCK": B}), + "broadcast_rows": ({"out": "*i32"}, {"BLOCK": B}), + "exit_wrap": ( + {"flag": "*i32", "data": "*i32", "out": "*i32", "big": "i32"}, + {"BLOCK": B}, + ), + "cas_in_loop": ({"flag": "*i32", "n": "i32"}, {}), + "launch_wrap": ({"x": "*i32", "out": "*i32", "big": "i32"}, {"BLOCK": B}), + "cas_cast_operand": ({"flag": "*i32", "data": "*i32"}, {}), + "cas_arith_operand": ({"flag": "*i32", "data": "*i32"}, {}), + "loop_under_guard": ({"out": "*i32"}, {"BLOCK": B}), + "scatter": ({"idx": "*i32", "out": "*i32"}, {"BLOCK": B}), + "store_then_load": ({"hist": "*i32", "out": "*i32"}, {"N": B}), + "store_fence_load": ({"hist": "*i32", "out": "*i32"}, {"N": B}), + "store_flag_fence_load": ( + {"hist": "*i32", "out": "*i32", "flag": "i32"}, + {"N": B}, + ), + "loop_rows_fenced": ({"t": "*i32", "u": "*i32", "n": "i32"}, {"BLOCK": B}), + "loop_buffer_fenced": ({"t": "*i32", "u": "*i32", "n": "i32"}, {"BLOCK": B}), + "loop_buffer_two_fences": ({"t": "*i32", "u": "*i32", "n": "i32"}, {"BLOCK": B}), + "inplace": ({"x": "*i32"}, {"BLOCK": B}), + "inplace_shift": ({"x": "*i32"}, {"BLOCK": B}), + "comm_comp": ( + {"sem": "*i32", "payload": "*fp32", "out": "*fp32"}, + {"BLOCK": B, "SEM": "acquire", "EXPECT": 1, "ALL_POLL": True}, + ), + "counted_while": ({"x": "*i32", "n": "i32"}, {}), + "cas_lock": ({"lock": "*i32", "data": "*i32"}, {"BLOCK": B, "SEM": "release"}), + "cas_handoff": ( + {"flag": "*i32", "data": "*i32", "out": "*i32"}, + {"PSEM": "release", "CSEM": "acquire", "SCOPE": "gpu"}, + ), + "cas_bool_operand": ({"flag": "*i32", "out": "*i32"}, {}), + "early_return": ({"x": "*fp32", "out": "*fp32", "T": "i32"}, {"BLOCK": B}), + "sequential_loops": ( + {"x": "*fp32", "out": "*fp32", "n": "i32"}, + {"BLOCK": B, "SHIFT": 1}, + ), + "read_loop_then_inplace": ( + {"x": "*fp32", "y": "*fp32", "n": "i32"}, + {"BLOCK": B, "IN_LOOP": True}, + ), + "grouped_swizzle": ({"out": "*i32", "n_m": "i32", "n_n": "i32"}, {"GROUP": 4}), + "nested_loops": ( + {"x": "*fp32", "out": "*fp32", "n": "i32", "m": "i32"}, + {"BLOCK": B}, + ), +} + + +@functools.cache +def ttir(name: str, **constexprs: object) -> str: + """The TTIR of kernel ``name``, its SPECS constexprs updated by + ``constexprs``.""" + from ..ir import ttir_corpus + + sig, consts = SPECS[name] + consts = {**consts, **constexprs} + full = {**sig, **{k: "constexpr" for k in consts}} + return ttir_corpus._ttir(globals()[name], full, consts, cc=89) + + +def graph(name: str, **constexprs: object) -> AccessGraph: + """Kernel ``name`` read as the compiled race detector reads it.""" + return parse_ttir(ttir(name, **constexprs), multipath=True, while_regions=True) diff --git a/tests/unit/race_detector/test_compiled_analysis.py b/tests/unit/race_detector/test_compiled_analysis.py new file mode 100644 index 00000000..7827166d --- /dev/null +++ b/tests/unit/race_detector/test_compiled_analysis.py @@ -0,0 +1,370 @@ +"""The compiled race detector's analysis of host-compiled TTIR: one race +per class, a proof per rung, fences, awaits, CAS operands, dependency order, +multipath graphs and the width obligations (``analysis.analyze_graph`` over +``compiled_kernels``, compiled at test time).""" + +from __future__ import annotations + +import pytest +import torch + +from tilelens.clients.race_detector.compiled.analysis import analyze_graph +from tilelens.clients.race_detector.compiled.leaves import LaunchTensor +from tilelens.clients.race_detector.hb_common import UnsupportedSymbolicRaceQuery +from tilelens.clients.race_detector.two_copy_symbolic_hb_solver import ( + TwoCopySymbolicHBSolver, +) +from tilelens.ir.launch import tensor_facts + +from . import compiled_kernels as K + +B = K.B + + +def i32(n: int) -> torch.Tensor: + return torch.zeros(n, dtype=torch.int32) + + +def f32(n: int) -> torch.Tensor: + return torch.zeros(n, dtype=torch.float32) + + +def run(name, tensors, params=None, grid=(2, 1, 1), snap=(), timeout_ms=None, **cx): + """Analyze kernel ``name`` for a launch: ``tensors`` (the caller keeps + them alive, so no two share an address), ``snap`` the ones whose + pre-launch contents are captured.""" + launch = { + k: LaunchTensor( + tensor_facts(t), tuple(t.reshape(-1).tolist()) if k in snap else None + ) + for k, t in tensors.items() + } + kw = ( + {} + if timeout_ms is None + else dict(t0_timeout_ms=timeout_ms, t1_timeout_ms=timeout_ms) + ) + return analyze_graph(K.graph(name, **cx), params or {}, launch, grid, **kw) + + +def races(out) -> set[tuple[str, tuple[int, int], tuple[str, str]]]: + return { + ( + r.race_type.value, + tuple(sorted((r.witness_grid_a[0], r.witness_grid_b[0]))), + (r.first_record.tensor_name, r.second_record.tensor_name), + ) + for r in out.reports + } + + +@pytest.mark.parametrize( + "name, race", + [ + ("waw", ("WAW", "out", "out")), + ("raw", ("RAW", "x", "x")), + ("war", ("WAR", "x", "x")), + ("atomic_vs_load", ("RAW", "x", "x")), + ], +) +def test_each_race_class_between_two_programs(name, race): + tensors = {"x": i32(B), "out": i32(4 * B)} + out = run(name, {k: v for k, v in tensors.items() if k in K.SPECS[name][0]}) + assert out.status == "races" and not out.widened + kind, a, b = race + assert {(t, pids, names) for t, pids, names in races(out)} == { + (kind, (0, 1), (a, b)) + } + # the witness lines are the kernel's source lines + lines = {r.first_record.source_location[1] for r in out.reports} + assert all(line > 0 for line in lines) + + +def test_rungs(): + # T0: any parameters and grid, with the width assumption it needs + out = run("add", {"x": f32(1024), "y": f32(1024), "out": f32(1024)}, {"n": 1000}) + assert (out.status, out.rung, out.overflow) == ("proved", "T0", (32,)) + out = run("atomic_only", {"x": i32(B)}, grid=(4, 1, 1)) + assert (out.status, out.rung, out.overflow) == ("proved", "T0", ()) + # pid * stride is nonlinear with the params symbolic: T1 + tensors = {"x": f32(64 * B), "out": f32(64 * B)} + out = run("strided", tensors, {"stride": B}, grid=(4, 1, 1)) + assert (out.status, out.rung) == ("proved", "T1") + out = run("strided", tensors, {"stride": B // 2}, grid=(4, 1, 1)) + assert ("WAW", (0, 1), ("out", "out")) in races(out) + # race-free only up to two programs: the launch rung, with the + # any-grid reports kept as evidence only + out = run("two_slots", {"out": i32(2 * B)}) + assert (out.status, out.rung, out.reports) == ("proved", "T1-launch", []) + assert out.grid_fragile + out = run("two_slots", {"out": i32(2 * B)}, grid=(3, 1, 1)) + assert races(out) == {("WAW", (0, 2), ("out", "out"))} + # T0's one encoding per tensor needs disjoint tensors: aliased views + # race through the T1 byte addresses + x = f32(4 * B + 1) + tensors = {"x": x[: 4 * B], "y": f32(4 * B), "out": x[1:]} + out = run("add", tensors, {"n": 4 * B}, grid=(4, 1, 1)) + assert ("WAR", (0, 1), ("x", "out")) in races(out) + # a mask on loaded flags, no T0: the content-free attempt proves it for + # any contents, so the proof does not claim to read them + tensors = {"flags": torch.ones(4 * B, dtype=torch.int32), "out": i32(4 * B)} + out = run("masked_table", tensors, {"stride": B}, grid=(4, 1, 1), snap=("flags",)) + assert (out.status, out.rung, out.content) == ("proved", "T1", False) + # with no snapshot the mask is free: a race over it is only possible + out = run("table_mask", {"flags": f32(4 * B), "out": i32(B)}, grid=(4, 1, 1)) + assert (out.status, out.reports) == ("races", []) and out.widened + + +def test_loaded_index_reads_the_launch_snapshot(): + perm = torch.randperm(4 * B, generator=torch.Generator().manual_seed(0)) + tensors = {"idx": perm.to(torch.int32), "out": i32(4 * B)} + out = run("scatter", tensors, grid=(4, 1, 1), snap=("idx",)) + assert (out.status, out.rung, out.content) == ("proved", "T1", True) + dup = dict(tensors, idx=tensors["idx"].clone()) + dup["idx"][5] = dup["idx"][40] # programs 0 and 2 store one element + out = run("scatter", dup, grid=(4, 1, 1), snap=("idx",)) + assert out.content and races(out) == {("WAW", (0, 2), ("out", "out"))} + out = run("scatter", tensors, grid=(4, 1, 1)) + assert out.status == "unsupported" + assert out.refusal.kind == "indirect-address" + assert "no usable snapshot" in out.refusal.message + # a view of a larger allocation may legally read outside its values + buf = i32(8 * B) + buf[: 4 * B] = tensors["idx"] + out = run( + "scatter", + {"idx": buf[: 4 * B], "out": i32(4 * B)}, + grid=(4, 1, 1), + snap=("idx",), + ) + assert out.refusal.kind == "indirect-address" and "view" in out.refusal.message + # a table the kernel writes is no snapshot of what it reads (program 1 + # stores out[5] after program 0 published idx[0] = 5) + tensors = {"flag": i32(1), "idx": i32(1), "out": i32(8)} + out = run("handoff_index", tensors, snap=("flag", "idx")) + assert ( + out.refusal.kind == "indirect-address" and "writes 'idx'" in out.refusal.message + ) + # masked-off lanes read `other`: lanes >= 8 of every program store out[0] + idx = torch.arange(1, 4 * B + 1, dtype=torch.int32) + out = run( + "masked_gather_other", + {"idx": idx, "out": i32(5 * B)}, + grid=(4, 1, 1), + snap=("idx",), + ) + assert "WAW" in {t for t, _, _ in races(out)} + + +@pytest.mark.parametrize( + "name, params, tensors, expected", + [ + # one program: a store and a reload of another lane's element + ("store_then_load", {}, ("hist", "out"), "races"), + # ordered by a global barrier (any grid races: the launch rung) + ("store_fence_load", {}, ("hist", "out"), "proved"), + # a conditional barrier orders nothing + ("store_flag_fence_load", {"flag": 1}, ("hist", "out"), "races"), + # in a loop: the barrier orders one iteration's store and reload ... + ("loop_rows_fenced", {"n": 4}, ("t", "u"), "proved"), + # ... and not one iteration's reload against the next one's store + ("loop_buffer_fenced", {"n": 4}, ("t", "u"), "races"), + ("loop_buffer_two_fences", {"n": 4}, ("t", "u"), "proved"), + ], +) +def test_fences_order_one_program(name, params, tensors, expected): + out = run(name, {k: i32(4 * B) for k in tensors}, params, grid=(1, 1, 1)) + assert out.status == expected + if expected == "races": + assert races(out) == {("RAW", (0, 0), (tensors[0], tensors[0]))} + + +def test_dependency_orders_a_read_modify_write_at_one_lane(): + assert run("inplace", {"x": i32(4 * B)}, grid=(4, 1, 1)).status == "proved" + out = run("inplace_shift", {"x": i32(4 * B + 1)}, grid=(4, 1, 1)) + assert races(out) == {("WAR", (0, 1), ("x", "x"))} + # through an atomic that runs wherever the store does, the order composes + tensors = {"x": f32(4 * B), "out": f32(4 * B)} + assert run("atomic_chain", tensors, {"n": 4 * B}, grid=(4, 1, 1)).status == "proved" + # where the atomic is masked off, the stored value is not the load's + out = run("masked_chain", {"x": i32(B), "y": i32(B)}, {"flag": 0}, grid=(1, 1, 1)) + assert races(out) == {("WAR", (0, 0), ("x", "x"))} + + +@pytest.mark.parametrize( + "name, grid, race", + [ + # B lanes store B values to out[pid] + ("one_address", 2, True), + ("one_address_masked", 1, False), + # the columns of a broadcast row pointer store to one element + ("broadcast_rows", 1, True), + ("tile2d", 1, False), + ], +) +def test_lanes_not_named_by_a_variable_collide(name, grid, race): + out = run(name, {"out": i32(B * B)}, grid=(grid, 1, 1)) + assert races(out) == ({("WAW", (0, 0), ("out", "out"))} if race else set()) + assert race or out.status == "proved" + + +@pytest.mark.parametrize( + "cx, payload_race", + [ + ({}, False), + ({"SEM": "relaxed"}, True), # the poll acquires nothing + ({"EXPECT": 0}, True), # it exits on the initial value + ({"ALL_POLL": False}, True), # program 2 never polls + ], +) +def test_an_awaited_flag_orders_the_payload(cx, payload_race): + tensors = {"sem": i32(1), "payload": f32(B), "out": f32(2 * B)} + out = run("comm_comp", tensors, grid=(3, 1, 1), snap=("sem",), **cx) + assert out.assumes_termination + if payload_race: + assert ("RAW", "payload") in {(t, n[0]) for t, _, n in races(out)} + else: + assert out.status == "proved" + + +def test_a_wrapping_exit_value_drops_the_await(): + # program 4 expects (4 * 2**30) // 2**30, which wraps to 0: it does not + # wait for program 0's publication + tensors = {"flag": i32(1), "data": i32(B), "out": i32(5 * B)} + out = run("exit_wrap", tensors, {"big": 2**30}, grid=(5, 1, 1), snap=("flag",)) + assert ("RAW", (0, 4), ("data", "data")) in races(out) + out = run("exit_wrap", tensors, {"big": 2**30}, grid=(2, 1, 1), snap=("flag",)) + assert (out.status, out.rung) == ("proved", "T1-launch") + + +def test_a_while_that_is_no_spin_refuses(): + out = run("counted_while", {"x": i32(B)}, {"n": 4}, grid=(1, 1, 1)) + assert out.status == "unsupported" and out.refusal.kind == "spin-shape" + assert "carries values" in out.refusal.message + + +def test_cas_lock_and_operands(): + tensors = {"lock": i32(1), "data": i32(B)} + out = run("cas_lock", tensors, snap=("lock",)) + assert (out.status, out.assumes_termination) == ("proved", True) + out = run("cas_lock", tensors, snap=("lock",), SEM="relaxed") + assert ("WAW", (0, 1), ("data", "data")) in races(out) + # (pid == 0) as an int32 operand: the zero extension of a comparison + out = run("cas_bool_operand", {"flag": i32(1), "out": i32(4)}, snap=("flag",)) + assert out.status == "proved" + for name, why in ( + ("cas_cast_operand", "changes an integer's width"), + ("cas_arith_operand", "arithmetic"), + ): + out = run(name, {"flag": i32(1), "data": i32(1)}, snap=("flag", "data")) + assert out.status == "unsupported" and out.refusal.kind == "cas-value" + assert why in out.refusal.message + out = run("cas_in_loop", {"flag": i32(1)}, {"n": 2}, snap=("flag",)) + assert out.refusal.kind == "control-flow" + + +@pytest.mark.parametrize( + "cx, data_race", + [ + ({}, False), + ({"PSEM": "relaxed"}, True), + ({"CSEM": "relaxed"}, True), + ({"SCOPE": "cta"}, True), + ], +) +def test_a_failed_cas_acquires_a_cas_publication(cx, data_race): + tensors = {"flag": i32(1), "data": i32(1), "out": i32(4)} + out = run("cas_handoff", tensors, snap=("flag",), **cx) + assert ( + ("RAW", (0, 1), ("data", "data")) in races(out) + if data_race + else out.status == "proved" + ) + + +def test_multipath_graphs(): + out = run("early_return", {"x": f32(4 * B), "out": f32(4 * B)}, {"T": 4 * B}) + assert (out.status, out.rung) == ("proved", "T0") + tensors = {"x": f32(B), "out": f32(32 * B)} + out = run("sequential_loops", tensors, {"n": 2}, grid=(4, 1, 1)) + assert (out.status, out.rung) == ("proved", "T1") + out = run("sequential_loops", tensors, {"n": 2}, grid=(4, 1, 1), SHIFT=2) + assert ("WAW", "out") in {(t, n[0]) for t, _, n in races(out)} + tensors = {"x": f32(64 * B), "out": f32(64 * B)} + out = run("nested_loops", tensors, {"n": 2, "m": 3}, grid=(4, 1, 1)) + assert out.status == "proved" + # an in-place update beside a loop it is not in: one execution, whose + # load comes before its store (the other loop's iteration is no copy's) + for in_loop in (True, False): + tensors = {"x": f32(8 * B), "y": f32(4 * B)} + out = run( + "read_loop_then_inplace", tensors, {"n": 3}, grid=(4, 1, 1), IN_LOOP=in_loop + ) + assert (out.status, out.rung) == ("proved", "T0") + + +def test_width_obligations(): + # A loop under a branch: its bounds wrap only where it is not reached. + out = run("loop_under_guard", {"out": i32(8 * B)}, grid=(8, 1, 1)) + assert (out.status, out.rung, out.overflow) == ("proved", "T0", ()) + # A width premise holds only where a Select arm reads its term: pid * + # 2**30 wraps at pid 2, where the other arm stores out[0] as pid 0 does. + out = run("where_wrap", {"out": i32(4)}, grid=(3, 1, 1)) + assert races(out) == {("WAW", (0, 2), ("out", "out"))} + # Assuming no overflow, two programs at most store: an any-grid proof, + # qualified, that holds at a launch where nothing wraps ... + out = run("wrapped_slots", {"out": i32(3)}, grid=(2, 1, 1)) + assert (out.status, out.rung, out.overflow) == ("proved", "T0", (32,)) + # ... and refuses at one where program 2 wraps to element 0. + out = run("wrapped_slots", {"out": i32(3)}, grid=(3, 1, 1)) + assert (out.status, out.refusal.kind) == ("unsupported", "integer-overflow") + # The qualifier names every width that can wrap, though no one + # execution shows both. + out = run("two_widths", {"out": i32(8)}, grid=(2, 1, 1)) + assert (out.status, out.rung, out.overflow) == ("proved", "T0", (32, 64)) + # The launch rung discharges them: pid * 2**31 wraps at this launch. + tensors = {"x": i32(2), "out": i32(2 * B)} + assert run("launch_wrap", tensors, {"big": 2**29}).rung == "T1-launch" + out = run("launch_wrap", tensors, {"big": 2**30}) + assert out.status == "unsupported" + assert out.refusal.kind == "integer-overflow" + + +def test_an_undecided_query_is_a_timeout(): + perm = torch.arange(4 * B, dtype=torch.int32) + perm[5] = 40 + tensors = {"idx": perm, "out": i32(4 * B)} + out = run("scatter", tensors, grid=(4, 1, 1), snap=("idx",), timeout_ms=1) + assert out.status == "unsupported" and out.refusal.kind == "timeout" + + +def test_an_undecided_launch_grid_requery_is_a_timeout(monkeypatch): + """When the requery at the launch's grid is not decided, the any-grid + races (which may lie outside the launch's grid) are evidence only.""" + found = TwoCopySymbolicHBSolver.find_races + + def requery_times_out(self, **kw): + if self.only_pairs is not None: + raise UnsupportedSymbolicRaceQuery( + "Z3 could not decide", reason_unknown="timeout" + ) + return found(self, **kw) + + monkeypatch.setattr(TwoCopySymbolicHBSolver, "find_races", requery_times_out) + out = run("two_slots", {"out": i32(2 * B)}) + assert (out.status, out.refusal.kind, out.reports) == ("unsupported", "timeout", []) + assert "launch's grid was not decided" in out.refusal.message and out.grid_fragile + + +def test_a_cancelled_query_is_a_keyboard_interrupt(monkeypatch): + """Z3's Ctrl+C handler cancels a query as unknown; the analysis raises + instead of reading it as undecided and asking the next rung.""" + + def interrupted(self, **kw): + raise UnsupportedSymbolicRaceQuery( + "Z3 could not decide", reason_unknown="interrupted from keyboard" + ) + + monkeypatch.setattr(TwoCopySymbolicHBSolver, "find_races", interrupted) + with pytest.raises(KeyboardInterrupt): + run("waw", {"out": i32(B)}) diff --git a/tests/unit/race_detector/test_compiled_rules.py b/tests/unit/race_detector/test_compiled_rules.py new file mode 100644 index 00000000..05f0a801 --- /dev/null +++ b/tests/unit/race_detector/test_compiled_rules.py @@ -0,0 +1,355 @@ +"""The compiled race detector's leaves, await rules and fence rule.""" + +from __future__ import annotations + +from dataclasses import replace + +import pytest +import torch +from z3 import And, Solver, simplify, substitute, unsat + +from tilelens.clients.race_detector.compiled.analysis import analyze_graph +from tilelens.clients.race_detector.compiled.awaits import awaited_accesses +from tilelens.clients.race_detector.compiled.global_records import Encoder +from tilelens.clients.race_detector.compiled.leaves import ( + LaunchTensor, + RaceLeaves, + RaceRefusal, +) +from tilelens.clients.race_detector.two_copy_symbolic_hb_solver import ( + TwoCopySymbolicHBSolver, +) +from tilelens.clients import pids +from tilelens.ir.launch import tensor_facts +from tilelens.ir.lowering import truncating_divrem +from tilelens.ir.ttir_reader import Arange, Bin, Cmp, Const, DataDep, Param, Pid +from tilelens.ir.ttir_reader import Select as TSelect +from tilelens.ir.ttir_reader import parse_ttir + +from . import compiled_kernels as K + + +def _launch(**tensors) -> dict[str, LaunchTensor]: + return { + k: LaunchTensor(tensor_facts(t), tuple(t.tolist())) for k, t in tensors.items() + } + + +def test_the_solver_renames_a_lowered_pid_per_copy(): + """RaceLeaves' pids are the variables the two-copy solver renames: a + pid-dependent address must change under each copy's rename, or both + copies would be one program and every cross-program race missed.""" + out = torch.zeros(4 * K.B, dtype=torch.int32) + enc = Encoder(K.graph("raw"), {}, _launch(x=out[: K.B], out=out), {}).encode() + record = next(r for r in enc.records if "pid_0" in str(r.addr_expr)) + solver = TwoCopySymbolicHBSolver( + enc.records, grid=(2, 1, 1), arange_dict=enc.arange_dict + ) + for ctx in (solver.ctx_a, solver.ctx_b): + renamed = substitute(record.addr_expr, *ctx.pid_substitutions) + assert not renamed.eq(record.addr_expr) + assert f"pid_{ctx.label}_0" in str(renamed) + + +def test_lanes_are_keyed_by_dim_and_extent(): + graph = K.graph("tile2d") + leaves = RaceLeaves(graph, {}, {}) + leaves.value(graph.accesses[0].offset) + # one make_range read along two dims: two lanes, each in [0, B) + assert sorted(leaves.arange_dict) == [(0, K.B, "ttir", 0), (0, K.B, "ttir", 1)] + # the reader's tile shape; each dimension named by its lane + assert graph.accesses[0].shape == (K.B, K.B) + position = Encoder(graph, {}, {}, {}).lane_position(0) + assert [str(v) for v in position] == [f"ttir_lane_d{d}_n{K.B}" for d in (0, 1)] + + +def test_params_snapshots_and_unmodeled_values(): + graph = K.graph("scatter") + leaves = RaceLeaves(replace(graph), {"idx": 0}, {}) + # a launch value is read signed at the argument's width + strided = K.graph("strided") + assert ( + RaceLeaves(strided, {"stride": 2**32 - 1}, {}).param(Param("stride")).as_long() + == -1 + ) + assert ( + str(RaceLeaves(strided, None, {}).param(Param("stride"))) == "ttir_param_stride" + ) + with pytest.raises(RaceRefusal) as exc: + leaves.data_dep(DataDep("loaded value")) + assert exc.value.kind == "unmodeled-value" + # a snapshot is read at the load's width, signed, or not at all + load = graph.accesses[0] + for dtype, bits, first, problem in ( + (torch.int64, 32, 1, "read as i32"), + (torch.int8, 8, -1, None), + (torch.uint8, 8, 127, None), + (torch.uint8, 8, 128, "signed i8 reading"), + (torch.bool, 1, 1, None), + ): + t = torch.tensor([first, 0], dtype=dtype) + leaves = RaceLeaves(graph, {}, _launch(idx=t)) + access = replace(load, elem_bits=bits) + found = leaves.snapshot_problem(access) + assert (found is None) if problem is None else problem in found + if problem is None: + assert leaves.snapshot(access)[1][0] == first + + +_P, _NEG = Pid(0), Bin("-", Const(0), Pid(0)) + + +@pytest.mark.parametrize( + "term, stride, branches", + [ + (Bin("//", _P, Const(4)), 1, False), + (Bin("u%", _P, Const(-3)), 1, False), + (Bin("//", Arange("r", 0, 16, 0), Const(4)), 1, False), + (Bin("%", _P, Param("stride")), 0, True), # the remainder at b == 0 + (Bin("//", _NEG, Const(3)), 1, False), + (Bin("//", _NEG, Param("stride")), -3, False), + (Bin("%", _NEG, Const(-3)), 1, False), + (Bin("%", Bin("-", _P, Const(5)), Const(4)), 1, True), # sign unknown + ( + Bin("//", TSelect(Cmp("slt", _P, Const(2)), Const(3), _P), Const(2)), + 1, + False, + ), + ], +) +def test_division_drops_sign_branches_only_where_exact(term, stride, branches): + """RaceLeaves' ``//`` / ``%`` are Euclidean where the operands' bounds + make them the IR's truncating ones, and equal to the lowering's default + under the domains every query assumes.""" + leaves = RaceLeaves(K.graph("strided"), {"stride": stride}, {}) + bounded = leaves.value(term) + a, b = leaves.value(term.a), leaves.value(term.b) + exact = truncating_divrem(term, a, b) + # a branch beyond the operands' own + assert ( + str(bounded).count("If") > str(a).count("If") + str(b).count("If") + ) == branches + s = Solver() + s.add(pids.pid_var(0) >= 0, bounded != exact) + for (start, end, *_), (var, _) in leaves.arange_dict.items(): + s.add(And(var >= start, var < end)) + assert s.check() == unsat + + +def _module(*body: str, flag: str = "i32") -> str: + return ( + "module {\n tt.func public @k(%flag: !tt.ptr<" + + flag + + '> loc("flag"("k.py":1:1)), ' + '%out: !tt.ptr loc("out"("k.py":1:1))) attributes {noinline = false} {\n ' + "%c0 = arith.constant 0 : i32\n %c1 = arith.constant 1 : i32\n " + "%t = arith.constant true\n " + + "\n ".join(body) + + "\n tt.return\n }\n}\n" + ) + + +def _spin(*cond: str, body: str = "scf.yield", flag: str = "i32") -> str: + return _module("scf.while : () -> () {", *cond, "} do {", body, "}", flag=flag) + + +_POLL = "%o = tt.load %flag {isVolatile = true} : !tt.ptr" +_NE = "%c = arith.cmpi ne, %o, %c1 : i32" +_RMW = "%o = tt.atomic_rmw {op}, acquire, gpu, %flag, {v}, %t : (!tt.ptr, i32, i1) -> i32" + + +@pytest.mark.parametrize( + "text", + [ + _spin(_POLL, _NE, "scf.condition(%c)"), + _spin(_RMW.format(op="add", v="%c0"), _NE, "scf.condition(%c)"), + _spin(_RMW.format(op="xor", v="%c0"), _NE, "scf.condition(%c)"), + _spin( + "%o = tt.atomic_cas acquire, gpu, %flag, %c0, %c1 : (!tt.ptr, i32, i32) -> i32", + "%c = arith.cmpi ne, %o, %c0 : i32", + "scf.condition(%c)", + ), + ], +) +def test_a_spin_awaits_its_exit_value(text): + graph = parse_ttir(text, multipath=True, while_regions=True) + ((index, exit_pred),) = awaited_accesses(graph).items() + assert index == 0 and type(exit_pred).__name__ == "Not" + + +@pytest.mark.parametrize( + "text, why", + [ + ( + _module( + "%r = scf.while (%a = %c0) : (i32) -> i32 {", + _POLL, + _NE, + "scf.condition(%c) %o : i32", + "} do {", + "^bb0(%b: i32):", + "scf.yield %b : i32", + "}", + ), + "carries values", + ), + ( + _spin( + "scf.while : () -> () {", + _POLL, + _NE, + "scf.condition(%c)", + "} do {", + "scf.yield", + "}", + "scf.condition(%t)", + ), + "nested", + ), + ( + _spin( + _POLL, + _NE, + "scf.condition(%c)", + body="tt.store %out, %c1 : !tt.ptr\n scf.yield", + ), + "scf.yield only", + ), + ( + _spin(_POLL, "tt.store %out, %c1 : !tt.ptr", _NE, "scf.condition(%c)"), + "stores", + ), + ( + _spin(_POLL, "%p = tt.load %out : !tt.ptr", _NE, "scf.condition(%c)"), + "exactly one", + ), + ( + _spin( + "%o = tt.load %flag : !tt.ptr", + "%one = arith.constant 1.0 : f32", + "%c = arith.cmpf one, %o, %one : f32", + "scf.condition(%c)", + flag="f32", + ), + "float", + ), + ( + _spin(_RMW.format(op="add", v="%c1"), _NE, "scf.condition(%c)"), + "writes the awaited", + ), + ( + _spin(_RMW.format(op="exch", v="%c0"), _NE, "scf.condition(%c)"), + "writes the awaited", + ), + ( + _spin(_POLL, "%c = arith.trunci %o : i32 to i1", "scf.condition(%c)"), + "not a comparison", + ), + ( + _spin( + _POLL, + "%e = arith.addi %o, %c1 : i32", + "%c = arith.cmpi ne, %o, %e : i32", + "scf.condition(%c)", + ), + "does not depend on it", + ), + ], +) +def test_a_while_that_is_no_spin_is_spin_shape(text, why): + graph = parse_ttir(text, multipath=True, while_regions=True) + with pytest.raises(RaceRefusal) as exc: + awaited_accesses(graph) + assert exc.value.kind == "spin-shape" and why in exc.value.message + + +@pytest.mark.parametrize( + "reset, race", + [ + ("", None), + # a cta-scope reset of the flag races the failed polls (corner A) + ("%r = tt.atomic_rmw exch, relaxed, cta, %flag, %c0, %t", "WAW"), + ("%r = tt.atomic_rmw exch, relaxed, gpu, %flag, %c0, %t", None), + # a plain read of the flag races them too (corner B) + ("%f = tt.load %flag : !tt.ptr", "WAR"), + ], +) +def test_the_failed_polls_of_an_await_keep_their_races(reset, race): + """Program 0 writes, then releases a flag that program 1 awaits; what + program 0 does to the flag before the release meets the failed polls, + which only the pre-exit representative stands for.""" + if reset.startswith("%r"): + reset += " : (!tt.ptr, i32, i1) -> i32" + text = _module( + "%pid = tt.get_program_id x : i32", + "%isp = arith.cmpi eq, %pid, %c0 : i32", + "scf.if %isp {", + "tt.store %out, %c1 : !tt.ptr", + reset, + "ttg.barrier all", + "%x = tt.atomic_rmw exch, release, gpu, %flag, %c1, %t" + " : (!tt.ptr, i32, i1) -> i32", + "} else {", + "scf.while : () -> () {", + _RMW.format(op="add", v="%c0"), + _NE, + "scf.condition(%c)", + "} do {", + "scf.yield", + "}", + "ttg.barrier all", + "%v = tt.load %out : !tt.ptr", + "}", + ) + graph = parse_ttir(text, multipath=True, while_regions=True) + zero = torch.zeros(1, dtype=torch.int32) + out = analyze_graph(graph, {}, _launch(flag=zero, out=zero.clone()), (2, 1, 1)) + if race is None: + assert out.status == "proved" + return + (report,) = out.reports + assert report.race_type.value == race + assert report.second_record.debug_name.endswith(":pre-exit") + + +@pytest.mark.parametrize( + "space, ordered", + [ + ("all", True), + ("global_read|global_write", True), + ("local", False), + ("global_read", False), + ], +) +def test_only_a_global_barrier_orders_global_memory(space, ordered): + text = _module( + "%r = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32>", + "%s = tt.splat %out : !tt.ptr -> tensor<16x!tt.ptr>", + "%p = tt.addptr %s, %r : tensor<16x!tt.ptr>, tensor<16xi32>", + "tt.store %p, %r : tensor<16x!tt.ptr>", + f"ttg.barrier {space}", + "%q = tt.addptr %out, %c1 : !tt.ptr, i32", + "%v = tt.load %q : !tt.ptr", + ) + graph = parse_ttir(text, multipath=True, while_regions=True) + out = torch.zeros(16, dtype=torch.int32) + flag = torch.zeros(1, dtype=torch.int32) + outcome = analyze_graph(graph, {}, _launch(out=out, flag=flag), (1, 1, 1)) + assert (outcome.status == "proved") is ordered + + +def test_fence_epochs_count_the_barriers_before_an_execution(): + def epochs(name, params, **tensors): + encoder = Encoder(K.graph(name), params, _launch(**tensors), {}) + enc = encoder.encode() + out = [enc.fences.epoch(r) for r in enc.records] + return [e if isinstance(e, int) else simplify(e) for e in out] + + t = torch.zeros(4 * K.B, dtype=torch.int32) + store, load, _ = epochs("loop_buffer_fenced", {"n": 4}, t=t, u=t.clone()) + assert str(store) == "ttir_loop_k" and str(load) == "1 + ttir_loop_k" + before, after, _ = epochs( + "store_flag_fence_load", {"flag": 1}, hist=t, out=t.clone() + ) + assert before == 0 + assert "fence_epoch_unknown" in str(after) # behind a conditional barrier diff --git a/tests/unit/test_race_detector.py b/tests/unit/test_race_detector.py index bc079477..fd563758 100644 --- a/tests/unit/test_race_detector.py +++ b/tests/unit/test_race_detector.py @@ -5,7 +5,7 @@ from z3 import substitute import tilelens -from tilelens.clients import RaceDetector +from tilelens.clients import CompiledRaceDetector, RaceDetector from tilelens.clients.race_detector.race_detector import ( SymbolicRaceDetector, NullRaceDetector, @@ -31,11 +31,15 @@ def test_race_detector_factory_toggle(): cfg.enable_race_detector = saved -def test_compile_mode_is_not_available_yet(monkeypatch): +def test_compile_mode_is_the_compiled_race_detector(monkeypatch): monkeypatch.setattr(cfg, "enable_race_detector", True) - for make in (RaceDetector, SymbolicRaceDetector): - with pytest.raises(NotImplementedError, match="compiled race detector"): - make(compile=True) + det = RaceDetector(compile=True, target="cuda:90") + assert type(det) is CompiledRaceDetector and isinstance(det, RaceDetector) + assert (det.ir_target.arch, det.LAUNCH, det.IR_STAGES) == (90, "skip", {"ttir"}) + with pytest.raises(TypeError, match=r"RaceDetector\(compile=True\)"): + SymbolicRaceDetector(compile=True) + monkeypatch.setattr(cfg, "enable_race_detector", False) + assert isinstance(RaceDetector(compile=True), NullRaceDetector) @triton.jit diff --git a/tilelens/clients/__init__.py b/tilelens/clients/__init__.py index 54f21113..6167e7d8 100644 --- a/tilelens/clients/__init__.py +++ b/tilelens/clients/__init__.py @@ -12,6 +12,19 @@ "RaceType": ("tilelens.clients.race_detector.data", "RaceType"), "RaceRecord": ("tilelens.clients.race_detector.data", "RaceRecord"), "RaceAccess": ("tilelens.clients.race_detector.data", "RaceAccess"), + "CompiledRaceDetector": ( + "tilelens.clients.race_detector.compiled.client", + "CompiledRaceDetector", + ), + "CompiledRaceRecord": ( + "tilelens.clients.race_detector.data", + "CompiledRaceRecord", + ), + "RaceVerdict": ("tilelens.clients.race_detector.data", "RaceVerdict"), + "RaceConfigVerdict": ( + "tilelens.clients.race_detector.data", + "RaceConfigVerdict", + ), "Sanitizer": ("tilelens.clients.sanitizer.sanitizer", "Sanitizer"), "CompiledSanitizer": ( "tilelens.clients.sanitizer.compiled.client", diff --git a/tilelens/clients/pids.py b/tilelens/clients/pids.py new file mode 100644 index 00000000..df13bc9b --- /dev/null +++ b/tilelens/clients/pids.py @@ -0,0 +1,46 @@ +"""The program-id and grid variables every symbolic client names alike. + +One place defines how a program id and a grid size along an axis are +spelled as Z3 variables, so the eager engine (``SymbolicExpr.PID0-2``), the +compiled sanitizer and the race detector's compiled leaves speak one +vocabulary. Only the spelling is shared, not the context: ``pid_var`` and +``grid_var`` take the client's ``ctx`` (None: Z3's main context, where the +race detector's solver renames pids per program copy). How a client renames +or bounds these variables is its own policy. +""" + +from __future__ import annotations + +from z3 import ArithRef, Context, Int + +PID_AXES = (0, 1, 2) + + +def pid_name(axis: int) -> str: + """The name of the program id along ``axis``.""" + return f"pid_{axis}" + + +def pid_var(axis: int, ctx: Context | None = None) -> ArithRef: + """The program id along ``axis`` as a Z3 Int in ``ctx``.""" + return Int(pid_name(axis), ctx) + + +def grid_name(axis: int) -> str: + """The name of the grid size along ``axis``.""" + return f"grid_{axis}" + + +def grid_var(axis: int, ctx: Context | None = None) -> ArithRef: + """The grid size along ``axis`` as a Z3 Int in ``ctx``.""" + return Int(grid_name(axis), ctx) + + +def __getattr__(name: str) -> ArithRef: + """``PID0``, ``PID1``, ``PID2``: the program ids in Z3's main context, + made on first use (importing this module creates no Z3 term, so a + client in its own contexts sees Z3 exactly as before).""" + if name in ("PID0", "PID1", "PID2"): + value = globals()[name] = pid_var(int(name[-1])) + return value + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/tilelens/clients/race_detector/compiled/__init__.py b/tilelens/clients/race_detector/compiled/__init__.py new file mode 100644 index 00000000..cb0168b4 --- /dev/null +++ b/tilelens/clients/race_detector/compiled/__init__.py @@ -0,0 +1,8 @@ +"""The compiled race detector: race analysis over host-compiled TTIR. + +``leaves`` gives the TTIR reader's terms their meaning for the two-copy +solver (``RaceLeaves``), ``awaits`` and ``fences`` hold the race detector's +await and fence rules over the reader's neutral facts, ``global_records`` +lowers an ``AccessGraph`` to the solver's records, and ``analysis`` picks +the proof rung for one compiled kernel. +""" diff --git a/tilelens/clients/race_detector/compiled/analysis.py b/tilelens/clients/race_detector/compiled/analysis.py new file mode 100644 index 00000000..03348d4a --- /dev/null +++ b/tilelens/clients/race_detector/compiled/analysis.py @@ -0,0 +1,443 @@ +"""The proof rung of one compiled kernel: the least concretization that +decides it. + +* T0, any parameters and any grid along the axes the kernel reads: tried + when the launch's tensors are pairwise disjoint (one encoding per tensor + stands for non-aliased arguments) and every term is linear with the + params symbolic; any SAT, refusal or undecided query falls through. +* T1, this launch's parameters, any grid: first with every loaded value + free (the content-free attempt, when a term reads one: its UNSAT is a + proof for any contents, its SAT is no report), then with the launch's + snapshots (a proof that read one holds for these contents). +* T1-launch, this launch's parameters and grid: when the any-grid query + finds a race, the race pairs are asked again with the grid pinned; UNSAT + there proves the kernel as launched, and the any-grid reports are kept + as grid-fragility evidence, never as races. When Z3 does not decide that + requery, the config is undecided (``timeout`` / ``unknown``), with the + any-grid reports as evidence only. + +A proof needs a feasible base (some execution satisfies the premises), or +it would be vacuous. An any-grid proof assumes the width obligations hold +(``Outcome.overflow`` names the widths when some can fail in its domain); +they are discharged exactly at the launch (its parameters and grid), and a +proof whose obligations can fail there refuses as ``integer-overflow``, +whatever its rung: the unbounded reading it proved is not what this launch +computes. + +Z3 context: every term lives in Z3's main context, which is not +thread-safe; the caller serializes analyses (``MAIN_CONTEXT_LOCK``) and +drops the outcome's terms before releasing it. A Ctrl+C that cancels a +query raises KeyboardInterrupt. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +from z3 import Or, Solver, is_true, sat, unknown + +from ... import pids +from ....ir.ttir_reader import AccessGraph, Term +from ..hb_common import CountingScopeUnsupported, UnsupportedSymbolicRaceQuery +from ..two_copy_symbolic_hb_solver import TwoCopySymbolicHBSolver +from .awaits import awaited_accesses +from .global_records import ( + Encoder, + Encoding, + content_free_view, + graph_mentions_loaded, + grid_pins, + symbolic_grid, + t0_linearity_gate, +) +from .leaves import LaunchTensor, RaceKind, RaceRefusal + +T0_TIMEOUT_MS = 10_000 +T1_TIMEOUT_MS = 120_000 +# What Z3 answers for a query its Ctrl+C handler cancelled (Python's own +# handler does not run during the query). +_INTERRUPTED = "interrupted from keyboard" + + +@dataclass +class Outcome: + """What one kernel's analysis decided. ``status`` is "proved", "races" + (``reports`` exact, ``widened`` from over-approximated records), + "vacuous" (no execution satisfies the premises) or "unsupported" + (``refusal``).""" + + status: str + rung: str | None = None # "T0" | "T1" | "T1-launch" for a proof + # The verdict read the launch's tensor contents (a snapshot). + content: bool = False + # The kernel awaits: the verdict assumes its spin loops terminate. + assumes_termination: bool = False + # The widths whose obligations an any-grid proof assumes (some can fail + # in its domain). + overflow: tuple[int, ...] = () + reports: list[Any] = field(default_factory=list) + widened: list[Any] = field(default_factory=list) + # The any-grid reports behind a T1-launch proof. + grid_fragile: list[Any] = field(default_factory=list) + refusal: RaceRefusal | None = None + + +def _unsupported(r: RaceRefusal, termination: bool = False) -> Outcome: + # A caught refusal's traceback holds frames full of Z3 terms. + return Outcome( + "unsupported", refusal=r.with_traceback(None), assumes_termination=termination + ) + + +def _check_interrupt(reason: str | None) -> None: + if reason == _INTERRUPTED: + raise KeyboardInterrupt(f"Z3 query cancelled ({reason})") + + +def _undecided(e: UnsupportedSymbolicRaceQuery) -> RaceRefusal: + _check_interrupt(e.reason_unknown) + kind = RaceKind.TIMEOUT if e.reason_unknown == "timeout" else RaceKind.UNKNOWN + return RaceRefusal(kind, f"solver: {e}") + + +def disjoint_tensors(graph: AccessGraph, tensors: Mapping[str, LaunchTensor]) -> bool: + """Every accessed tensor has a verified allocation, and they are pairwise + disjoint: the non-aliasing the per-tensor T0 encoding assumes.""" + intervals = [] + for name in {a.base_param for a in graph.accesses}: + t = tensors.get(name) + interval = None if t is None else t.facts.allocation_interval() + if interval is None: + return False + intervals.append(interval) + intervals.sort() + return all(s2 >= e1 for (_, e1), (s2, _) in zip(intervals, intervals[1:])) + + +class _Analysis: + def __init__( + self, + graph: AccessGraph, + params: Mapping[str, int], + tensors: Mapping[str, LaunchTensor], + launch_grid: tuple[int, int, int] | None, + awaits: Mapping[int, Term], + t0_timeout_ms: int, + t1_timeout_ms: int, + ) -> None: + self.graph = graph + self.params = params + self.tensors = tensors + self.grid = launch_grid + self.awaits = awaits + self.t0_ms = t0_timeout_ms + self.t1_ms = t1_timeout_ms + self.termination = bool(awaits) + + def solver(self, enc: Encoding, timeout_ms: int, **kw: Any) -> Any: + kw.setdefault("extra_assumptions", enc.assumptions) + return TwoCopySymbolicHBSolver( + enc.records, + arange_dict=enc.arange_dict, + fence_between=enc.fences.between, + fence_epoch=enc.fences.epoch, + timeout_ms=timeout_ms, + **kw, + ) + + def proved(self, rung: str, enc: Encoding | None, overflow=()) -> Outcome: + return Outcome( + "proved", + rung=rung, + content=bool(enc and enc.content_qualified), + assumes_termination=self.termination, + overflow=tuple(overflow), + ) + + # ── width obligations ── + + def overflow( + self, encoder: Encoder, grid: tuple[Any, Any, Any], timeout_ms: int + ) -> tuple[int, ...] | None: + """The widths of the obligations that can fail in the domain of + ``grid`` (its symbolic dims free): () when none can, None when Z3 + cannot tell. Each width is asked until none is left that can fail + (one model need not show them all).""" + failing = encoder.violations + found: set[int] = set() + s = Solver() + s.set("timeout", timeout_ms) + s.add(*encoder.domain()) + s.add( + *(pids.grid_var(i) == d for i, d in enumerate(grid) if isinstance(d, int)) + ) + while True: + rest = [(bits, f) for bits, f in failing if bits not in found] + if not rest: + break + s.push() + s.add(Or(*(f for _, f in rest))) + result = s.check() + if result == unknown: + _check_interrupt(s.reason_unknown()) + return None + if result != sat: + break + model = s.model() + found |= {bits for bits, f in rest if is_true(model.eval(f, True))} + s.pop() + return tuple(sorted(found)) + + def any_grid_proof( + self, rung: str, encoder: Encoder, enc: Encoding, grid: Any, ms: int + ) -> Outcome | None: + """The proof at ``rung``, qualified by the widths whose obligations + can fail in its domain; those must hold at the launch (a T0 proof + is checked there with a launch encoding: None when it cannot be + built, and the next rung decides).""" + widths = self.overflow(encoder, grid, ms) + if widths is None: # undecided: assume every width the proof uses + widths = tuple(sorted({bits for bits, _ in encoder.violations})) + if widths and self.grid is not None: + at_launch = encoder + if rung == "T0": + try: + at_launch = Encoder( + self.graph, self.params, self.tensors, self.awaits + ) + at_launch.encode() + except RaceRefusal: + return None + refused = self.launch_widths(at_launch) + if refused is not None: + return refused + return self.proved(rung, enc, widths) + + def launch_widths(self, encoder: Encoder) -> Outcome | None: + """The refusal when a width obligation of ``encoder``'s records can + fail at the launch (its parameters and grid), or Z3 cannot tell; + None when they all hold.""" + assert self.grid is not None + widths = self.overflow(encoder, self.grid, self.t1_ms) + if widths is None: + return _unsupported( + RaceRefusal( + RaceKind.TIMEOUT, + "Z3 could not decide the width obligations at the launch", + ), + self.termination, + ) + if widths: + bits = "/".join(f"i{b}" for b in widths) + return _unsupported( + RaceRefusal( + RaceKind.INTEGER_OVERFLOW, + f"an {bits} value can overflow at this launch, where the " + "race-freedom proof needs the unbounded reading", + ), + self.termination, + ) + return None + + def launch_proof(self, encoder: Encoder, enc: Encoding, **kw: Any) -> Outcome: + """A proof at the launch's grid: its width obligations must hold + there exactly.""" + refused = self.launch_widths(encoder) + if refused is not None: + return refused + out = self.proved("T1-launch", enc) + out.grid_fragile = list(kw.get("grid_fragile", ())) + return out + + # ── rungs ── + + def run(self) -> Outcome: + if disjoint_tensors(self.graph, self.tensors) and t0_linearity_gate( + self.graph, self.awaits + ): + t0 = self.t0() + if t0 is not None: + return t0 + if graph_mentions_loaded(self.graph): + free = self.t1(blind=True) + if free is not None: + return free + out = self.t1(blind=False) + assert out is not None + return out + + def t0(self) -> Outcome | None: + """Every per-tensor group UNSAT over a feasible base, or None.""" + try: + encoder = Encoder(self.graph, None, {}, self.awaits) + groups: dict[str, list] = {} + for seq, access in enumerate(self.graph.accesses): + groups.setdefault(access.base_param, []).append( + (seq, content_free_view(access, self.graph)) + ) + last = None + for items in groups.values(): + if all(a.kind == "load" for _, a in items): + continue # reads never conflict + enc = encoder.encode(items) + solver = self.solver( + enc, self.t0_ms, grid=symbolic_grid(enc, self.grid, t0=True) + ) + if solver.find_races() or not solver.check_feasibility(): + return None + last = enc + except RaceRefusal: + return None + except UnsupportedSymbolicRaceQuery as e: + _check_interrupt(e.reason_unknown) + return None + if last is None: + return self.proved("T0", None) + return self.any_grid_proof( + "T0", encoder, last, symbolic_grid(last, self.grid, t0=True), self.t0_ms + ) + + def t1(self, blind: bool) -> Outcome | None: + """The T1 rungs; with ``blind`` (the content-free attempt) only a + proof, else None.""" + try: + encoder = Encoder( + self.graph, self.params, self.tensors, self.awaits, blind=blind + ) + enc = encoder.encode() + grid = symbolic_grid(enc, self.grid) + scoped = False + try: + solver = self.solver( + enc, self.t1_ms, grid=grid, enum_fallback_grid=self.grid + ) + except CountingScopeUnsupported: + if self.grid is None: + raise + # Retry every query at this launch, every counting guard kept. + solver = self.solver( + enc, + self.t1_ms, + grid=self.grid, + extra_assumptions=grid_pins(grid, self.grid) + enc.assumptions, + enum_fallback_grid=self.grid, + launch_ceiling=True, + ) + scoped = True + found = solver.find_races() + if found: + if blind: + return None + return self.races(encoder, enc, grid, found, scoped) + extra = () + if solver.enum_used: + assert self.grid is not None # the fallback needs a launch grid + extra = grid_pins(grid, self.grid) + feasible = solver.check_feasibility(extra=extra) + except RaceRefusal as r: + return None if blind else _unsupported(r, self.termination) + except UnsupportedSymbolicRaceQuery as e: + refusal = _undecided(e) # a Ctrl+C raises, blind or not + return None if blind else _unsupported(refusal, self.termination) + if not feasible: + if blind: + return None + return Outcome( + "vacuous", + refusal=RaceRefusal( + RaceKind.VACUOUS, + "the premises admit no execution, so the race-freedom " + "proof would be vacuous", + ), + assumes_termination=self.termination, + ) + out: Outcome | None + if scoped or solver.enum_used: + out = self.launch_proof(encoder, enc) + else: + out = self.any_grid_proof("T1", encoder, enc, grid, self.t1_ms) + assert out is not None # a T1 encoder is a launch encoding + # A refusal of the content-free attempt leaves the snapshots to try. + return None if blind and out.status != "proved" else out + + def races( + self, encoder: Encoder, enc: Encoding, grid: Any, found: list, scoped: bool + ) -> Outcome: + """Any-grid reports: asked again at the launch's grid, unless the + query already ran there. When that requery is not decided, the + config is undecided: the any-grid reports may lie outside the + launch's grid, so they are kept as evidence only (``grid_fragile``), + never as races of this launch.""" + if not scoped and self.grid is not None: + pinned = self.requery(enc, grid, found) + if isinstance(pinned, RaceRefusal): + out = _unsupported(pinned, self.termination) + out.grid_fragile = list(found) + return out + if not pinned: + return self.launch_proof(encoder, enc, grid_fragile=found) + found = pinned + out = Outcome( + "races", content=enc.content_qualified, assumes_termination=self.termination + ) + for rep in found: + ids = {rep.first.event_id, rep.second.event_id} + (out.widened if ids & enc.uncertain else out.reports).append(rep) + return out + + def requery(self, enc: Encoding, grid: Any, prior: list) -> list | RaceRefusal: + """The pairs behind ``prior`` at the launch's grid (a pair with no + any-grid report stays UNSAT under the pins); the refusal saying why + when Z3 does not decide it.""" + assert self.grid is not None + pairs = frozenset( + ( + min(r.first.event_id, r.second.event_id), + max(r.first.event_id, r.second.event_id), + ) + for r in prior + ) + try: + return self.solver( + enc, + self.t1_ms, + grid=grid, + extra_assumptions=grid_pins(grid, self.grid) + enc.assumptions, + only_pairs=pairs, + enum_fallback_grid=self.grid, + launch_ceiling=True, + ).find_races() + except UnsupportedSymbolicRaceQuery as e: + undecided = _undecided(e) # a Ctrl+C raises + return RaceRefusal( + undecided.kind, + f"the race query at the launch's grid was not decided " + f"({undecided.message}); the {len(prior)} race(s) found at " + "another grid are evidence only", + ) + + +def analyze_graph( + graph: AccessGraph, + params: Mapping[str, int], + tensors: Mapping[str, LaunchTensor], + launch_grid: tuple[int, int, int] | None, + *, + t0_timeout_ms: int = T0_TIMEOUT_MS, + t1_timeout_ms: int = T1_TIMEOUT_MS, +) -> Outcome: + """Decide one compiled kernel (a graph read with ``multipath=True`` and + ``while_regions=True``) for a launch's ``params``, ``tensors`` and + ``launch_grid`` (None: no launch rung).""" + try: + awaits = awaited_accesses(graph) + except RaceRefusal as r: + return _unsupported(r) + return _Analysis( + graph, params, tensors, launch_grid, awaits, t0_timeout_ms, t1_timeout_ms + ).run() + + +__all__ = ["Outcome", "analyze_graph", "disjoint_tensors"] diff --git a/tilelens/clients/race_detector/compiled/awaits.py b/tilelens/clients/race_detector/compiled/awaits.py new file mode 100644 index 00000000..3e9042e8 --- /dev/null +++ b/tilelens/clients/race_detector/compiled/awaits.py @@ -0,0 +1,101 @@ +"""The await abstraction: which ``scf.while`` loops the race detector reads +as a spin on one location. + +The reader reports each ``scf.while`` as a neutral ``WhileRegion``. A spin +``while load(flag) != 1: pass`` (or on a CAS or an identity RMW) is read as +ONE poll that observed its exit value: the loop terminates, and the last +read saw ``Not(cond)``. The failed iterations are dropped, and a pre-exit +representative of the poll (``global_records``) keeps the races they could +have. Anything else refuses as ``spin-shape``, naming the rule it breaks +(an ``scf.condition`` outside a while's condition region is the TTIR +reader's refusal, and a condition it cannot read breaks rule 8): + +1. the loop carries no values (no iter args or results); +2. it is not nested in another while (inside an scf.for it may be); +3. its body is ``scf.yield`` only; +4. its condition region writes nothing (no store); +5. its condition region reads exactly one location; +6. that location is not float-typed (observations are integers); +7. an RMW poll writes back what it read (``add`` / ``or`` / ``xor`` of 0); +8. the condition is a comparison; +9. of the poll's observation against a value that does not depend on it. +""" + +from __future__ import annotations + +from ....ir.ttir_reader import ( + AccessGraph, + Cmp, + Const, + Not, + Observed, + Term, + WhileRegion, + observed_indices, +) +from .leaves import RaceKind, RaceRefusal + +_IDENTITY_RMW = ("add", "or", "xor") + + +def awaited_accesses(graph: AccessGraph) -> dict[int, Term]: + """The awaited polls of ``graph``: access index -> exit predicate. + Raises a ``spin-shape`` RaceRefusal for a while that is not a spin.""" + for w in graph.whiles: # outer first, as the loops open + if w.carries: + raise _refuse(w, "the loop carries values (iter args or results)") + if w.parent is not None: + raise _refuse( + w, "a spin nested in another scf.while is not the await shape" + ) + return dict(_spin(graph, w) for w in graph.whiles) + + +def _refuse(w: WhileRegion, why: str) -> RaceRefusal: + return RaceRefusal( + RaceKind.SPIN_SHAPE, + f"line {w.line_no} (scf.while): {why}", + line_no=w.line_no, + loc=w.loc, + ) + + +def _spin(graph: AccessGraph, w: WhileRegion) -> tuple[int, Term]: + def refuse(why: str) -> RaceRefusal: + return _refuse(w, why) + + extra = [name for name in w.after_ops if name != "scf.yield"] + if extra: + raise refuse(f"the loop body must be scf.yield only, found {extra[0]}") + reads = [graph.accesses[i] for i in w.before] + if any(a.kind == "store" for a in reads): + raise refuse("the spin condition stores") + if len(reads) != 1: + raise refuse( + "the spin condition must read exactly one location " + f"(found {len(reads)} accesses)" + ) + index, poll = w.before[0], reads[0] + if poll.elem_float: + raise refuse("the awaited location is float-typed") + if poll.kind == "atomic_rmw": + op = poll.atomic.rmw_op if poll.atomic is not None else None + if op not in _IDENTITY_RMW or poll.atomic_val != Const(0): + # A mutating poll writes on every dropped iteration: the loop + # could end by observing its own writes. + raise refuse( + f"the spin re-read writes the awaited location (atomic {op} " + "with a nonzero operand)" + ) + cond = w.cond + if not isinstance(cond, Cmp): + raise refuse("the spin condition is not a comparison over the awaited read") + own = Observed(index) + a_is, b_is = cond.a == own, cond.b == own + expected = cond.b if a_is else cond.a + if a_is == b_is or index in observed_indices(expected, graph): + raise refuse( + "the spin condition must compare the awaited read against a value " + "that does not depend on it" + ) + return index, Not(cond) diff --git a/tilelens/clients/race_detector/compiled/client.py b/tilelens/clients/race_detector/compiled/client.py new file mode 100644 index 00000000..7c04eb0b --- /dev/null +++ b/tilelens/clients/race_detector/compiled/client.py @@ -0,0 +1,674 @@ +"""The compiled race detector client (``RaceDetector(compile=True)``). + +An ``IRClient`` that reads each compiled config's TTIR (with the reader's +multipath and while-region facts) and decides, with ``analysis.analyze_graph``, +whether two program instances of the launch can race on global memory. The +kernel is not run (``LAUNCH = "skip"``), so output tensors are left untouched; +it is compiled on the host for the client's target (``target=``, else +``TILELENS_IR_TARGET``, else ``cuda:89``), so no GPU is needed. + +Every config the launch compiled is decided, autotune benchmark configs +included, each against the bindings the core delivered for it. Configs whose +TTIR (minus its source locations), scalar arguments, tensor facts and +contents, and grid are the same (e.g. ones that differ only in ``num_warps``) +share one solve. A config's verdict is a ``RaceConfigVerdict``: "ok" with the +proof's rung and qualifiers, "races", "unsupported" (a refusal, a Z3 timeout +or unknown, or only possible races under over-approximation) or "vacuous". +The launch's ``RaceVerdict`` is the union: "races" if a config races, else +"unsupported", else "vacuous", else "ok". No launch-level time budget: each +solver has its own timeout (``T0_TIMEOUT_MS``, ``T1_TIMEOUT_MS``). + +A config that failed to compile is classified by tilelens.ir.failures, as in +the compiled sanitizer: an error no target compiles past is only a note, any +other one makes the config unsupported (``"compile-failed"``), and a host +compile that could not run is ``"host-compile-unavailable"``. + +A loaded integer value is read from the launch's pre-launch contents of its +tensor (a snapshot, taken in ``before_launch``): only for an integer or bool, +contiguous tensor of at most ``SNAPSHOT_MAX_ELEMENTS`` elements, and while the +launch's snapshots stay under ``SNAPSHOT_CAP_BYTES``. A verdict that read one +holds for those contents ("reads tensor contents"). + +What a verdict covers: this launch as compiled for the target. Since no +kernel runs, a later launch whose arguments the host computes from an earlier +one's outputs is analyzed with the values the untraced program would not +have used; likewise an autotune ``Config.pre_hook`` does not run, so a +verdict that reads tensor contents read them as the launch began, not as a +config's pre_hook would have left them. + +Each race found becomes a ``CompiledRaceRecord`` (plain data: no Z3 terms, no +tensors), followed by the launch's ``RaceVerdict`` in ``Launch.records``; +``tilelens.save()`` keeps both. Every race is printed as the launch +finalizes, as the eager race detector prints them; under ``verbose`` or the +tile-race CLI, so is each part of a launch that was not decided (once per +specialization, kind and site) and a race-free launch's claim. + +The analysis runs in Z3's main context, which is not thread-safe: solves are +serialized process-wide (``hb_common.MAIN_CONTEXT_LOCK``), so traced +launches on several host threads take turns. +""" + +from __future__ import annotations + +from collections.abc import Hashable, Mapping +from dataclasses import dataclass, replace +from typing import Any, ClassVar + +from ....core.client import LaunchCall, LaunchEvent +from ....core.config import config as cfg +from ....core.host_compile import parse_ir_target +from ....ir.capture import ( + ArtifactLog, + CompiledSpecialization, + ParseCache, + configs_of, + content_key, + saveable_config, +) +from ....ir.client import IRClient +from ....ir.failures import ( + compiles_nowhere, + failure_refusal, + none_compiled, + nothing_compiled, + nowhere_note, +) +from ....ir.launch import LaunchBinding, TensorFacts +from ....ir.verdict import Refusal +from ..hb_common import MAIN_CONTEXT_LOCK +from ..data import ( + CompiledRaceRecord, + RaceAccess, + RaceConfigVerdict, + RaceReport, + RaceType, + RaceVerdict, +) +from ..race_detector import RaceDetector, print_race_record +from .analysis import T0_TIMEOUT_MS, T1_TIMEOUT_MS, Outcome, analyze_graph +from .leaves import LaunchTensor, RaceKind, RaceRefusal + +# How to analyze a kernel for another target, as a refusal says it. +_NAME_A_TARGET = "RaceDetector(compile=True, target=...), or TILELENS_IR_TARGET" + +# What a proof at each rung holds for, weakest last. +_SCOPES = { + "T0": "any-params-any-grid", + "T1": "this-params-any-grid", + "T1-launch": "this-params-this-grid", +} +_STATUS_ORDER = ("races", "unsupported", "vacuous", "ok") + + +@dataclass(frozen=True) +class _Solved: + """One solve's result, shared by the configs whose key it has.""" + + status: str + rung: str | None = None + qualifiers: tuple[str, ...] = () + refusal: Refusal | None = None + races: tuple[CompiledRaceRecord, ...] = () + grid_fragile: tuple[CompiledRaceRecord, ...] = () + # Possible races under over-approximation, withheld beside the races. + withheld: int = 0 + + +class CompiledRaceDetector(IRClient): + """Global-memory race detection over a kernel's compiled TTIR. + + ``last_reports`` (the last launch's races), ``last_verdict`` and + ``last_status`` (its status, None before a launch is finalized) are the + compatibility view of the last launch. + """ + + NAME = "compiled_race_detector" + LOG_TAG: ClassVar[str] = "CompiledRaceDetector" + IR_STAGES: ClassVar[frozenset[str]] = frozenset({"ttir"}) + LAUNCH = "skip" + + # Per-solver Z3 timeouts of the any-parameters rung and of the others. + T0_TIMEOUT_MS: ClassVar[int] = T0_TIMEOUT_MS + T1_TIMEOUT_MS: ClassVar[int] = T1_TIMEOUT_MS + # The largest tensor whose contents are read, and the launch's total. + SNAPSHOT_MAX_ELEMENTS: ClassVar[int] = 16384 + SNAPSHOT_CAP_BYTES: ClassVar[int] = 256 * 1024 * 1024 + + def __init__(self, abort_on_error: bool = False, *, target: Any = None) -> None: + """``target``: what the kernels are compiled for (see + tilelens.core.host_compile.parse_ir_target); None for the configured + default. ``abort_on_error`` is the RaceDetector factory's, unused.""" + super().__init__() + self.ir_target = None if target is None else parse_ir_target(target) + self.abort_on_error = abort_on_error + self.parses = ParseCache() + self.last_reports: list[CompiledRaceRecord] = [] + # id(binding) -> (binding, its tensors with their snapshots); holding + # the binding keeps its id unique. + self._contents: dict[int, tuple[LaunchBinding, dict[str, LaunchTensor]]] = {} + self._kernel = "kernel" + # What the launch's snapshots may still take (SNAPSHOT_CAP_BYTES). + self._budget = self.SNAPSHOT_CAP_BYTES + # The undecided parts already printed (see print_undecided). + self._printed: set[Hashable] = set() + + @property + def last_status(self) -> str | None: + verdict = self.last_verdict + return None if verdict is None else verdict.status + + # ── launch lifecycle ───────────────────────────────────────────── + + def begin_launch(self, call: LaunchCall) -> None: + super().begin_launch(call) + self.last_reports = [] + self._contents = {} + self._budget = self.SNAPSHOT_CAP_BYTES + self._kernel = getattr(call.jit_fn, "__name__", None) or "kernel" + + def abort_launch(self, exc: BaseException) -> None: + super().abort_launch(exc) + self._contents = {} + + def before_launch(self, event: LaunchEvent) -> None: + super().before_launch(event) + # The snapshots are taken for the binding just recorded. + binding = self.artifacts.last_binding + assert binding is not None + self._contents[id(binding)] = (binding, self._snapshots(event, binding)) + + def finalize(self) -> list: + try: + out = super().finalize() + finally: + self._contents = {} + verdict = self.last_verdict + assert isinstance(verdict, RaceVerdict) + for record in self.last_reports: + print_race_record(record) + # What was not analyzed always prints (once per site), so a quiet + # launch never reads as race-free. A proof prints where someone + # asked to see it: verbose, or tile-race, which reports every launch. + print_undecided(verdict, self._printed) + if verdict.status == "ok" and (cfg.verbose or cfg.cli_active): + print(f"[{self.LOG_TAG}] {self._kernel}: race-free{_claim(verdict)}") + return out + + # ── snapshots ──────────────────────────────────────────────────── + + def _snapshots( + self, event: LaunchEvent, binding: LaunchBinding + ) -> dict[str, LaunchTensor]: + """Each tensor argument's facts and, where it can be read, its + pre-launch contents (no kernel of the launch has run yet).""" + out: dict[str, LaunchTensor] = {} + for name, facts in binding.tensors.items(): + value = _described(event.bound_args.get(name)) + values: tuple[int, ...] | None = None + nbytes = facts.numel * facts.elem_size + try: + why = _unreadable(value, facts, self.SNAPSHOT_MAX_ELEMENTS) + if not why and nbytes > self._budget: + why = ( + f"over the launch's {self.SNAPSHOT_CAP_BYTES}-byte snapshot cap" + ) + if not why: + values = tuple( + int(v) for v in value.detach().reshape(-1).cpu().tolist() + ) + self._budget -= nbytes + except Exception as exc: + values, why = None, f"capture failed: {type(exc).__name__}" + out[name] = LaunchTensor(facts, values, why) + return out + + # ── the analysis ───────────────────────────────────────────────── + + def analyze_launch(self, log: ArtifactLog) -> tuple[list, RaceVerdict]: + failed: list[RaceConfigVerdict] = [] + notes: list[str] = [] + for failure in log.failures: + if compiles_nowhere(failure.error): + notes.append(nowhere_note(failure, saveable_config(failure.config))) + else: + failed.append( + RaceConfigVerdict( + None, + saveable_config(failure.config), + "unsupported", + failure_refusal(failure, _NAME_A_TARGET), + ) + ) + specializations = log.specializations + if not specializations: + if failed: + refusal = failed[0].refusal + elif log.failures: + refusal = none_compiled(log) + else: + refusal = Refusal( + RaceKind.NO_COMPILED_KERNEL, + nothing_compiled(log, "RaceDetector()"), + ) + return [], RaceVerdict( + self.NAME, + "unsupported", + refusal=refusal, + per_config=tuple(failed), + notes=tuple(notes), + ) + solved: dict[Hashable, _Solved] = {} + records: list[CompiledRaceRecord] = [] + per_config: list[RaceConfigVerdict] = [] + for spec in specializations: + for found, verdict, withheld in self._analyze_specialization(spec, solved): + records += found + per_config.append(verdict) + if withheld: + notes.append( + f"config {verdict.config}: {withheld} more possible " + "race(s) under over-approximation were withheld" + ) + per_config += failed + status = _union(c.status for c in per_config) + refusal = next((c.refusal for c in per_config if c.refusal is not None), None) + rung = _weakest(c.rung for c in per_config) + scope = _SCOPES[rung] if status == "ok" and rung is not None else None + self.last_reports = records + return records, RaceVerdict( + self.NAME, + status, + scope=scope, + refusal=refusal, + per_config=tuple(per_config), + notes=tuple(notes), + qualifiers=_merged(c.qualifiers for c in per_config), + ) + + def _analyze_specialization( + self, spec: CompiledSpecialization, solved: dict[Hashable, _Solved] + ) -> list[tuple[list[CompiledRaceRecord], RaceConfigVerdict, int]]: + configs = configs_of(spec) + try: + graph = self._graph_of(spec) + except Exception as exc: + graph = Refusal(RaceKind.INTERNAL_ERROR, _describe(exc)) + if isinstance(graph, Refusal): + return [ + ( + [], + RaceConfigVerdict( + spec.specialization, config, "unsupported", graph + ), + 0, + ) + for config, _ in configs + ] + text = _without_locs(spec.artifacts.stages["ttir"]) + return [ + self._analyze_config( + spec.specialization, graph, text, config, bindings, solved + ) + for config, bindings in configs + ] + + def _graph_of(self, spec: CompiledSpecialization) -> Any: + """The access graph read from ``spec``'s TTIR, or the Refusal saying + why there is none.""" + text = spec.artifacts.stages.get("ttir") + if not isinstance(text, str): + why = spec.artifacts.error or "the compiled kernel holds no TTIR" + return Refusal(RaceKind.NO_TTIR, why) + outcome = self.parses.get(text, multipath=True, while_regions=True) + if outcome.refusal is not None: + refused = Refusal.from_exception(outcome.refusal) + where = _site(refused.line_no, refused.loc) + return replace(refused, message=f"{where}: {refused.message}") + if outcome.error is not None: + return Refusal( + RaceKind.INTERNAL_ERROR, f"the TTIR reader failed: {outcome.error}" + ) + return outcome.graph + + def _analyze_config( + self, + specialization: Hashable, + graph: Any, + text: str, + config: dict[str, Any], + bindings: list[LaunchBinding], + solved: dict[Hashable, _Solved], + ) -> tuple[list[CompiledRaceRecord], RaceConfigVerdict, int]: + if not bindings: + refusal = Refusal( + RaceKind.INTERNAL_ERROR, + "no binding was delivered for this kernel, so nothing was analyzed", + ) + return ( + [], + RaceConfigVerdict(specialization, config, "unsupported", refusal), + 0, + ) + results: list[_Solved] = [] + for binding in bindings: + contents = self._contents_of(binding) + key = (text, _binding_key(binding, contents)) + hit = solved.get(key) + if hit is None: + hit = solved[key] = self._solve(graph, binding, contents) + results.append(hit) + status = _union(r.status for r in results) + races = [replace(r, config=config) for s in results for r in s.races] + verdict = RaceConfigVerdict( + specialization, + config, + status, + next((r.refusal for r in results if r.refusal is not None), None), + n_reports=len(races), + rung=_weakest(r.rung for r in results) if status == "ok" else None, + qualifiers=_merged(r.qualifiers for r in results), + grid_fragile=tuple( + replace(r, config=config) for s in results for r in s.grid_fragile + ), + ) + return races, verdict, sum(r.withheld for r in results) + + def _contents_of(self, binding: LaunchBinding) -> dict[str, LaunchTensor]: + entry = self._contents.get(id(binding)) + if entry is not None and entry[0] is binding: + return entry[1] + # A binding no before_launch of this client recorded: facts only. + return { + name: LaunchTensor(facts, None, "no snapshot was taken") + for name, facts in binding.tensors.items() + } + + def _solve( + self, graph: Any, binding: LaunchBinding, contents: dict[str, LaunchTensor] + ) -> _Solved: + """One analysis, in Z3's main context: serialized with every other + race-detector solve of the process, and plain data by the time the + lock is released (no Z3 term outlives it, not even in a + traceback).""" + with MAIN_CONTEXT_LOCK: + try: + out = analyze_graph( + graph, + binding.params, + contents, + binding.grid, + t0_timeout_ms=self.T0_TIMEOUT_MS, + t1_timeout_ms=self.T1_TIMEOUT_MS, + ) + try: + return _settle(out, binding) + finally: + del out + except Exception as exc: + # A bug in this config's analysis; the other configs still + # count. + exc.__traceback__ = None + return _Solved( + "unsupported", + refusal=Refusal(RaceKind.INTERNAL_ERROR, _describe(exc)), + ) + + def on_analysis_error(self, exc: Exception) -> RaceVerdict: + self.last_reports = [] + return RaceVerdict( + self.NAME, + "unsupported", + refusal=Refusal(RaceKind.INTERNAL_ERROR, _describe(exc)), + ) + + +# ─────────────────────────── helpers ─────────────────────────── + + +def _describe(exc: BaseException) -> str: + return f"{type(exc).__name__}: {exc}" + + +def _site(line_no: int | None, loc: Any) -> str: + if loc is not None: + return f"{loc.file}:{loc.line}" + return f"TTIR line {line_no}" if line_no is not None else "the kernel" + + +def _described(value: Any) -> Any: + # A descriptor-style argument is bound as its .base tensor (see + # tilelens.ir.launch.bind_launch). + base = getattr(value, "base", None) + if base is not None and hasattr(base, "data_ptr"): + return base + return value + + +def _unreadable(value: Any, facts: TensorFacts, max_elements: int) -> str: + """Why ``value``'s contents are not read ("" when they are).""" + dtype = getattr(value, "dtype", None) + if dtype is None or not hasattr(value, "detach"): + return "not a torch tensor" + if getattr(dtype, "is_floating_point", True) or getattr(dtype, "is_complex", False): + return f"float dtype {facts.dtype}" + if any(getattr(value, flag, lambda: False)() for flag in ("is_conj", "is_neg")): + return "lazy value transform" + if not facts.contiguous: + return "non-contiguous" + if facts.numel == 0: + return "empty" + if facts.numel > max_elements: + # "too large" is what the address refusal keys its kind on. + return f"too large ({facts.numel} elements, bound {max_elements})" + return "" + + +def _strip_locs(line: str) -> str: + out: list[str] = [] + i = 0 + while True: + j = line.find("loc(", i) + if j < 0: + out.append(line[i:]) + return "".join(out) + if j > 0 and (line[j - 1].isalnum() or line[j - 1] in "_#.@%$"): + out.append(line[i : j + 4]) + i = j + 4 + continue + out.append(line[i:j].rstrip()) + depth, k, quoted = 0, j + 3, False + while k < len(line): + c = line[k] + if quoted: + if c == "\\": + k += 1 + elif c == '"': + quoted = False + elif c == '"': + quoted = True + elif c == "(": + depth += 1 + elif c == ")": + depth -= 1 + if depth == 0: + break + k += 1 + i = k + 1 + + +def _without_locs(text: str) -> str: + """The content key of ``text`` without its source locations (configs of + one kernel print the same ones).""" + lines = (_strip_locs(line) for line in text.splitlines()) + return content_key("\n".join(ln for ln in lines if not ln.startswith("#loc"))) + + +def _binding_key( + binding: LaunchBinding, contents: Mapping[str, LaunchTensor] +) -> Hashable: + """What the analysis reads of a binding.""" + return ( + tuple(sorted(binding.params.items())), + tuple(sorted((n, t.facts, t.values, t.why) for n, t in contents.items())), + binding.grid, + binding.error, + ) + + +def _union(statuses: Any) -> str: + seen = set(statuses) + return next(s for s in _STATUS_ORDER if s in seen) + + +def _weakest(rungs: Any) -> str | None: + order = list(_SCOPES) + found = [r for r in rungs if r is not None] + return max(found, key=order.index) if found else None + + +def _merged(groups: Any) -> tuple[str, ...]: + out: list[str] = [] + for group in groups: + out += [q for q in group if q not in out] + return tuple(out) + + +def _qualifiers(out: Outcome) -> tuple[str, ...]: + found: list[str] = [] + if out.overflow: + widths = "/".join(f"i{bits}" for bits in out.overflow) + found.append(f"assumes no {widths} overflow") + if out.assumes_termination: + found.append("assumes termination") + if out.content: + found.append("reads tensor contents") + return tuple(found) + + +def _settle(out: Outcome, binding: LaunchBinding) -> _Solved: + # A qualifier says what a verdict assumed: an undecided one has none. + qualifiers = _qualifiers(out) if out.status in ("proved", "races") else () + if out.status == "proved": + return _Solved( + "ok", + rung=out.rung, + qualifiers=qualifiers, + grid_fragile=tuple(_record(r, binding) for r in out.grid_fragile), + ) + if out.status == "races": + races = tuple(_record(r, binding) for r in out.reports) + if races: + return _Solved( + "races", qualifiers=qualifiers, races=races, withheld=len(out.widened) + ) + first = _record(out.widened[0], binding) + return _Solved( + "unsupported", + refusal=Refusal( + RaceKind.POTENTIAL_RACE, + f"{len(out.widened)} possible race(s) under over-approximation " + "(a dropped mask or a loaded value with no snapshot), not a " + f"witness; the first: {_pair(first)}", + ), + ) + refusal = out.refusal or RaceRefusal(RaceKind.OTHER, "no verdict") + return _Solved( + out.status, + refusal=Refusal.from_exception(refusal), + grid_fragile=tuple(_record(r, binding) for r in out.grid_fragile), + ) + + +def _pair(record: CompiledRaceRecord) -> str: + def side(access: RaceAccess) -> str: + where = access.source_location + line = f" (line {where[1]})" if where is not None else "" + return f"{access.op} of {access.tensor_name}{line}" + + return f"{record.race_type} between {side(record.first)} and {side(record.second)}" + + +def _record(report: RaceReport, binding: LaunchBinding) -> CompiledRaceRecord: + race_type = report.race_type + writes = ( + race_type in (RaceType.WAW, RaceType.RAW), + race_type in (RaceType.WAW, RaceType.WAR), + ) + assert report.witness_addrs is not None + accesses = [] + for event, event_writes, addr in zip( + (report.first, report.second), writes, report.witness_addrs + ): + record = event.record + facts = binding.tensors.get(record.tensor_name) + element = ( + None + if facts is None + else (addr - facts.data_ptr) // max(1, record.elem_size) + ) + accesses.append( + RaceAccess( + op=record.op_type.name, + writes=event_writes, + tensor_name=record.tensor_name, + source_location=record.source_location, + element=element, + ) + ) + assert report.witness_grid_a is not None and report.witness_grid_b is not None + return CompiledRaceRecord( + race_type=race_type, + first=accesses[0], + second=accesses[1], + programs=(report.witness_grid_a, report.witness_grid_b), + reason=report.reason, + ) + + +def _claim(verdict: RaceVerdict) -> str: + parts = [verdict.scope or "", *verdict.qualifiers] + return f" ({'; '.join(p for p in parts if p)})" + + +# ─────────────────────────── reporting ─────────────────────────── + + +def print_undecided(verdict: RaceVerdict, printed: set[Hashable]) -> None: + """Print one line per part of a launch the compiled race detector did + not decide (each refused config, else the launch's own refusal) that is + not in ``printed`` yet, and add it there: once per specialization, kind + and site; a config that compiled to nothing once per config and + message. The launch's notes follow, each once.""" + refused = [ + (c.specialization, c.config, c.refusal) + for c in verdict.per_config + if c.refusal is not None and c.status != "races" + ] + if not refused and verdict.refusal is not None and verdict.status != "races": + refused = [(None, {}, verdict.refusal)] + for specialization, config, refusal in refused: + ident: Hashable = specialization + if specialization is None: + ident = ( + tuple(sorted((name, repr(value)) for name, value in config.items())), + refusal.message, + ) + key = (ident, refusal.kind, refusal.line_no, refusal.loc) + if key in printed: + continue + printed.add(key) + where = f" (config {config})" if config else "" + print( + f"[{CompiledRaceDetector.LOG_TAG}] not analyzed{where}: " + f"{refusal.kind}: {refusal.message}" + ) + for note in verdict.notes: + if ("note", note) in printed: + continue + printed.add(("note", note)) + print(f"[{CompiledRaceDetector.LOG_TAG}] note: {note}") + + +# A race detector mode like the eager one: isinstance(..., RaceDetector) +# holds, so the factory's ENABLE_RACE_DETECTOR=0 rule reads it alike. +RaceDetector.register(CompiledRaceDetector) diff --git a/tilelens/clients/race_detector/compiled/fences.py b/tilelens/clients/race_detector/compiled/fences.py new file mode 100644 index 00000000..d8f287b8 --- /dev/null +++ b/tilelens/clients/race_detector/compiled/fences.py @@ -0,0 +1,155 @@ +"""The compiled race detector's fence rule, the two-copy solver's +``fence_between`` and ``fence_epoch`` hooks over the reader's Fence facts. + +A ``ttg.barrier`` (``tl.debug_barrier()``) orders global memory only when +its ``addr_space`` has both the global_read (2) and global_write (4) bits. +Such a fence ``f`` orders two accesses ``a < f < b`` of one program +instance when its block path is a prefix of both accesses' block paths: +then it runs whenever both do (it is not under a branch either one escapes, +and not in a loop either one is outside of). In a loop, ``f`` orders the +two accesses of one iteration; ``fence_between`` (the solver asks it for +every pair of executions) holds only when every pair of executions is +ordered, so a fence in a loop also needs one on the back edge: another +fence of the loop's body outside ``[a, b]``. The pairs of different +iterations that a single fence leaves unordered the solver queries at +equal fence epochs (``epoch``): the number of ordering fences the instance +has run before the execution, exact where every fence that may run before +it is unconditional and the loop trip counts it depends on are constants, +else unconstrained. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from z3 import ArithRef, Int, IntVal + +from ....ir.ttir_reader import AccessEvent, AccessGraph, Fence +from ..data import AccessEventRecord + +_GLOBAL = 2 | 4 # global_read | global_write + + +def _prefix(p: tuple[int, ...], q: tuple[int, ...]) -> bool: + return len(p) <= len(q) and q[: len(p)] == p + + +def _common(p: tuple[str, ...], q: tuple[str, ...]) -> tuple[str, ...]: + n = 0 + while n < len(p) and n < len(q) and p[n] == q[n]: + n += 1 + return p[:n] + + +class FenceRule: + """The rule for one graph. ``loops`` maps a loop name to its + iteration-index variable and its trip count (None when not a constant + for the encoding).""" + + def __init__( + self, graph: AccessGraph, loops: Mapping[str, tuple[ArithRef, int | None]] + ) -> None: + self.graph = graph + self.loops = loops + self.fences = [f for f in graph.fences if f.addr_space & _GLOBAL == _GLOBAL] + # Fences inside an scf.while run a data-dependent number of times. + self._in_while = [ + any( + _prefix(w.block_path + (block,), f.block_path) + for w in graph.whiles + for block in w.blocks + ) + for f in self.fences + ] + self._epochs: dict[int, Any] = {} + + def _access(self, record: AccessEventRecord) -> AccessEvent: + return self.graph.accesses[record.program_seq] + + def between(self, first: AccessEventRecord, second: AccessEventRecord) -> bool: + """Every execution of ``first`` is ordered with every execution of + ``second`` by a fence between them (``first`` comes first in program + order; in a loop, a later iteration's ``first`` comes after an + earlier one's ``second``).""" + a, b = self._access(first), self._access(second) + lo, hi = first.program_seq, second.program_seq + shared = [ + f + for f in self.fences + if _prefix(f.block_path, a.block_path) + and _prefix(f.block_path, b.block_path) + ] + for f in shared: + if not lo < f.position < hi: + continue + if not f.loops: + return True + # In a loop: a back-edge fence of the same loop body orders the + # later iterations' executions of ``first`` after ``second``. + if any( + g is not f + and g.loops[: len(f.loops)] == f.loops + and not lo < g.position < hi + for g in shared + ): + return True + return False + + def epoch(self, record: AccessEventRecord) -> Any: + """The number of ordering fences the instance ran before this + execution of ``record``: an int, a term over its loops' iteration + indices, or an unconstrained value where that is not exact.""" + seq = record.program_seq + if seq not in self._epochs: + self._epochs[seq] = self._epoch(seq) + return self._epochs[seq] + + def _epoch(self, seq: int) -> Any: + access = self.graph.accesses[seq] + own = access.loops + total: Any = 0 + for f, in_while in zip(self.fences, self._in_while): + count = self._count(f, in_while, seq, own) + if count is None: + return Int(f"fence_epoch_unknown_{seq}") + total = count if isinstance(total, int) and total == 0 else total + count + return total + + def _count( + self, f: Fence, in_while: bool, seq: int, own: tuple[str, ...] + ) -> Any | None: + """How many times ``f`` ran before the execution of access ``seq`` + (None: not exact).""" + common = _common(f.loops, own) + before = f.position < seq + if not common and not before: + return 0 # runs only after the access, in no loop they share + if in_while or f.guarded or f.path is not None: + return None + # Executions of f per iteration of the innermost common loop. + per = 1 + for name in f.loops[len(common) :]: + trips = self._trips(name) + if trips is None: + return None + per *= trips + # Complete iterations of the innermost common body before this one. + earlier: Any = 0 + for i, name in enumerate(common): + weight = 1 + for inner in common[i + 1 :]: + trips = self._trips(inner) + if trips is None: + return None + weight *= trips + k = self.loops[name][0] + earlier = k * weight if i == 0 else earlier + k * weight + current = per if before else 0 + if isinstance(earlier, int): + return earlier * per + current + return earlier * IntVal(per) + current if per != 1 else earlier + current + + def _trips(self, name: str) -> int | None: + found = self.loops.get(name) + return None if found is None else found[1] diff --git a/tilelens/clients/race_detector/compiled/global_records.py b/tilelens/clients/race_detector/compiled/global_records.py new file mode 100644 index 00000000..3e654753 --- /dev/null +++ b/tilelens/clients/race_detector/compiled/global_records.py @@ -0,0 +1,1109 @@ +"""AccessGraph -> the two-copy solver's records (global memory). + +``Encoder`` lowers every access of the TTIR reader's ``AccessGraph`` into +the ``AccessEventRecord`` shape the eager race detector gives the +``TwoCopySymbolicHBSolver``, through ``RaceLeaves``: + +* ``addr_expr`` is a byte address: ``data_ptr + offset * elem_size`` for a + launch (the T1 rungs), the byte offset from the tensor's own base at the + any-parameters rung (T0, where accesses of different tensors never + conflict: one encoding per tensor); +* ``mask`` and ``path`` make the record's ``active``; each enclosing loop's + iteration premise, the in-bounds premise, the loaded values' domains and + the width premises are local constraints (they gate the record's + activity, and stay out of the solver's feasibility base); +* an integer RMW outside loops and a CAS are value-modeled: the access's + observation is the record's ``old_value``; +* an awaited poll (``awaits``) gets its exit predicate as a premise of every + record (where the predicate's own width obligations hold), and a pre-exit + representative; +* a tile access (``AccessEvent.shape``) carries its shape and, when a lane + variable names each of its dimensions, its ``lane_position``; otherwise + the solver reads its lanes as not named (distinct lanes at any + assignment) and it anchors no dependency. + +The in-bounds premise: a record's address stays inside its tensor's +allocation. Without it an unbounded grid would stray into other tensors' +addresses; a race reachable only through an out-of-bounds access is the +compiled sanitizer's finding, not this model's. + +The width premises: the reader's integer terms are exact while their +``width_obligations`` hold. A record is active only where its obligations +hold, each one only where a Select arm reads its term (``GuardDag``), so a +proof at an any-grid rung assumes no overflow, which ``violation`` lets the +caller check (and discharge exactly at the launch). + +Over-approximated records (a dropped mask, an unmodeled branch, a free +observation or loaded value in a mask or path, a loop bound on a free +loaded value) are listed in ``uncertain``: an UNSAT over them still proves, +a SAT touching one is no witness. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Any + +from z3 import And, ArithRef, BoolRef, BoolVal, ExprRef, Implies, Int, IntVal, Not +from z3 import simplify + +from ....core.data import AtomicCas, AtomicRMW, Load, Store +from ....ir.lowering import GuardDag, children, fold +from ....ir.ttir_reader import ( + AccessEvent, + AccessGraph, + Arange, + Bin, + BoolBin, + Cmp, + Const, + IntCast, + IterArgOffset, + Loaded, + LoopInfo, + LoopVar, + NumPrograms, + Observed, + Param, + Pid, + Select, + Term, + WidthObligation, + width_obligations, +) +from ....ir.ttir_reader import Not as TNot +from ... import pids +from ..data import AccessEventRecord +from .fences import FenceRule +from .leaves import ( + LaunchTensor, + RaceKind, + RaceLeaves, + RaceRefusal, + loaded_leaves, + observed_leaves, + signed, + terms_below, +) + +_KNOWN_SEMS = ("relaxed", "acquire", "release", "acq_rel") +# The printer's RMW spellings -> the solver's op names. +_RMW_OPS = {"exch": "xchg"} +# Pre-launch values the solver's atomic read-from and counting read +# (its own initial-value cap). +INIT_VALUES_MAX_ELEMENTS = 1024 + + +@dataclass +class Encoding: + """The solver's input for one encoding (or one T0 tensor group).""" + + records: list[AccessEventRecord] + arange_dict: dict[Any, Any] + # Event ids of over-approximated records (see the module docstring). + uncertain: set[int] + # The snapshot equalities (``extra_assumptions``). + assumptions: tuple[BoolRef, ...] + # Some record read a snapshot: a verdict holds for these contents. + content_qualified: bool + pid_axes: frozenset[int] + has_atomics: bool + fences: FenceRule + + +@dataclass(frozen=True) +class _Tile: + """A tile record's ``symbolic_expr``: the solver reads only its shape.""" + + shape: tuple[int, ...] + + +@dataclass(frozen=True) +class _Loop: + var: ArithRef # the iteration index k + premise: BoolRef # k is an iteration that runs + trips: int | None # the trip count, when a constant + domain: tuple[BoolRef, ...] # its bounds' loaded values' domains + + +class _InitValues: + """The tensor surface the solver's initial-value source reads, over the + pre-launch values (the solver never touches a real tensor here).""" + + def __init__(self, facts: Any, values: tuple[int, ...]) -> None: + self._facts = facts + self._values = values + + def data_ptr(self) -> int: + return self._facts.data_ptr + + def element_size(self) -> int: + return self._facts.elem_size + + def numel(self) -> int: + return self._facts.numel + + def is_contiguous(self) -> bool: + return self._facts.contiguous + + def reshape(self, *_shape: Any) -> _InitValues: + return self + + def tolist(self) -> list[int]: + return list(self._values) + + +# ─────────────────────────── term predicates ─────────────────────────── + + +def _loop_names(graph: AccessGraph, access: AccessEvent) -> tuple[str, ...]: + if access.loops: + return access.loops + return (graph.loop.loop_ssa,) if access.in_loop and graph.loop else () + + +def access_terms(access: AccessEvent) -> tuple[Term, ...]: + return tuple( + t + for t in ( + access.offset, + access.mask, + access.path, + access.atomic_val, + access.atomic_cmp, + ) + if t is not None + ) + + +def graph_terms(graph: AccessGraph) -> list[object]: + terms: list[object] = [] + for a in graph.accesses: + terms += access_terms(a) + for lp in graph.loops: + terms += [lp.lower, lp.upper, lp.step] + for info in graph.iter_args: + terms += [info.offset0, info.delta] + return terms + + +def mentions_loaded(term: object, graph: AccessGraph) -> bool: + return any(isinstance(t, Loaded) for t in terms_below((term,), graph)) + + +def _conjuncts(term: Term) -> list[Term]: + out, stack = [], [term] + while stack: + t = stack.pop() + if isinstance(t, BoolBin) and t.op == "and": + stack += [t.b, t.a] + else: + out.append(t) + return out + + +def _same(a: object, b: object) -> bool: + if a is b: + return True + try: + return a == b + except RecursionError: + return False + + +def _runs_wherever(d: AccessEvent, a: AccessEvent) -> bool: + """``d`` runs at every lane and iteration where ``a`` does: no dropped + mask or unmodeled branch, its loops enclose ``a``, and its mask and path + conjuncts are among ``a``'s. (Where ``d`` is masked off, a value ``a`` + reads from it does not come from ``d``'s sources.)""" + if d.mask_dropped or d.guarded or a.loops[: len(d.loops)] != d.loops: + return False + have = [c for t in (a.mask, a.path) if t is not None for c in _conjuncts(t)] + return all( + (isinstance(c, Const) and bool(c.value)) or any(_same(c, h) for h in have) + for t in (d.mask, d.path) + if t is not None + for c in _conjuncts(t) + ) + + +def _var_ids(terms: Sequence[Any]) -> set[int]: + """The ids of every node of ``terms`` (Z3 interns them: a variable has + one id wherever it occurs).""" + seen: set[int] = set() + stack = [t for t in terms if isinstance(t, ExprRef)] + while stack: + t = stack.pop() + i = t.get_id() + if i not in seen: + seen.add(i) + stack.extend(t.children()) + return seen + + +# Symbolic leaves: at T0 every leaf but constants; at T1 the params are +# concrete. +_T0_SYMBOLIC: tuple[type, ...] = (Pid, Param, Arange, LoopVar, IterArgOffset) +_T0_SYMBOLIC += (Observed, NumPrograms, Loaded) +_T1_SYMBOLIC = tuple(c for c in _T0_SYMBOLIC if c is not Param) + + +def linear(term: object, graph: AccessGraph, symbolic: tuple[type, ...]) -> bool: + """No product of two symbolic terms and no symbolic divisor (``symbolic`` + names the leaves a rung leaves symbolic), the shape Z3 decides.""" + + def apply(t: object, kids: Sequence[tuple[bool, bool]]) -> tuple[bool, bool]: + sym = isinstance(t, symbolic) or any(s for s, _ in kids) + lin = all(ok for _, ok in kids) + if isinstance(t, Bin) and t.op == "*": + lin = lin and not (kids[0][0] and kids[1][0]) + elif isinstance(t, Bin) and t.op in ("//", "%", "u//", "u%"): + lin = lin and not kids[1][0] + elif isinstance(t, (LoopVar, IterArgOffset)): + # lower + k * step, offset0 + k * delta + lin = kids[0][1] and not kids[1][0] + return sym, lin + + return fold(term, graph, apply, {})[1] + + +def _cas_operand_ok(term: object, graph: AccessGraph) -> bool: + """A CAS operand with no arithmetic and no width change: a literal, an + observation, a comparison of atoms, its i1 zero extension, or a Select + of those (an observed i32 + 1 must not become unbounded addition where + it decides the CAS).""" + atoms = (Const, Pid, Arange, Observed) + + def apply(t: object, kids: Sequence[tuple[bool, bool]]) -> tuple[bool, bool]: + # (a condition, a value) + if isinstance(t, Const): + return True, True + if isinstance(t, Observed): + return False, True + if isinstance(t, Cmp): + ok = isinstance(t.a, atoms) and isinstance(t.b, atoms) + return ok, ok + if isinstance(t, TNot): + return kids[0][0], kids[0][0] + if isinstance(t, BoolBin): + ok = kids[0][0] and kids[1][0] + return ok, ok + if isinstance(t, IntCast) and t.kind == "extui" and t.src_bits == 1: + return False, kids[0][0] + if isinstance(t, Select): + return False, kids[0][0] and kids[1][1] and kids[2][1] + return False, False + + return fold(term, graph, apply, {})[1] + + +def _widens(term: object, graph: AccessGraph) -> bool: + """``term`` changes an integer's width (an IntCast other than an i1's + zero extension).""" + return any( + isinstance(t, IntCast) and not (t.kind == "extui" and t.src_bits == 1) + for t in terms_below((term,), graph) + ) + + +def content_free_view(access: AccessEvent, graph: AccessGraph) -> AccessEvent: + """The access with every loaded value free: a mask or path conjunct + that reads one is dropped (the solver could pick its value either way, + and it may be nonlinear), and the access is flagged over-approximated. + An address on a loaded value is left alone (it refuses).""" + + def strip(t: Term | None) -> tuple[Term | None, bool]: + if t is None: + return None, False + conj = _conjuncts(t) + kept = [c for c in conj if not mentions_loaded(c, graph)] + if len(kept) == len(conj): + return t, False + if not kept: + return None, True + out = kept[0] + for c in kept[1:]: + out = BoolBin("and", out, c) + return out, True + + mask, m_dropped = strip(access.mask) + path, p_dropped = strip(access.path) + if not (m_dropped or p_dropped): + return access + return replace( + access, + mask=mask, + path=path, + mask_dropped=access.mask_dropped or m_dropped, + guarded=access.guarded or p_dropped, + ) + + +def t0_linearity_gate(graph: AccessGraph, awaits: Mapping[int, Term]) -> bool: + """The any-parameters rung is tried only when every address, mask, path + and exit predicate is linear with the params symbolic.""" + for i, access in enumerate(graph.accesses): + view = content_free_view(access, graph) + terms = [view.offset, view.mask, view.path, awaits.get(i)] + if not all(linear(t, graph, _T0_SYMBOLIC) for t in terms if t is not None): + return False + return True + + +def written_sources( + graph: AccessGraph, tensors: Mapping[str, LaunchTensor] +) -> dict[str, str]: + """The tensors loaded values are read from that the kernel may write: + their snapshot does not stand for what the load reads (the static model + cannot order the instances), so their values stay free.""" + sources = { + graph.accesses[leaf.access_index].base_param + for leaf in loaded_leaves(graph_terms(graph), graph) + } + written = {a.base_param for a in graph.accesses if a.kind != "load"} + + def interval(name: str) -> tuple[int, int] | None: + t = tensors.get(name) + return None if t is None else t.facts.allocation_interval() + + out: dict[str, str] = {} + for src in sorted(sources): + si = interval(src) + for w in sorted(written): + wi = interval(w) + if src == w or (si and wi and max(si[0], wi[0]) < min(si[1], wi[1])): + out[src] = f"the kernel writes {w!r}, which overlaps the source" + break + return out + + +def symbolic_grid( + enc: Encoding, launch_grid: tuple[int, int, int] | None, t0: bool = False +) -> tuple[Any, Any, Any]: + """The grid of an any-grid rung: symbolic along the axes the kernel + reads; along an unread axis the launch's extent (blocks that differ + only there behave alike, but a launch along it is still checked), or + symbolic for a kernel with atomics without a launch extent or at T0 + (their blocks need not behave alike).""" + + def dim(i: int) -> Any: + if i in enc.pid_axes or (enc.has_atomics and (t0 or launch_grid is None)): + return pids.grid_var(i) + return int(launch_grid[i]) if launch_grid is not None else 1 + + return (dim(0), dim(1), dim(2)) + + +def grid_pins(grid: tuple[Any, Any, Any], launch: tuple[int, int, int]) -> tuple: + """Equalities pinning the symbolic dims of ``grid`` to the launch.""" + return tuple( + d == IntVal(launch[i]) for i, d in enumerate(grid) if not isinstance(d, int) + ) + + +def _fits(value: ArithRef, bits: int, is_signed: bool) -> BoolRef: + if is_signed: + half = 1 << (bits - 1) + return And(value >= -half, value < half) + return And(value >= 0, value < (1 << bits)) + + +def _is_increment(ob: WidthObligation, loop: LoopInfo) -> bool: + """The loop's increment obligation (``upper - 1 + step``), which holds + only where the loop runs.""" + t = ob.term + return ( + isinstance(t, Bin) + and t.op == "+" + and t.b is loop.step + and isinstance(t.a, Bin) + and t.a.op == "-" + and t.a.a is loop.upper + ) + + +def _guard_operands(t: object, graph: AccessGraph) -> tuple: + if isinstance(t, Loaded): + return tuple(k for k in (t.offset, t.mask, t.other) if k is not None) + return children(t, graph) + + +# ─────────────────────────── the encoder ─────────────────────────── + + +class Encoder: + """One encoding of ``graph``: ``params`` None is the any-parameters + rung (T0); otherwise the launch's params and tensors (T1), and + ``blind`` reads no snapshot (the content-free attempt).""" + + def __init__( + self, + graph: AccessGraph, + params: Mapping[str, int] | None, + tensors: Mapping[str, LaunchTensor], + awaits: Mapping[int, Term], + *, + blind: bool = False, + ) -> None: + self.graph = graph + self.t0 = params is None + self.tensors = tensors + self.awaits = awaits + unusable = {} if self.t0 else written_sources(graph, tensors) + self.leaves = RaceLeaves(graph, params, tensors, blind=blind, unusable=unusable) + self.modeled = { + i + for i, a in enumerate(graph.accesses) + if a.kind in ("atomic_rmw", "atomic_cas") + and not a.elem_float + and not a.in_loop + } + # Dependency order is causal, so it composes: an access that uses + # an atomic's result at a lane comes after the loads that atomic + # used at that lane, where the atomic runs (the reader lists direct + # sources only). + self.deps: list[tuple[int, ...]] = [] + for a in graph.accesses: + found = set(a.deps) + for d in a.deps: + if _runs_wherever(graph.accesses[d], a): + found.update(self.deps[d]) + self.deps.append(tuple(sorted(found))) + # access index -> its lane position (see lane_position) + self._positions: dict[int, tuple[Any, ...] | None] = {} + self.loops: dict[str, _Loop] = {} + self.free_bound_loops: set[str] = set() + for index, lp in enumerate(graph.loops): + self._bind_loop(lp, index) + self.fences = FenceRule( + graph, {name: (b.var, b.trips) for name, b in self.loops.items()} + ) + # (width, where one obligation of an encoded record or an await's + # exit predicate fails) + self.violations: list[tuple[int, BoolRef]] = [] + self.await_premises, self.await_obs = self._await_premises() + + # ── loops ── + + def _bind_loop(self, lp: LoopInfo, index: int) -> None: + leaves = self.leaves + try: + lower = leaves.value(lp.lower) + upper = leaves.value(lp.upper) + step = leaves.value(lp.step) + except RaceRefusal as r: + raise r.at(lp.line_no, lp.loc) + bound_leaves = loaded_leaves((lp.lower, lp.upper), self.graph) + if any( + self._free(leaf) or self._pad_reaches(leaf, ()) for leaf in bound_leaves + ): + # The iteration range is over-approximated. + self.free_bound_loops.add(lp.loop_ssa) + domain = leaves.domain_premises(bound_leaves) + lower_c, upper_c, step_c = (_as_long(v) for v in (lower, upper, step)) + + def refuse(why: str) -> RaceRefusal: + return RaceRefusal(RaceKind.LOOP_BOUND, why, line_no=lp.line_no, loc=lp.loc) + + # A symbolic k * step is nonlinear; scf.for needs a positive step. + if step_c is None: + raise refuse("the loop step is not a compile-time constant") + if step_c <= 0: + raise refuse(f"the loop step is {step_c}; only positive steps are modeled") + k = Int("ttir_loop_k" if index == 0 else f"ttir_loop_k{index}") + trips: int | None = None + if lower_c is not None and upper_c is not None: + trips = max(0, -(-(upper_c - lower_c) // step_c)) + premise = And(k >= 0, k < trips) + else: + # A bound that stays symbolic: at T1 only one linear in the + # pids, lanes and iterations (a grid-stride or triangular loop). + if not self.t0 and not all( + linear(t, self.graph, _T1_SYMBOLIC) for t in (lp.lower, lp.upper) + ): + raise refuse( + "the loop bounds are neither constant at the launch nor " + "linear in the program ids and iterations" + ) + premise = And(k >= 0, lower + k * IntVal(step_c) < upper) + self.loops[lp.loop_ssa] = _Loop(k, premise, trips, domain) + # A single iteration reads its bounds' lower value directly. + leaves.iterations[lp.loop_ssa] = IntVal(0) if trips == 1 else k + + def iteration_premises(self, access: AccessEvent) -> tuple[BoolRef, ...]: + out: list[BoolRef] = [] + for name in _loop_names(self.graph, access): + b = self.loops[name] + out += [b.premise, *b.domain] + return tuple(out) + + def zero_trip(self, access: AccessEvent) -> bool: + return any( + self.loops[name].trips == 0 for name in _loop_names(self.graph, access) + ) + + # ── loaded values ── + + def _free(self, leaf: Loaded) -> str | None: + return self.leaves.snapshot_problem(self.graph.accesses[leaf.access_index]) + + def _pad_reaches(self, leaf: Loaded, guards: Sequence[Term]) -> bool: + """The load's masked-off lanes have no value (a mask without + ``other``) and the consumer's own mask and path do not repeat the + load's mask (which would keep those lanes inactive).""" + if leaf.mask is None or leaf.other is not None or self._free(leaf): + return False + have = [c for g in guards for c in _conjuncts(g)] + return not all(any(_same(c, h) for h in have) for c in _conjuncts(leaf.mask)) + + def widened(self, access: AccessEvent, exit_pred: Term | None) -> bool: + """The access's activity or address is over-approximated.""" + if access.mask_dropped or access.guarded: + return True + for t in (access.mask, access.path): + if t is not None and observed_leaves((t,), self.graph) - self.modeled: + return True # a free observation + guards = [t for t in (access.mask, access.path) if t is not None] + for t in (access.offset, access.mask, access.path, exit_pred): + if t is None: + continue + for leaf in loaded_leaves((t,), self.graph): + if t is not access.offset and self._free(leaf): + return True + if self._pad_reaches(leaf, guards): + return True + return any( + name in self.free_bound_loops for name in _loop_names(self.graph, access) + ) + + # ── awaits ── + + def _await_premises(self) -> tuple[tuple[BoolRef, ...], tuple[ArithRef, ...]]: + """For each await, ``reach(poll) -> exit``: every terminating run + satisfies it, so it holds in every record's activity (asserting it + on the poll alone would let a model leave the poll inactive and drop + its synchronization). The exit predicate is exact only where its + width obligations hold: the premise asks them too, and where the + poll is reached they join ``violations``. A guarded await gives + none.""" + graph, leaves = self.graph, self.leaves + premises, obs = [], [] + for index, exit_pred in self.awaits.items(): + access = graph.accesses[index] + if access.guarded: + continue + try: + exit_b = leaves.cond(exit_pred) + guard = [ + leaves.cond(t) for t in (access.mask, access.path) if t is not None + ] + # The predicate's own obligations, as a mask's. + probe = replace( + access, + loops=(), + in_loop=False, + path=None, + mask=exit_pred, + offset=Const(0), + ) + dag = GuardDag([exit_pred], graph, leaves.cond, _guard_operands) + reach = And( + *guard, + *self.iteration_premises(access), + *leaves.domain_premises(loaded_leaves((exit_pred,), graph)), + ) + held: list[BoolRef] = [] + for ob in width_obligations(graph, probe): + premise = self._held(dag, ob) + held.append(premise) + self.violations.append((ob.bits, And(reach, Not(premise)))) + except RaceRefusal as r: + raise r.at(access.line_no, access.loc) + if guard or held: + premises.append(Implies(And(*guard, *held), exit_b)) + else: + premises.append(exit_b) + obs.append(leaves.observation(index)) + return tuple(premises), tuple(obs) + + def _await_dependencies(self, seq: int, access: AccessEvent) -> tuple[int, ...]: + if access.guarded or access.mask_dropped: + return () + return tuple( + i + for i in self.awaits + if i < seq + and not self.graph.accesses[i].guarded + and not self.graph.accesses[i].mask_dropped + ) + + # ── width premises ── + + def _held(self, dag: GuardDag, ob: WidthObligation) -> BoolRef: + """Where obligation ``ob`` holds or no Select arm reads its term.""" + fits = _fits(self.leaves.value(ob.term), ob.bits, ob.signed) + # A signed quotient that does not fit is undefined in either arm. + undefined = isinstance(ob.term, Bin) and ob.term.op == "//" + return fits if undefined else dag.premise(ob.term, fits) + + def _width( + self, access: AccessEvent + ) -> tuple[BoolRef | None, list[tuple[int, BoolRef]]]: + """The access's width premise (each obligation where a Select arm + reads its term) and where it fails while it matters: a loop's + obligations where the loop is reached (its increment where it + runs), the path's at every iteration, the mask's where the path + holds, the offset's where the access runs.""" + graph, leaves = self.graph, self.leaves + names = _loop_names(graph, access) + loops: list[LoopInfo] = [] + for n in names: + lp = graph.loop_info(n) or graph.loop + assert lp is not None + loops.append(lp) + roots: list[object] = [] + for lp in loops: + roots += [lp.lower, lp.upper, lp.step] + roots += [access.path, access.mask, access.offset] + dag = GuardDag(roots, graph, leaves.cond, _guard_operands) + path = leaves.cond(access.path) if access.path is not None else BoolVal(True) + mask = leaves.cond(access.mask) if access.mask is not None else BoolVal(True) + iters = [self.loops[n].premise for n in names] + held: list[BoolRef] = [] + failing: list[tuple[int, BoolRef]] = [] + + def need(ob: WidthObligation, where: BoolRef) -> None: + premise = self._held(dag, ob) + held.append(premise) + failing.append((ob.bits, And(where, Not(premise)))) + + for j, (name, lp) in enumerate(zip(names, loops)): + probe = replace( + access, + loops=(name,), + in_loop=True, + path=None, + mask=None, + offset=Const(0), + ) + reach = And(path, *iters[:j]) + for ob in width_obligations(graph, probe): + if ob.role == "loop": + runs = _is_increment(ob, lp) + need(ob, And(reach, iters[j]) if runs else reach) + where = { + "path": And(*iters) if iters else BoolVal(True), + "mask": And(*iters, path), + "offset": And(*iters, path, mask), + } + flat = replace(access, loops=(), in_loop=False) + for ob in width_obligations(graph, flat): + need(ob, where[ob.role]) + return (And(*held) if held else None), failing + + # ── lane positions ── + + def lane_position(self, seq: int) -> tuple[Any, ...] | None: + """The block position of access ``seq``'s lanes: per dimension of + its tile, the lane variable of that ``(dim, extent)`` (0 along an + extent-1 dimension). None when a dimension is named by no lane + variable its address and activity read (a broadcast pointer, a block + on one address), or when they do not lower.""" + if seq in self._positions: + return self._positions[seq] + access = self.graph.accesses[seq] + leaves = self.leaves + position: tuple[Any, ...] | None = None + try: + terms = [leaves.value(access.offset)] + terms += [ + leaves.cond(t) for t in (access.mask, access.path) if t is not None + ] + except RaceRefusal: + terms = [] + if terms: + ids = _var_ids(terms) + rank = len(access.shape) + coords: list[Any] = [] + for d, n in enumerate(access.shape): + if n == 1: + coords.append(IntVal(0)) + continue + # a rank-1 tile's aranges are unplaced (dim -1) + keys = [(-1, n), (0, n)] if rank == 1 else [(d, n)] + found = [ + leaves.lanes[k] + for k in keys + if k in leaves.lanes and leaves.lanes[k].get_id() in ids + ] + if not found: + coords = [] + break + coords.append(found[0]) + position = tuple(coords) if coords else None + self._positions[seq] = position + return position + + def _tile(self, seq: int) -> bool: + return any(n > 1 for n in self.graph.accesses[seq].shape) + + def _anchors(self, seq: int) -> tuple[int, ...]: + """The dependency anchors of access ``seq``: a tile anchor only + with a lane position (the solver compares positions through it; + a block without one has no position to be ordered at).""" + return tuple( + d + for d in self.deps[seq] + if not self._tile(d) or self.lane_position(d) is not None + ) + + # ── one record ── + + def record(self, seq: int, access: AccessEvent) -> AccessEventRecord: + try: + return self._record(seq, access) + except RaceRefusal as r: + raise r.at(access.line_no, access.loc) + + def _record(self, seq: int, access: AccessEvent) -> AccessEventRecord: + graph, leaves = self.graph, self.leaves + where = f"line {access.line_no}" + meta = None if self.t0 else self.tensors.get(access.base_param) + if not self.t0: + if meta is None: + raise RaceRefusal( + RaceKind.MISSING_BINDING, + f"pointer argument {access.base_param!r} has no tensor binding", + ) + if meta.facts.allocation_interval() is None: + layout = "" if meta.facts.contiguous else "non-contiguous " + raise RaceRefusal( + RaceKind.MISSING_BINDING, + f"{layout}tensor {access.base_param!r}: the in-bounds " + "premise needs a verified allocation extent", + ) + if access.elem_bits <= 0: + raise RaceRefusal( + RaceKind.OTHER, f"unknown element width for {access.base_param!r}" + ) + elem = (access.elem_bits + 7) // 8 # an i1 is a byte + old_value: Any = None + rmw_op: str | None = None + rmw_operand: Any = None + cas_cmp: Any = None + cas_new: Any = None + exit_pred = self.awaits.get(seq) + if access.kind == "atomic_rmw" and seq in self.modeled: + assert access.atomic is not None + old_value = leaves.observation(seq) + op = access.atomic.rmw_op + rmw_op = _RMW_OPS.get(op, op) if op else None + if access.atomic_val is not None: + try: + rmw_operand = leaves.value(access.atomic_val) + except RaceRefusal: + rmw_operand = None # an unmodeled operand: the write stays open + elif access.kind == "atomic_cas": + cas_cmp, cas_new = self._cas(seq, access, exit_pred is not None) + old_value = leaves.observation(seq) + if observed_leaves((access.offset,), graph) - self.modeled: + raise RaceRefusal( + RaceKind.INDIRECT_ADDRESS, + f"{where}: the address depends on an atomic observation that is " + "not value-modeled (a float or loop-carried atomic)", + ) + # (a free loop bound only widens the loop's accesses) + for leaf in loaded_leaves((access.offset,), graph, loops=False): + why = self._free(leaf) + if why is not None: + raise RaceRefusal( + RaceKind.SNAPSHOT_BOUND + if why.startswith("too large") + else RaceKind.INDIRECT_ADDRESS, + f"{where}: the address depends on a loaded value with no " + f"usable snapshot ({why})", + ) + offset = leaves.value(access.offset) + bounds: tuple[BoolRef, ...] = () + if meta is None: + addr = offset * IntVal(elem) + else: + facts = meta.facts + if facts.elem_size != elem: + raise RaceRefusal( + RaceKind.OTHER, + f"element width mismatch for {access.base_param!r}: the TTIR " + f"reads {elem} bytes, the launch tensor has {facts.elem_size}", + ) + addr = IntVal(facts.data_ptr) + offset * IntVal(elem) + start, end = facts.allocation_interval() # type: ignore[misc] + bounds = (addr >= IntVal(start), addr + IntVal(elem) <= IntVal(end)) + active: Any = True + conds = [leaves.cond(t) for t in (access.mask, access.path) if t is not None] + if conds: + active = conds[0] if len(conds) == 1 else And(*conds) + + atomic = access.atomic + scope: str | None = None + reads: Any = None + writes: Any = None + if access.kind in ("atomic_rmw", "atomic_cas"): + assert atomic is not None + sem: Any = atomic.sem if atomic.sem in _KNOWN_SEMS else "relaxed" + scope = atomic.scope + is_atomic = True + access_mode: Any = "read" + if access.kind == "atomic_rmw": + op_type: type = AtomicRMW + atomic_kind: Any = "rmw" + reads, writes = True, True + else: + # The solver's CAS lowering writes only on success. + op_type, atomic_kind = AtomicCas, "cas" + reads = True + else: + sem, is_atomic, atomic_kind = "plain", False, "none" + op_type = Store if access.kind == "store" else Load + access_mode = "write" if access.kind == "store" else "read" + + value_terms = [access.mask, access.path, exit_pred] + if access.kind == "atomic_cas": + value_terms += [access.atomic_cmp, access.atomic_val] + roots = [access.offset, *(t for t in value_terms if t is not None)] + loaded = loaded_leaves(roots, graph) + iteration = self.iteration_premises(access) + width, failing = self._width(access) + domain = leaves.domain_premises(loaded) + self.violations += [(bits, And(*domain, f)) for bits, f in failing] + # Every observation a record reads is renamed per copy: the solver + # renames only the copy-local variables of the records it is given. + # Of the loops, only the record's own: the solver reads them as the + # record's iterators, and one of a loop the access is not in would + # make one execution two. The loops of an awaited poll too, whose + # premise rides every record. + names = dict.fromkeys(_loop_names(graph, access)) + for index in self.awaits: + names.update(dict.fromkeys(_loop_names(graph, graph.accesses[index]))) + copy_local = tuple(self.loops[n].var for n in names) + copy_local += tuple( + leaves.observation(i) for i in sorted(observed_leaves(roots, graph)) + ) + if old_value is not None: + copy_local += (old_value,) + copy_local += self.await_obs + leaves.pads + tensor = None + if old_value is not None and meta is not None: + tensor = _init_values(meta) + source = ( + (access.loc.file, access.loc.line, graph.kernel_name) + if access.loc is not None + else None + ) + tile = _Tile(tuple(access.shape)) if self._tile(seq) else None + return AccessEventRecord( + op_type=op_type, + access_mode=access_mode, + tensor=tensor, # type: ignore[arg-type] + tensor_name=access.base_param, + symbolic_expr=tile, + addr_expr=addr, + # The await premises hold in every execution; the iteration + # ranges gate only this record (several loops may run for + # disjoint instances). + premises=self.await_premises, + local_constraints=bounds + + iteration + + domain + + ((width,) if width is not None else ()), + source_location=source, + program_seq=seq, + dep_loads=self._anchors(seq), + lane_position=self.lane_position(seq) if tile else None, + causal_constraints=iteration, + await_observations=( + tuple(sorted(observed_leaves((exit_pred,), graph))) + if exit_pred is not None + else () + ), + await_dependencies=self._await_dependencies(seq, access), + debug_name=f"{graph.kernel_name}:ttir{access.line_no}:{access.kind}", + active=active, + reads=reads, + writes=writes, + is_atomic=is_atomic, + atomic_kind=atomic_kind, + sem=sem, + scope=scope, + old_value=old_value, + rmw_op=rmw_op, + rmw_operand=rmw_operand, + cas_cmp_value=cas_cmp, + cas_new_value=cas_new, + event_id=seq, + elem_size=elem, + copy_local_vars=copy_local, + ) + + def _cas(self, seq: int, access: AccessEvent, awaited: bool) -> tuple[Any, Any]: + """The CAS's compare and new values. Neither may be dropped (they + decide success and the read-from), so an operand outside the exact + integer fragment refuses.""" + graph = self.graph + where = f"line {access.line_no}" + + def refuse(why: str) -> RaceRefusal: + return RaceRefusal(RaceKind.CAS_VALUE, f"{where}: {why}") + + if access.in_loop: + raise RaceRefusal( + RaceKind.CONTROL_FLOW, + f"{where}: CAS inside scf.for (one observation cannot stand " + "for one per iteration)", + ) + if access.elem_float: + raise refuse("a float-typed CAS is outside the integer value model") + if access.atomic_cmp is None or access.atomic_val is None: + raise refuse("the CAS compare or new value is not modelable") + if not awaited: + for operand in (access.atomic_cmp, access.atomic_val): + if loaded_leaves((operand,), graph): + raise refuse("a CAS operand depends on a plain loaded value") + observed = observed_leaves((operand,), graph) + if observed - self.modeled: + raise refuse( + "a CAS operand depends on an atomic observation without " + "an integer value model" + ) + if _widens(operand, graph): + raise refuse("a CAS operand changes an integer's width") + if not _cas_operand_ok(operand, graph): + raise refuse( + "CAS operand arithmetic is outside the modeled integer " + "fragment" + ) + if any( + graph.accesses[i].elem_bits != access.elem_bits for i in observed + ): + raise refuse("a CAS operand changes an observation's width") + return self.leaves.value(access.atomic_cmp), self.leaves.value( + access.atomic_val + ) + + # ── encodings ── + + def domain(self) -> tuple[BoolRef, ...]: + """The ranges every query assumes, for the obligation check.""" + out: list[BoolRef] = [] + for axis in pids.PID_AXES: + g = pids.grid_var(axis) + out += [pids.pid_var(axis) >= 0, pids.pid_var(axis) < g, g >= 1] + for (start, end, *_), (var, _) in self.leaves.arange_dict.items(): + out += [var >= start, var < end] + return tuple(out) + tuple(self.leaves.assumptions) + self.await_premises + + def encode( + self, accesses: Sequence[tuple[int, AccessEvent]] | None = None + ) -> Encoding: + """The records of ``accesses`` (by default every access of the + graph; an access of a zero-trip loop runs never and is left out).""" + graph = self.graph + items = list(enumerate(graph.accesses)) if accesses is None else list(accesses) + records: list[AccessEventRecord] = [] + uncertain: set[int] = set() + # Pre-exit representatives get ids above every access's. + next_rep = len(graph.accesses) + for seq, access in items: + if self.zero_trip(access): + continue + rec = self.record(seq, access) + records.append(rec) + widened = self.widened(access, self.awaits.get(seq)) + if widened: + uncertain.add(seq) + if seq in self.awaits: + records.append(pre_exit_representative(rec, access, next_rep)) + if widened: + uncertain.add(next_rep) + next_rep += 1 + return Encoding( + records=records, + arange_dict=dict(self.leaves.arange_dict), + uncertain=uncertain, + assumptions=tuple(self.leaves.assumptions), + content_qualified=self.leaves.used_snapshot, + pid_axes=frozenset(graph.pid_axes), + has_atomics=any(a.atomic is not None for a in graph.accesses), + fences=self.fences, + ) + + +def _as_long(v: ArithRef) -> int | None: + try: + return simplify(v).as_long() + except Exception: + return None + + +def _init_values(meta: LaunchTensor) -> _InitValues | None: + """The pre-launch values an atomic's read-from may start from: a small + signed integer tensor's (an unsigned one wraps by definition).""" + facts = meta.facts + if ( + meta.values is None + or not facts.contiguous + or facts.numel > INIT_VALUES_MAX_ELEMENTS + or facts.dtype.startswith("torch.uint") + ): + return None + bits = facts.elem_size * 8 + return _InitValues(facts, tuple(signed(int(v), bits) for v in meta.values)) + + +def pre_exit_representative( + poll: AccessEventRecord, access: AccessEvent, event_id: int +) -> AccessEventRecord: + """One record for every failed iteration of an awaited poll: the poll's + footprint, activity, sem, scope and program_seq (unordered with the + poll, ordered like it against everything else) under a fresh id, with + no value model, so no read-from or synchronization reaches it.""" + if access.kind == "atomic_cas": + # A failed CAS reads and writes nothing: a read-only RMW record. + extra: dict[str, Any] = dict( + op_type=AtomicRMW, atomic_kind="rmw", reads=True, writes=False + ) + elif access.kind == "atomic_rmw": + # An identity RMW writes back what it read on every iteration. + extra = dict(reads=True, writes=True) + else: + extra = dict(reads=None, writes=None) + return replace( + poll, + event_id=event_id, + debug_name=f"{poll.debug_name}:pre-exit", + pre_exit=True, + old_value=None, + rmw_op=None, + rmw_operand=None, + cas_cmp_value=None, + cas_new_value=None, + **extra, + ) + + +def graph_mentions_loaded(graph: AccessGraph) -> bool: + return bool(loaded_leaves(graph_terms(graph), graph)) + + +__all__ = [ + "Encoder", + "Encoding", + "content_free_view", + "graph_mentions_loaded", + "grid_pins", + "linear", + "pre_exit_representative", + "symbolic_grid", + "t0_linearity_gate", + "written_sources", +] diff --git a/tilelens/clients/race_detector/compiled/leaves.py b/tilelens/clients/race_detector/compiled/leaves.py new file mode 100644 index 00000000..5965b4e3 --- /dev/null +++ b/tilelens/clients/race_detector/compiled/leaves.py @@ -0,0 +1,635 @@ +"""The race detector's meaning of the TTIR reader's leaves. + +``RaceLeaves`` is the race detector's ``TermLeaves``: it lowers the +reader's terms in Z3's main context (``ctx`` None), where the two-copy +solver renames each program copy's variables with ``z3.substitute``: + +* a program id is ``tilelens.clients.pids.pid_var(axis)`` and the grid size + ``grid_var(axis)``, the names the solver renames and bounds; +* an arange is ``start + lane``, one lane variable per ``(dim, extent)`` + (the reader's contract, see ``Arange``), registered in the solver's + arange dict so each copy gets its own lane in ``[0, extent)``; +* a loop's iteration index is the variable its binding gives + (``iterations``), copy-local in the records; +* an atomic observation is one interned variable per access; +* a scalar argument is a shared free Int when ``params`` is None (no + launch: the any-params rung) and the launch's value, read signed at the + argument's width, otherwise; +* a loaded value is a ``Select`` over the launch's pre-launch contents of + its tensor (a snapshot), read signed at the load's own element width (an + unsigned table with a value the signed reading makes negative has no + usable snapshot); without a usable snapshot it is free (a per-load array, + copy-local), and the access that reads it is over-approximated; +* an unmodeled value refuses; +* ``//`` and ``%`` are Z3's Euclidean ones where the operands' signs make + them the IR's truncating ones (``BoundedDivRem``), else the lowering's + default. + +What cannot be modeled raises :class:`RaceRefusal` with a ``RaceKind``. +""" + +from __future__ import annotations + +import weakref +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from enum import Enum +from typing import Any + +from z3 import ( + And, + ArithRef, + Array, + BoolRef, + If, + Implies, + Int, + IntSort, + IntVal, + Select, + simplify, +) + +from ....ir import failures +from ....ir.launch import TensorFacts +from ....ir.lowering import Lowerer, children, fold, truncating_divrem +from ....ir.ttir_reader import ( + AccessEvent, + AccessGraph, + Arange, + Bin, + BoolBin, + Cmp, + Const, + DataDep, + IntCast, + Loaded, + LoopVar, + Not, + NumPrograms, + Observed, + Param, + Pid, +) +from ....ir.ttir_reader import Select as TSelect +from ... import pids + + +class RaceKind(str, Enum): + """Why the compiled race detector left a kernel undecided (its own + kinds, apart from the reader's ``TTIRKind``).""" + + # An scf.while that is not the await shape. + SPIN_SHAPE = "spin-shape" + # A CAS whose operands are outside the integer value model. + CAS_VALUE = "cas-value" + # A CAS inside a loop (one observation cannot stand for one per + # iteration). + CONTROL_FLOW = "control-flow" + # An address built on an observation or a loaded value that is not + # value-modeled. + INDIRECT_ADDRESS = "indirect-address" + # An address built on a loaded value whose table is too large to read. + SNAPSHOT_BOUND = "snapshot-bound" + # A value the reader marked unmodelable reached a modeled term. + UNMODELED_VALUE = "unmodeled-value" + # A kernel argument, tensor fact or the grid the analysis needs is not + # bound (or unusable) for the launch. + MISSING_BINDING = "missing-binding" + # A loop whose bounds or step the model cannot iterate. + LOOP_BOUND = "loop-bound" + # A width obligation can fail at the launch, where the proof needs it. + INTEGER_OVERFLOW = "integer-overflow" + # Z3 answered unknown: a timeout, or another reason. + TIMEOUT = "timeout" + UNKNOWN = "unknown" + # Every race found is between over-approximated records (a dropped + # mask, a free loaded value): a possible race, not a witness. + POTENTIAL_RACE = "potential-race" + # The premises admit no execution, so a proof would be vacuous. + VACUOUS = "vacuous" + OTHER = "other" + # The client's (compiled/client.py, which classifies a failed compile + # with tilelens.ir.failures, whose two kinds these are): the launch + # compiled nothing to read; a config failed to compile for the IR + # target, or none compiled; a compiled kernel holds no TTIR; the host + # compile could not run; the analysis itself failed (a bug, contained). + NO_COMPILED_KERNEL = "no-compiled-kernel" + COMPILE_FAILED = failures.COMPILE_FAILED + NO_TTIR = "no-ttir" + HOST_COMPILE_UNAVAILABLE = failures.HOST_COMPILE_UNAVAILABLE + INTERNAL_ERROR = "internal-error" + + def __str__(self) -> str: + return self.value + + def __format__(self, spec: str) -> str: + return format(self.value, spec) + + +class RaceRefusal(Exception): + """The compiled race detector cannot decide a kernel: ``kind`` (a + ``RaceKind``), ``message``, and the TTIR ``line_no`` / user ``loc`` of + the access it refused at (None when not tied to one).""" + + def __init__( + self, + kind: RaceKind | str, + message: str, + *, + line_no: int | None = None, + loc: Any = None, + ) -> None: + super().__init__(message) + self.kind = RaceKind(kind) + self.message = message + self.line_no = line_no + self.loc = loc + + def at(self, line_no: int | None, loc: Any) -> RaceRefusal: + """Tie the refusal to an access when it is tied to none yet.""" + if self.line_no is None and self.loc is None: + self.line_no, self.loc = line_no, loc + return self + + def __reduce__(self): + return ( + type(self), + (self.kind, self.message), + {"line_no": self.line_no, "loc": self.loc}, + ) + + +@dataclass(frozen=True) +class LaunchTensor: + """What the analysis knows of one tensor argument at the launch: its + facts and, for a small integer tensor, its pre-launch contents + (``values``, flat in the view's element order, as the dtype's values; + None with ``why`` saying why not).""" + + facts: TensorFacts + values: tuple[int, ...] | None = None + why: str = "" + + +def signed(value: int, bits: int) -> int: + """The IR's signed reading of a ``bits``-wide value (an i1 stays 0/1).""" + if bits <= 1: + return value + half = 1 << (bits - 1) + return (value + half) % (1 << bits) - half + + +def param_value(graph: AccessGraph, params: Mapping[str, int], name: str) -> int: + """A scalar argument's launch value, read signed at its width.""" + arg = graph.arg(name) + return signed(int(params[name]), arg.int_bits if arg else 0) + + +Bounds = tuple[int | None, int | None] + + +class BoundedDivRem: + """The race detector's ``//`` and ``%`` (the Lowerer's ``divrem``): + Z3's Euclidean ``/`` and ``%`` where the operands' signs make them the + IR's truncating ones, which saves the lowering's sign branches (a + grouped program-id swizzle, ``pid // width`` and ``pid % group``, is + several times faster to solve). The signs come from inclusive bounds + over the domains every query assumes (``bounds``); otherwise the + lowering's default: + + * a dividend ``>= 0``: ``a / b`` and ``a % b`` (Euclidean and truncating + agree for either sign of ``b``; ``%`` is guarded at ``b == 0``, where + the truncating remainder is ``a``, unless ``b``'s sign is known); + * a dividend ``<= 0`` with a divisor of known sign: the same on ``-a``, + the sign re-applied. + """ + + def __init__(self, graph: AccessGraph, params: Mapping[str, int] | None) -> None: + self.graph = graph + self.params = params + self._memo: dict[int, tuple[object, Any]] = {} + + def bounds(self, term: object) -> Bounds: + """Inclusive bounds on ``term``'s lowered value (None: unknown), + from the domains every query assumes and nothing else: constants, + the launch's parameters (none at T0), program ids ``>= 0``, the grid + ``>= 1``, lanes in ``[0, extent)``, and the arithmetic over them. + Never a mask, a path, a launch pin or a loop's range (a loop's + iteration range holds only in its own accesses' records).""" + return fold(term, self.graph, self._bounds, self._memo) + + def _bounds(self, t: object, kids: Sequence[Bounds]) -> Bounds: + if isinstance(t, Const): + return t.value, t.value + if isinstance(t, Param): + if self.params is None or t.name not in self.params: + return None, None + v = param_value(self.graph, self.params, t.name) + return v, v + if isinstance(t, Pid): + return 0, None + if isinstance(t, NumPrograms): + return 1, None + if isinstance(t, Arange): + return (t.start, t.end - 1) if t.start < t.end else (None, None) + if isinstance(t, (Cmp, BoolBin, Not)): + return 0, 1 # a Bool, read as 0/1 in an integer position + if isinstance(t, IntCast): + return kids[0] # lowered as its operand + if isinstance(t, TSelect): + (lo_a, hi_a), (lo_b, hi_b) = kids[1], kids[2] + return ( + None if lo_a is None or lo_b is None else min(lo_a, lo_b), + None if hi_a is None or hi_b is None else max(hi_a, hi_b), + ) + if isinstance(t, Bin): + return _bin_bounds(t.op.lstrip("u"), *kids[0], *kids[1]) + return None, None + + def __call__(self, t: Bin, a: ArithRef, b: ArithRef) -> ArithRef: + lo_a, hi_a = self.bounds(t.a) + lo_b, hi_b = self.bounds(t.b) + quotient = t.op.lstrip("u") == "//" + positive_b = lo_b is not None and lo_b > 0 + negative_b = hi_b is not None and hi_b < 0 + if lo_a is not None and lo_a >= 0: + if quotient: + return a / b + if positive_b or negative_b: + return a % b + return If(b == 0, a, a % b) + if (positive_b or negative_b) and hi_a is not None and hi_a <= 0: + d = b if positive_b else -b + if not quotient: + return -((-a) % d) + return -((-a) / d) if positive_b else (-a) / d + return truncating_divrem(t, a, b) + + +def _bin_bounds(op: str, lo_a: Any, hi_a: Any, lo_b: Any, hi_b: Any) -> Bounds: + """Bounds on ``a op b`` (``op`` signed; the unsigned ops lower as their + signed twins) from the operands'.""" + known = lo_a is not None and hi_a is not None and lo_b is not None + known = known and hi_b is not None + if op == "+": + return ( + None if lo_a is None or lo_b is None else lo_a + lo_b, + None if hi_a is None or hi_b is None else hi_a + hi_b, + ) + if op == "-": + return ( + None if lo_a is None or hi_b is None else lo_a - hi_b, + None if hi_a is None or lo_b is None else hi_a - lo_b, + ) + if op == "min": + highs = [v for v in (hi_a, hi_b) if v is not None] + return ( + None if lo_a is None or lo_b is None else min(lo_a, lo_b), + min(highs) if highs else None, + ) + if op == "max": + lows = [v for v in (lo_a, lo_b) if v is not None] + return ( + max(lows) if lows else None, + None if hi_a is None or hi_b is None else max(hi_a, hi_b), + ) + if op == "*": + if known: + products = [x * y for x in (lo_a, hi_a) for y in (lo_b, hi_b)] + return min(products), max(products) + if lo_a == hi_a == 0 or lo_b == hi_b == 0: + return 0, 0 + nonneg_a, nonneg_b = ( + lo_a is not None and lo_a >= 0, + lo_b is not None and lo_b >= 0, + ) + nonpos_a, nonpos_b = ( + hi_a is not None and hi_a <= 0, + hi_b is not None and hi_b <= 0, + ) + if nonneg_a and nonneg_b: + return lo_a * lo_b, None + if nonpos_a and nonpos_b: + return hi_a * hi_b, None + if nonneg_a and nonpos_b: + return None, lo_a * hi_b + if nonpos_a and nonneg_b: + return None, hi_a * lo_b + return None, None + if op == "//": + if lo_a is not None and lo_a >= 0 and lo_b is not None and lo_b > 0: + return ( + 0 if hi_b is None else lo_a // hi_b, + None if hi_a is None else hi_a // lo_b, + ) + return None, None + if op == "%": + if lo_b is not None and hi_b is not None and lo_b > 0: + return ( + 0 if lo_a is not None and lo_a >= 0 else 1 - hi_b, + 0 if hi_a is not None and hi_a <= 0 else hi_b - 1, + ) + if lo_a is not None and lo_a >= 0: + # the truncating remainder has the dividend's sign (a at b == 0) + return 0, None + return None, None + return None, None + + +# Representation only: a non-affine table up to this many elements is read +# as an If-chain, a larger one through the array alone (same values). +_SMALL_SNAPSHOT_ELEMENTS = 32 + + +def snapshot_lookup(values: Sequence[int], index: ArithRef, fallback: Any) -> Any: + """``values[index]`` as arithmetic or a small finite choice where that + is exact, else ``fallback`` (the Select over the snapshot array, whose + equalities the solver asserts). Outside the table the value stays the + fallback: no value is invented for an element that was not read.""" + if not values: + return fallback + index = simplify(index) + try: + concrete = index.as_long() + except AttributeError: + concrete = None + if concrete is not None: + return IntVal(values[concrete]) if 0 <= concrete < len(values) else fallback + step = values[1] - values[0] if len(values) > 1 else 0 + if all(v == values[0] + i * step for i, v in enumerate(values)): + return If( + And(index >= 0, index < len(values)), + IntVal(values[0]) + index * step, + fallback, + ) + if len(values) > _SMALL_SNAPSHOT_ELEMENTS: + return fallback + value = fallback + for i in reversed(range(len(values))): + value = If(index == i, IntVal(values[i]), value) + return value + + +def terms_below( + roots: Sequence[object], graph: AccessGraph, *, loops: bool = True +) -> Iterator[object]: + """Every term reachable from ``roots`` (the lowering's ``children``, and + a Loaded's offset, mask and other), each once by identity, pre-order; + without ``loops``, not into a loop variable's bounds. Iterative.""" + seen: set[int] = set() + stack = [r for r in reversed(roots) if r is not None] + while stack: + t = stack.pop() + if id(t) in seen: + continue + seen.add(id(t)) + yield t + if isinstance(t, Loaded): + kids: tuple = tuple(k for k in (t.offset, t.mask, t.other) if k is not None) + elif isinstance(t, LoopVar) and not loops: + kids = () + else: + kids = children(t, graph) + stack.extend(reversed(kids)) + + +def loaded_leaves( + roots: Sequence[object], graph: AccessGraph, *, loops: bool = True +) -> list[Loaded]: + """The Loaded leaves ``roots`` read (each placement once).""" + return [t for t in terms_below(roots, graph, loops=loops) if isinstance(t, Loaded)] + + +def observed_leaves(roots: Sequence[object], graph: AccessGraph) -> set[int]: + """The access indices of the observations ``roots`` read.""" + return { + t.access_index for t in terms_below(roots, graph) if isinstance(t, Observed) + } + + +class RaceLeaves: + """The leaves of one encoding (see the module docstring): one launch + reading (``params`` None: any parameters), one Lowerer, and the state + the records read back (the arange dict, the snapshot equalities, the + loads left free).""" + + ctx = None + + def __init__( + self, + graph: AccessGraph, + params: Mapping[str, int] | None, + tensors: Mapping[str, LaunchTensor], + *, + blind: bool = False, + unusable: Mapping[str, str] | None = None, + ) -> None: + self.graph = graph + self.params = params + self.tensors = tensors + # The content-free reading: no snapshot is used. + self.blind = blind + # Sources whose contents may change during the launch (see + # global_records.written_sources). + self.unusable = dict(unusable or {}) + self.arange_dict: dict[Any, Any] = {} + self.iterations: dict[str, ArithRef] = {} + self._lanes: dict[tuple[int, int], ArithRef] = {} + self._observed: dict[int, ArithRef] = {} + self._params: dict[str, ArithRef] = {} + # Snapshot arrays by (base, width), their equalities, and per-load + # free arrays (copy-local: masked-off or unread lanes). + self._snapshots: dict[tuple[str, int], tuple[Any, tuple[int, ...]]] = {} + self.assumptions: list[BoolRef] = [] + self._pads: dict[int, Any] = {} + self.used_snapshot = False + # Loads whose value is free, with why. + self.free_loaded: dict[int, str] = {} + # The lowering reaches its leaves through a proxy, so the leaves and + # their terms are freed without the cyclic GC. + self.lowering = Lowerer( + graph, weakref.proxy(self), divrem=BoundedDivRem(graph, params) + ) + + # ── lowering ── + + def value(self, term: object) -> ArithRef: + return self.lowering.value(term) + + def cond(self, term: object) -> BoolRef: + return self.lowering.cond(term) + + @property + def lanes(self) -> Mapping[tuple[int, int], ArithRef]: + """The lane variables made so far, by ``(dim, extent)``.""" + return self._lanes + + @property + def pads(self) -> tuple[Any, ...]: + return tuple(self._pads.values()) + + def observation(self, index: int) -> ArithRef: + v = self._observed.get(index) + if v is None: + v = self._observed[index] = Int(f"ttir_obs_{index}") + return v + + def lane(self, dim: int, extent: int) -> ArithRef: + key = (dim, extent) + v = self._lanes.get(key) + if v is None: + v = self._lanes[key] = Int(f"ttir_lane_d{dim}_n{extent}") + # The solver's arange dict: key[0:2] is the lane's range, which + # it rebuilds per copy; the value's first item the variable. + self.arange_dict[(0, extent, "ttir", dim)] = (v, None) + return v + + # ── TermLeaves ── + + def param(self, t: Param) -> ArithRef: + if self.params is None: + v = self._params.get(t.name) + if v is None: + v = self._params[t.name] = Int(f"ttir_param_{t.name}") + return v + if t.name not in self.params: + raise RaceRefusal( + RaceKind.MISSING_BINDING, + f"scalar argument {t.name!r} has no launch binding", + ) + return IntVal(param_value(self.graph, self.params, t.name)) + + def pid(self, t: Pid) -> ArithRef: + return pids.pid_var(t.axis) + + def num_programs(self, t: NumPrograms) -> ArithRef: + return pids.grid_var(t.axis) + + def arange(self, t: Arange) -> ArithRef: + return IntVal(t.start) + self.lane(t.dim, t.end - t.start) + + def iteration(self, loop_ssa: str) -> ArithRef: + k = self.iterations.get(loop_ssa) + if k is None: + raise ValueError( + f"kernel {self.graph.kernel_name!r}: loop {loop_ssa!r} is not bound" + ) + return k + + def observed(self, t: Observed) -> ArithRef: + return self.observation(t.access_index) + + def data_dep(self, t: DataDep) -> ArithRef: + raise RaceRefusal(RaceKind.UNMODELED_VALUE, f"an unmodeled value ({t.why})") + + def loaded(self, t: Loaded) -> ArithRef: + """``If(mask, snapshot[offset], other)`` at the lane: a masked-off + lane holds ``other``, or, without one, a free value. The offset's + domain is not guarded here: each record that reads the value + carries ``mask -> 0 <= offset < numel`` (``domain_premises``).""" + index = t.access_index + off = self.value(t.offset) + access = self.graph.accesses[index] + problem = self.snapshot_problem(access) + if problem is not None: + self.free_loaded.setdefault(index, problem) + return Select(self.pad(index), off) + arr, values = self.snapshot(access) + self.used_snapshot = True + value = snapshot_lookup(values, off, Select(arr, off)) + if t.mask is None: + return value + other = ( + self.value(t.other) if t.other is not None else Select(self.pad(index), off) + ) + return If(self.cond(t.mask), value, other) + + # ── snapshots ── + + def pad(self, index: int) -> Any: + pad = self._pads.get(index) + if pad is None: + pad = self._pads[index] = Array(f"ttir_pad_{index}", IntSort(), IntSort()) + return pad + + def snapshot_problem(self, access: AccessEvent) -> str | None: + """Why the load ``access`` has no usable snapshot, or None. A + snapshot is read only at the load's own width (an i1 load of a bool + tensor included): a load through a pointer cast to another width + would read the bytes another way.""" + if self.params is None: + return "no launch at the any-parameters rung" + if self.blind: + return "content-free attempt" + base = access.base_param + tensor = self.tensors.get(base) + if tensor is None: + return "not captured" + if base in self.unusable: + return self.unusable[base] + facts = tensor.facts + if not facts.contiguous: + return "non-contiguous" + if tensor.values is None: + return tensor.why or "no snapshot" + if facts.allocation_interval() != ( + facts.data_ptr, + facts.data_ptr + facts.numel * facts.elem_size, + ): + # Consumers assume 0 <= offset < numel; a view of a larger + # allocation can legally read outside the captured values. + return "the snapshot covers only a view of its allocation" + bits = facts.elem_size * 8 + if access.elem_bits != bits and not ( + facts.dtype == "torch.bool" and access.elem_bits == 1 + ): + return f"read as i{access.elem_bits}, but the tensor holds {facts.dtype}" + if facts.dtype.startswith("torch.uint") and any( + v >= 1 << (bits - 1) for v in tensor.values + ): + # The IR's signless integer reads them negative: an extui of one + # would fail its width obligation, though the program computes + # it exactly. + return ( + f"{facts.dtype} values of 2**{bits - 1} or more, which the " + f"signed i{bits} reading does not keep" + ) + return None + + def snapshot(self, access: AccessEvent) -> tuple[Any, tuple[int, ...]]: + """The snapshot array of the load's tensor at its width, with its + values read signed at that width (their equalities join + ``assumptions`` once).""" + bits = access.elem_bits + key = (access.base_param, bits) + found = self._snapshots.get(key) + if found is not None: + return found + tensor = self.tensors[access.base_param] + assert tensor.values is not None + values = tuple(signed(int(v), bits) for v in tensor.values) + arr = Array(f"ttir_snap_{access.base_param}_i{bits}", IntSort(), IntSort()) + self.assumptions += [ + Select(arr, IntVal(i)) == IntVal(v) for i, v in enumerate(values) + ] + found = self._snapshots[key] = (arr, values) + return found + + def domain_premises(self, leaves: Sequence[Loaded]) -> tuple[BoolRef, ...]: + """For each snapshotted load among ``leaves``: where its mask holds, + its offset is inside the table (an instance whose load reads outside + its source is outside the model, as its own record says). A free + load needs none.""" + out: list[BoolRef] = [] + for leaf in leaves: + access = self.graph.accesses[leaf.access_index] + if self.snapshot_problem(access) is not None: + continue + numel = self.tensors[access.base_param].facts.numel + off = self.value(leaf.offset) + inside = And(off >= 0, off < numel) + out.append( + inside if leaf.mask is None else Implies(self.cond(leaf.mask), inside) + ) + return tuple(out) diff --git a/tilelens/clients/race_detector/data.py b/tilelens/clients/race_detector/data.py index d2d2039c..bd0b71b4 100644 --- a/tilelens/clients/race_detector/data.py +++ b/tilelens/clients/race_detector/data.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass, field from enum import Enum from typing import Any, Literal @@ -7,6 +8,7 @@ import torch from ...core.data import Op +from ...ir.verdict import ConfigVerdict, IRVerdict MemorySem = Literal["plain", "relaxed", "acquire", "release", "acq_rel"] @@ -187,12 +189,82 @@ def __post_init__(self) -> None: object.__setattr__(self, "race_type", RaceType(self.race_type).value) +@dataclass(frozen=True) +class CompiledRaceRecord(RaceRecord): + """A race the compiled race detector found in one config of a launch + (``config``: its config kwargs, as RaceConfigVerdict.config). An + access's ``element`` is counted from the launch tensor's data pointer.""" + + config: Mapping[str, Any] = field(default_factory=dict) + + __hash__ = None # type: ignore[assignment] + + def __post_init__(self) -> None: + super().__post_init__() + object.__setattr__(self, "config", dict(self.config)) + + +@dataclass(frozen=True) +class RaceConfigVerdict(ConfigVerdict): + """The compiled race detector's verdict on one config of a launch. + + ``status`` is "ok" (proved race-free), "races", "unsupported" or + "vacuous" (no execution satisfies the premises, so no proof is + claimed). For "ok", ``rung`` is the proof's: "T0" (any scalar + parameters and any grid), "T1" (this launch's parameters, any grid) + or "T1-launch" (this launch's parameters and grid). "Any grid" is any + extent along the grid axes the kernel reads (a program id or the + program count) and this launch's extent along the others (a kernel + with atomics: any extent along those too, except at T1 with a launch + grid). ``qualifiers`` name what a proof or a race assumes or read: + "assumes no i32 overflow" (the widths whose obligations can fail in + the proof's domain; they hold at this launch), "assumes termination" + (the kernel spins on a flag), "reads tensor contents" (it holds for + the launch's pre-launch contents of the tensors it read). ``grid_fragile`` are the races the + query found at some other grid behind a "T1-launch" proof, or behind an + "unsupported" config whose query at the launch's grid was not decided: + evidence only, never races of this launch.""" + + rung: str | None = None + qualifiers: tuple[str, ...] = () + grid_fragile: tuple[CompiledRaceRecord, ...] = () + + __hash__ = None # type: ignore[assignment] + + def __post_init__(self) -> None: + super().__post_init__() + object.__setattr__(self, "qualifiers", tuple(self.qualifiers)) + object.__setattr__(self, "grid_fragile", tuple(self.grid_fragile)) + + +@dataclass(frozen=True) +class RaceVerdict(IRVerdict): + """The compiled race detector's verdict on one traced launch: ``status`` + "races" if a config races, else "unsupported" if a config was not + decided, else "vacuous" if a config's proof would be vacuous, else + "ok". For "ok", ``scope`` says what the weakest config's proof holds + for ("any-params-any-grid", "this-params-any-grid" or + "this-params-this-grid"; "any grid" as RaceConfigVerdict.rung says), + and ``qualifiers`` gathers the configs'.""" + + qualifiers: tuple[str, ...] = () + + __hash__ = None # type: ignore[assignment] + + def __post_init__(self) -> None: + super().__post_init__() + object.__setattr__(self, "qualifiers", tuple(self.qualifiers)) + + __all__ = [ "AccessEventRecord", "AtomicKind", + "CompiledRaceRecord", "MemorySem", "RaceAccess", + "RaceConfigVerdict", "RaceRecord", "RaceReport", "RaceType", + "RaceVerdict", ] diff --git a/tilelens/clients/race_detector/hb_common.py b/tilelens/clients/race_detector/hb_common.py index 50180f3f..8881bff6 100644 --- a/tilelens/clients/race_detector/hb_common.py +++ b/tilelens/clients/race_detector/hb_common.py @@ -2,11 +2,18 @@ from __future__ import annotations +import threading from typing import Any, Callable, Iterable, Iterator from z3 import And, BoolSort, BoolVal, Not, Or, is_false, is_true, simplify, substitute from z3.z3 import BoolRef, IntNumRef +# Z3's main context, where the race detectors build and solve their +# records, is shared and not thread-safe: a race detector holds this lock +# around its solve (and the compiled one around building its records too), +# and drops every term it made before releasing it. +MAIN_CONTEXT_LOCK = threading.RLock() + class UnsupportedSymbolicRaceQuery(Exception): """Raised when a record cannot be lowered to Z3 or a query is undecided. diff --git a/tilelens/clients/race_detector/race_detector.py b/tilelens/clients/race_detector/race_detector.py index a71eb789..d5ade71c 100644 --- a/tilelens/clients/race_detector/race_detector.py +++ b/tilelens/clients/race_detector/race_detector.py @@ -59,6 +59,7 @@ ) from .data import ( AccessEventRecord, + CompiledRaceRecord, MemorySem, RaceAccess, RaceRecord, @@ -66,6 +67,7 @@ RaceType, ) from .hb_common import ( + MAIN_CONTEXT_LOCK, UnsupportedSymbolicRaceQuery, apply_sub, normalize_copy_local_vars, @@ -149,9 +151,11 @@ class PendingEvent(PendingCheck): class RaceDetector(Client): """Factory class that returns the concrete race-detector implementation: :class:`NullRaceDetector` when ``cfg.enable_race_detector`` is off, - otherwise :class:`SymbolicRaceDetector`, the eager detector, which - reasons about global memory across program instances from one - interpreter-driven symbolic capture. + otherwise ``RaceDetector()`` the eager one, :class:`SymbolicRaceDetector`, + which reasons about global memory across program instances from one + interpreter-driven symbolic capture, and ``RaceDetector(compile=True)`` + the compiled one (``CompiledRaceDetector``, a virtual subclass), which + reads each launch's host-compiled TTIR. """ NAME = "race_detector" @@ -162,6 +166,15 @@ def __new__(cls: type[RaceDetectorT], *args: Any, **kwargs: Any) -> RaceDetector if cls is RaceDetector: # The disable flag wins over the mode, as for the Sanitizer: # RaceDetector(compile=True) is a NullRaceDetector when it is off. + if kwargs.pop("compile", False) and cfg.enable_race_detector: + from .compiled.client import CompiledRaceDetector + + # Only a virtual RaceDetector subclass, so Python does not + # call its __init__ after __new__: call it here, without + # ``compile``. + compiled = object.__new__(CompiledRaceDetector) + CompiledRaceDetector.__init__(compiled, *args, **kwargs) + return cast(RaceDetectorT, compiled) target_cls = cast( type["RaceDetector"], SymbolicRaceDetector if cfg.enable_race_detector else NullRaceDetector, @@ -317,8 +330,12 @@ def print_race_record(record: RaceRecord) -> None: """Print one race the detector found.""" print("============================================================") print(" Data Race Detected ") + if isinstance(record, CompiledRaceRecord): + print(" (compiled race detector)") print("============================================================") print(f"Race Type: {record.race_type}") + if isinstance(record, CompiledRaceRecord) and record.config: + print(f"Config: {record.config}") for label, access, program in ( ("First", record.first, record.programs[0]), ("Second", record.second, record.programs[1]), @@ -350,9 +367,9 @@ class SymbolicRaceDetector(RaceDetector, SymbolicClient): # constructs this class with it. def __init__(self, abort_on_error: bool = False, *, compile: bool = False): if compile: - raise NotImplementedError( - "RaceDetector(compile=True): the compiled race detector is not " - "available yet; RaceDetector() is the eager one." + raise TypeError( + "SymbolicRaceDetector is the eager race detector; the compiled " + "one is RaceDetector(compile=True)" ) super().__init__(abort_on_error=abort_on_error) self.records: list[AccessEventRecord] = [] @@ -1032,6 +1049,10 @@ def finalize(self) -> list: self._clear_launch_runtime() def _solve(self) -> None: + with MAIN_CONTEXT_LOCK: + self._solve_locked() + + def _solve_locked(self) -> None: try: solver = TwoCopySymbolicHBSolver( self.records, diff --git a/tilelens/clients/race_detector/two_copy_symbolic_hb_solver.py b/tilelens/clients/race_detector/two_copy_symbolic_hb_solver.py index 6945db82..76ce228e 100644 --- a/tilelens/clients/race_detector/two_copy_symbolic_hb_solver.py +++ b/tilelens/clients/race_detector/two_copy_symbolic_hb_solver.py @@ -44,7 +44,8 @@ Z3 context: Every term lives in Z3's main context. Records must be lowered there (a term from another context fails with Z3's "Context mismatch"), and since - that context is shared and not thread-safe, the caller serializes solves. + that context is shared and not thread-safe, the caller serializes solves + (``hb_common.MAIN_CONTEXT_LOCK``). Intra-instance duplicate lanes: Cross-copy queries assert ``different_blocks``, so they can never witness @@ -152,6 +153,7 @@ ) from z3.z3 import BoolRef, IntNumRef, ModelRef +from .. import pids from .conflict_simplification import ( _PureSelectExpressionCache, conflict_impossible, @@ -258,11 +260,8 @@ class SymbolicMemoryEvent: causal_inputs: tuple[Any, ...] = () -def _import_symbolic_expr_pids(): - # Local import: SymbolicExpr is a heavy module; keep tests cheap. - from ..symbolic_engine import SymbolicExpr - - return (SymbolicExpr.PID0, SymbolicExpr.PID1, SymbolicExpr.PID2) +def _main_context_pids(): + return (pids.PID0, pids.PID1, pids.PID2) def _is_symbolic_dim(d: Any) -> bool: @@ -1412,7 +1411,7 @@ def _make_pid_vars(): @staticmethod def _make_pid_subs(pid_a, pid_b): - orig = _import_symbolic_expr_pids() + orig = _main_context_pids() sub_a = tuple((orig[i], pid_a[i]) for i in range(3)) sub_b = tuple((orig[i], pid_b[i]) for i in range(3)) return sub_a, sub_b @@ -2145,7 +2144,7 @@ def _fresh_context(self, tag: str) -> tuple[CopyContext, BoolRef]: ctx = CopyContext( label=tag, pid=pid, - pid_substitutions=tuple(zip(_import_symbolic_expr_pids(), pid)), + pid_substitutions=tuple(zip(_main_context_pids(), pid)), arange_substitutions=tuple(arange_subs), arange_constraints=(), copy_local_substitutions=copy_local, diff --git a/tilelens/clients/sanitizer/compiled/client.py b/tilelens/clients/sanitizer/compiled/client.py index cf15b8de..8effdfbe 100644 --- a/tilelens/clients/sanitizer/compiled/client.py +++ b/tilelens/clients/sanitizer/compiled/client.py @@ -83,7 +83,7 @@ import sys from collections import OrderedDict -from collections.abc import Hashable, Mapping +from collections.abc import Hashable from dataclasses import replace from typing import Any, ClassVar @@ -91,7 +91,13 @@ from ....core.config import config as cfg from ....core.data import Load, Store from ....core.host_compile import parse_ir_target -from ....ir.capture import ArtifactLog, CompiledSpecialization, ParseCache +from ....ir.capture import ( + ArtifactLog, + CompiledSpecialization, + ParseCache, + configs_of, + saveable_config, +) from ....ir.client import IRClient from ....ir.failures import ( compiles_nowhere, @@ -100,7 +106,7 @@ nothing_compiled, nowhere_note, ) -from ....ir.launch import LaunchBinding, TensorFacts, tensor_facts +from ....ir.launch import LaunchBinding from ....ir.verdict import ConfigVerdict, IRVerdict, Refusal from ....utils.traceback_utils import location_to_traceback_info from ..data import CompiledSanitizerRecord @@ -200,12 +206,12 @@ def analyze_launch(self, log: ArtifactLog) -> tuple[list, IRVerdict]: notes: list[str] = [] for failure in log.failures: if compiles_nowhere(failure.error): - notes.append(nowhere_note(failure, _saveable(failure.config))) + notes.append(nowhere_note(failure, saveable_config(failure.config))) else: failed.append( ConfigVerdict( None, - _saveable(failure.config), + saveable_config(failure.config), "unsupported", failure_refusal(failure, _NAME_A_TARGET), ) @@ -255,9 +261,9 @@ def _check_specialization( self, spec: CompiledSpecialization ) -> list[tuple[list[CompiledSanitizerRecord], ConfigVerdict]]: """One (records, ConfigVerdict) per config among ``spec``'s - bindings, in the order first seen (see _configs_of): each config is + bindings, in the order first seen (see configs_of): each config is checked against its own bindings only.""" - configs = _configs_of(spec) + configs = configs_of(spec) try: graph = self._graph_of(spec) except Exception as exc: @@ -362,83 +368,6 @@ def _describe(exc: BaseException) -> str: return f"{type(exc).__name__}: {exc}" -_PLAIN = (str, int, float, bool, type(None)) - - -def _tensor_of(value: Any) -> TensorFacts | None: - """``value``'s tensor facts, None if it is no (readable) tensor.""" - if not hasattr(value, "data_ptr"): - return None - try: - return tensor_facts(value) - except Exception: - return None - - -def _saveable(config: Mapping[str, Any]) -> dict[str, Any]: - """The config kwargs, each value a saved trace cannot hold put as text: - a tensor (e.g. a heuristic's view) described by its facts, never its - data, anything else (e.g. a heuristic's tl.dtype) as its repr.""" - - def describe(value: Any) -> Any: - if isinstance(value, _PLAIN): - return value - facts = _tensor_of(value) - if facts is None: - return repr(value) - return ( - f"" - ) - - return {name: describe(value) for name, value in config.items()} - - -def _config_key(config: Mapping[str, Any]) -> Hashable: - """What tells two bindings' config kwargs apart: a plain value by type - and value, a tuple item by item, a tensor by its data_ptr, shape, - strides and dtype (never its data), any other hashable value by type - and equality, anything else by identity (the bindings hold the values - while the keys are compared).""" - - def key(value: Any) -> Hashable: - if isinstance(value, float): - return (float, value.hex()) # -0.0 apart from 0.0, NaN equal - if isinstance(value, _PLAIN): - return (type(value), value) - if isinstance(value, tuple): - return (type(value), tuple(key(item) for item in value)) - facts = _tensor_of(value) - if facts is not None: - return ("tensor", facts.data_ptr, facts.shape, facts.strides, facts.dtype) - try: - hash(value) - except Exception: - return ("id", id(value)) - return (type(value), value) - - return tuple(sorted((name, key(value)) for name, value in config.items())) - - -def _configs_of( - spec: CompiledSpecialization, -) -> list[tuple[dict[str, Any], list[LaunchBinding]]]: - """``spec``'s bindings grouped by their config kwargs (see _config_key), - each group with its saveable config, in the order first seen, so the - first group's config is ``spec.config``. Configs that compile to one - kernel (e.g. differing only in a runtime int kwarg) share a - specialization but not their bindings.""" - groups: dict[Hashable, tuple[dict[str, Any], list[LaunchBinding]]] = {} - for binding in spec.bindings: - key = _config_key(binding.config) - group = groups.get(key) - if group is None: - groups[key] = (_saveable(binding.config), [binding]) - else: - group[1].append(binding) - return list(groups.values()) or [(_saveable(spec.config), [])] - - def _record( finding: Finding, kernel_name: str, binding: LaunchBinding, config: dict[str, Any] ) -> CompiledSanitizerRecord: diff --git a/tilelens/clients/sanitizer/compiled/oob.py b/tilelens/clients/sanitizer/compiled/oob.py index 0a30f328..ea666e1e 100644 --- a/tilelens/clients/sanitizer/compiled/oob.py +++ b/tilelens/clients/sanitizer/compiled/oob.py @@ -119,6 +119,7 @@ width_obligations, ) from ....ir.verdict import Refusal, SourceLocation +from ...pids import PID_AXES, pid_name, pid_var from ..data import CompiledFindingKind @@ -374,7 +375,7 @@ def __init__( self.ctx = ctx # Range premises of the variables created so far. self.premises: list[BoolRef] = [] - self.pids = tuple(Int(f"pid_{axis}", ctx) for axis in range(3)) + self.pids = tuple(pid_var(axis, ctx) for axis in PID_AXES) for pid, size in zip(self.pids, grid): self.premises += [pid >= 0, pid < size] # (dim, extent) -> the lane's position along that dim (see arange). @@ -1062,7 +1063,7 @@ def witness(self, env: _Env, model: ModelRef) -> dict[str, int]: def val(v: ArithRef) -> int: return model.eval(v, model_completion=True).as_long() - out = {f"pid_{axis}": val(pid) for axis, pid in enumerate(env.pids)} + out = {pid_name(axis): val(pid) for axis, pid in enumerate(env.pids)} for name, lane in env.lanes.items(): out[name] = val(lane) if env.k is not None: diff --git a/tilelens/clients/symbolic_engine.py b/tilelens/clients/symbolic_engine.py index 67f1c2aa..b4a3f1d7 100644 --- a/tilelens/clients/symbolic_engine.py +++ b/tilelens/clients/symbolic_engine.py @@ -106,6 +106,7 @@ type_spec, unpack_type_spec, ) +from . import pids from .utils import ( check_storage_contiguous, get_physical_addr_from_tensor_slice, @@ -457,9 +458,9 @@ class SymbolicExpr: ) SUPPORTED_OP_SET: ClassVar[frozenset[str]] = frozenset(SUPPORTED_OPS) - PID0: ClassVar[ArithRef] = Int("pid_0") - PID1: ClassVar[ArithRef] = Int("pid_1") - PID2: ClassVar[ArithRef] = Int("pid_2") + PID0: ClassVar[ArithRef] = pids.PID0 + PID1: ClassVar[ArithRef] = pids.PID1 + PID2: ClassVar[ArithRef] = pids.PID2 ARANGE_DICT: ClassVar[ dict[tuple[Any, ...], tuple[ArithRef, ConstraintConjunction]] @@ -1055,8 +1056,7 @@ def __init__(self, op: str, axis: Any): self.dtype = INT32 def _node_label_core(self) -> str: - axis_val = self.axis.to_py() - return f"pid_{axis_val}" + return pids.pid_name(self.axis.to_py()) def _to_z3_impl(self) -> tuple[Z3Expr, ConstraintConjunction]: axis_val = self.axis.to_py() diff --git a/tilelens/ir/capture.py b/tilelens/ir/capture.py index ddac874e..7b75e03e 100644 --- a/tilelens/ir/capture.py +++ b/tilelens/ir/capture.py @@ -5,6 +5,9 @@ metadata plus the LaunchBindings seen for it, and every compile failure. ``ParseCache`` runs a reader once per distinct text and keeps what it gave, a graph or a typed refusal, so a refusal's kind survives cache hits. +``configs_of`` groups a specialization's bindings by the config kwargs they +were called with, each config as values a saved trace can hold +(``saveable_config``). Mechanism only: which specialization counts, what an absent stage means and what a refusal or an error becomes are the client's calls. Neither class @@ -24,7 +27,13 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any -from .launch import LaunchBinding, bind_launch, config_kwargs +from .launch import ( + LaunchBinding, + TensorFacts, + bind_launch, + config_kwargs, + tensor_facts, +) if TYPE_CHECKING: from ..core.client import LaunchCall, LaunchEvent @@ -139,6 +148,7 @@ def reset(self, call: LaunchCall | None = None) -> None: Hashable, tuple[CompiledArtifacts, list[LaunchBinding]] ] = {} self._failures: list[CompileFailure] = [] + self._last: LaunchBinding | None = None def record(self, event: LaunchEvent) -> LaunchBinding: """Record one before_launch event; returns the binding recorded for @@ -150,8 +160,14 @@ def record(self, event: LaunchEvent) -> LaunchBinding: self._compiled[event.specialization] = (artifacts, [binding]) else: entry[1].append(binding) + self._last = binding return binding + @property + def last_binding(self) -> LaunchBinding | None: + """The binding the latest ``record`` recorded (None since reset).""" + return self._last + def record_failure(self, event: LaunchEvent) -> None: self._failures.append( CompileFailure( @@ -278,3 +294,82 @@ def get(self, text: str, /, **options: Hashable) -> ParseOutcome: outcome = ParseOutcome(refusal=_without_frames(exc, outer)) self._outcomes[key] = outcome return outcome + + +# ─────────────── the configs of a specialization ─────────────── + +_PLAIN = (str, int, float, bool, type(None)) + + +def _tensor_of(value: Any) -> TensorFacts | None: + """``value``'s tensor facts, None if it is no (readable) tensor.""" + if not hasattr(value, "data_ptr"): + return None + try: + return tensor_facts(value) + except Exception: + return None + + +def saveable_config(config: Mapping[str, Any]) -> dict[str, Any]: + """The config kwargs, each value a saved trace cannot hold put as text: + a tensor (e.g. a heuristic's view) described by its facts, never its + data, anything else (e.g. a heuristic's tl.dtype) as its repr.""" + + def describe(value: Any) -> Any: + if isinstance(value, _PLAIN): + return value + facts = _tensor_of(value) + if facts is None: + return repr(value) + return ( + f"" + ) + + return {name: describe(value) for name, value in config.items()} + + +def _config_key(config: Mapping[str, Any]) -> Hashable: + """What tells two bindings' config kwargs apart: a plain value by type + and value, a tuple item by item, a tensor by its data_ptr, shape, + strides and dtype (never its data), any other hashable value by type + and equality, anything else by identity (the bindings hold the values + while the keys are compared).""" + + def key(value: Any) -> Hashable: + if isinstance(value, float): + return (float, value.hex()) # -0.0 apart from 0.0, NaN equal + if isinstance(value, _PLAIN): + return (type(value), value) + if isinstance(value, tuple): + return (type(value), tuple(key(item) for item in value)) + facts = _tensor_of(value) + if facts is not None: + return ("tensor", facts.data_ptr, facts.shape, facts.strides, facts.dtype) + try: + hash(value) + except Exception: + return ("id", id(value)) + return (type(value), value) + + return tuple(sorted((name, key(value)) for name, value in config.items())) + + +def configs_of( + spec: CompiledSpecialization, +) -> list[tuple[dict[str, Any], list[LaunchBinding]]]: + """``spec``'s bindings grouped by their config kwargs (see _config_key), + each group with its saveable config, in the order first seen, so the + first group's config is ``spec.config``. Configs that compile to one + kernel (e.g. differing only in a runtime int kwarg) share a + specialization but not their bindings.""" + groups: dict[Hashable, tuple[dict[str, Any], list[LaunchBinding]]] = {} + for binding in spec.bindings: + key = _config_key(binding.config) + group = groups.get(key) + if group is None: + groups[key] = (saveable_config(binding.config), [binding]) + else: + group[1].append(binding) + return list(groups.values()) or [(saveable_config(spec.config), [])] diff --git a/tilelens/ir/lowering.py b/tilelens/ir/lowering.py index 646c8ffe..6ccdcd9a 100644 --- a/tilelens/ir/lowering.py +++ b/tilelens/ir/lowering.py @@ -16,7 +16,9 @@ * ``+ - *`` are unbounded Int arithmetic; ``//`` and ``%`` truncate toward zero (``arith.divsi`` / ``remsi``: the remainder takes the dividend's - sign); ``min`` / ``max`` pick an operand; + sign), encoded for any signs unless the client gives the Lowerer its own + ``divrem`` (an equal value it can encode more cheaply); ``min`` / + ``max`` pick an operand; * the unsigned ops (``u//``, ``u%``, ``umin``, ``umax``) and predicates (``ult``, ...) read as their signed twins, and an ``IntCast`` as its operand: exact only where the access's ``width_obligations`` hold, which @@ -137,6 +139,21 @@ def _divrem_ir(op: str, a: ArithRef, b: ArithRef) -> ArithRef: return q if op in ("//", "u//") else a - b * q +DIVREM_OPS = frozenset({"//", "u//", "%", "u%"}) + +# A client's encoding of a quotient or remainder (a Bin whose op is in +# DIVREM_OPS): (the term, its lowered dividend, its lowered divisor) -> a +# value equal to the IR's truncating one wherever the client's queries +# read it. +DivRem = Callable[[Bin, ArithRef, ArithRef], ArithRef] + + +def truncating_divrem(t: Bin, a: ArithRef, b: ArithRef) -> ArithRef: + """The default DivRem: the IR's truncating quotient or remainder, + encoded for any signs of ``a`` and ``b``.""" + return _divrem_ir(t.op, a, b) + + def _bin(op: str, a: ArithRef, b: ArithRef) -> ArithRef: # The unsigned ops read their operands unsigned; their width # obligations make both non-negative, where they equal the signed ones. @@ -146,8 +163,6 @@ def _bin(op: str, a: ArithRef, b: ArithRef) -> ArithRef: return a - b if op == "*": return a * b - if op in ("//", "u//", "%", "u%"): - return _divrem_ir(op, a, b) if op in ("min", "umin"): return If(a <= b, a, b) if op in ("max", "umax"): @@ -238,11 +253,18 @@ def fold( class Lowerer: """The lowering of one family of a client's queries: bound to one graph - and one leaves object (so one Z3 context), with one memo.""" + and one leaves object (so one Z3 context), with one memo; ``divrem`` + encodes ``//`` and ``%`` (by default for any signs).""" - def __init__(self, graph: AccessGraph, leaves: TermLeaves) -> None: + def __init__( + self, + graph: AccessGraph, + leaves: TermLeaves, + divrem: DivRem = truncating_divrem, + ) -> None: self.graph = graph self.leaves = leaves + self._divrem = divrem self._memo: dict[int, tuple[object, Any]] = {} def lower(self, term: object) -> Any: @@ -276,7 +298,10 @@ def _apply(self, t: object, kids: Sequence[Any]) -> Any: k = leaves.iteration(loop_ssa or _loop(self.graph, t, "").loop_ssa) return as_int(kids[0]) + k * as_int(kids[1]) if isinstance(t, Bin): - return _bin(t.op, as_int(kids[0]), as_int(kids[1])) + a, b = as_int(kids[0]), as_int(kids[1]) + if t.op in DIVREM_OPS: + return self._divrem(t, a, b) + return _bin(t.op, a, b) if isinstance(t, Cmp): return _cmp(t.pred, kids[0], kids[1]) if isinstance(t, BoolBin): diff --git a/tilelens/ir/ttir_reader.py b/tilelens/ir/ttir_reader.py index f631e5e9..e2263c62 100644 --- a/tilelens/ir/ttir_reader.py +++ b/tilelens/ir/ttir_reader.py @@ -547,6 +547,12 @@ class AccessEvent: # The enclosing scf.for loops (LoopInfo.loop_ssa), outermost first; # ``in_loop == bool(loops)`` (a hand-built graph may leave it empty). loops: tuple[str, ...] = () + # The tile shape of the access: the pointer operand's tensor type + # (``tensor<16x16x!tt.ptr>`` is (16, 16)); () for a scalar pointer + # and in a hand-built graph. The lanes of a dimension the offset does not + # vary along (a broadcast pointer, ``p + tl.zeros([B])``) share an + # address. + shape: tuple[int, ...] = () @dataclass(frozen=True) @@ -1500,6 +1506,7 @@ def access( deps=deps, block_path=op.path, loops=loops, + shape=parse_type(self.m.values[ptr_v].type).shape, ) ) diff --git a/tilelens/wrapper.py b/tilelens/wrapper.py index cf2d2eb8..cbfabe44 100644 --- a/tilelens/wrapper.py +++ b/tilelens/wrapper.py @@ -14,8 +14,8 @@ PROFILER_COMMAND = "tile-profiler" RACE_DETECTOR_COMMAND = "tile-race" -# tile-sanitizer's flag for the compiled sanitizer, given before the script -# name. +# tile-sanitizer's and tile-race's flag for their compiled mode, given before +# the script name. COMPILE_FLAG = "--compile" COMPILE_NOTE = ( f"[{COMPILE_FLAG}] kernel launches are compiled and checked, not run: " @@ -64,6 +64,11 @@ def race_detector_wrapper(kernel, *, frontend: str = "triton"): return tracer(kernel) +def compiled_race_detector_wrapper(kernel, *, frontend: str = "triton"): + tracer = tilelens.trace(client=RaceDetector(compile=True), frontend=frontend) + return tracer(kernel) + + def create_patched_jit(wrapper_func, original_jit, *, frontend: str): def _patched_jit(fn=None, **jit_kw): if fn is None: # @triton.jit(**opts) @@ -205,10 +210,18 @@ def apply_profiler(): def apply_race_detector(): """ - Apply the race detector wrapper to triton.jit and run the user script. + Apply the race detector wrapper to triton.jit and run the user script; with + ``--compile`` before the script, the compiled race detector's. """ _apply_wrapper( race_detector_wrapper, RACE_DETECTOR_COMMAND, - f"Usage: {RACE_DETECTOR_COMMAND} [args...]", + f"Usage: {RACE_DETECTOR_COMMAND} [{COMPILE_FLAG}] [args...]\n" + f" {COMPILE_FLAG} analyze each launch's compiled kernel " + "(RaceDetector(compile=True)) instead of interpreting it; kernels are " + "not run, so their outputs are never written. Kernels are compiled on " + "the host (no GPU needed) for TILELENS_IR_TARGET, by default cuda:89; " + "a config that fails to compile for it is reported as not analyzed, " + "and the script goes on.", + compile_wrapper=compiled_race_detector_wrapper, )