[FEAT] [RACE DETECTOR] Compiled mode: static shared-memory race detection over TTGIR - #476
Merged
Merged
Conversation
…tion over TTGIR Shared memory is invisible to the dynamic (interpreter-driven) mode — it is introduced by TritonGPU compiler passes. This adds the compile-mode counterpart from the design plan (race_detector_compiled_mode_plan.md): a static analyzer over the kernel specialization's TTGIR that either proves the cp.async pipeline race-free for ALL inputs, grids and trip counts, or returns a witness, or honestly reports unsupported. Components (triton_viz/clients/race_detector/compiled/): - ttgir_reader: line parser over the closed v1 op vocabulary (local_alloc/local_load/local_store/memdesc_index/ async_copy_global_to_local/async_commit_group/async_wait, scf.for iter_args/yield, the addi/cmpi/select rotation chains, loc aliases). Text instead of the pybind walk because attribute literals are opaque to the bindings and to_linear_layout SIGABRTs on shared encodings in the 3.6.0 wheel. Anything touching a memdesc outside the vocabulary (ttng/Hopper, barriers, nested loops) marks the kernel unsupported. - hb: the happens-before model. TTGIR has NO CTA barriers (Membar inserts bar.sync during lowering), so ordering comes from commit-group counting (async_wait num=N leaves at most the N newest groups outstanding) and the multibuffer rotation, whose (base + k) mod S closed form is extracted by abstract interpretation and CHECKED by exhaustive simulation of the parsed select chain — never trusted. Model boundary documented in the module docstring: v1 checks the RAW direction (copy -> load) under the Membar-barrier drift assumption; generic-generic pairs are Membar-ordered by construction. - layouts: blocked owner-coords and swizzled-shared XOR-offset closed forms (transcribed from the 3.6.x C++ sources), with the broadcast (tile > shape) case folded in as the bit-field modulo; differentially tested against an independent basis construction and, when available, triton's own LinearLayout. - smt_encoder: per (copy, load) pair, a Z3 query over symbolic iterations and a symbolic trip count — no loop unrolling: does the load's slot match a copy whose commit group the guarding wait does not cover? SAT yields a witness (iterations, slot, byte offset via the layout closed forms) and an optional SMT-LIB2 artifact; UNSAT over all pairs is a proof. Pathological inputs (SSA cycles, non-TTGIR text) degrade to unsupported, never crash or fake a proof. - client: acquires TTGIR through the runtime's own warmup (post_warmup_callback receives CompiledKernel.asm) so the analyzed IR carries the real divisibility specialization — a hand-built ASTSource silently disables the pipeliner. Analysis is cached per TTGIR hash; the interpreted grid run is skipped entirely (standalone client). Verified: the stock compiler pipelines for num_stages 2/3/4 are all proven race-free (the model extracts S/P/g/num per variant), the stock num=2 wait is exactly tight (num=3 already races), and every mutation class fires with a correct witness — weakened/deleted waits, shrunk stage dims, rotation off-by-one, dropped commit groups (exactly one report). Golden TTGIR dumps pin the printer format; dynamic-mode suites are untouched (234 passed).
…True) The compiled-mode backend was only reachable by importing CompiledRaceDetector directly. Route it through the public RaceDetector factory instead: RaceDetector(compile=True) dispatches __new__ to the compiled backend (flag on), the default / compile=False stays the dynamic SymbolicRaceDetector, and flag-off is NullRaceDetector either way. Extra keywords flow to the chosen backend (e.g. collect_smtlib=True). SymbolicRaceDetector.__init__ now tolerates the compile keyword: when __new__ returns a RaceDetector-subclass instance, Python re-invokes __init__ with the original kwargs (still carrying compile). The compiled backend is not a RaceDetector subclass, so no such re-invocation occurs for it. Trace-level tests use the new factory form behind an enable-flag fixture; the plan doc's mode-composition note matches the implemented switch.
… unsupported, never silent ok)
Review of the compiled-mode detector found several places where an
unmodeled construct yielded a confident "ok" instead of "unsupported" or
a report — the opposite of the tool's contract. This hardens each:
- async_wait operand tokens are now carried and used. The HB model gates
wait coverage per allocation: a wait that names explicit tokens but does
not await a load's allocation no longer "covers" that load by num count
alone. Resolution follows the iter_arg rotation (init + yield) and
wait-chaining; unresolvable tokens over-report (sound) rather than
proving. Dropping a wait operand now reports the uncovered allocation's
loads instead of reading as a proof. Stock pipeliner IR names tokens
consistent with num, so the proof is unchanged.
- async_commit_group raises UnsupportedTTGIR on a true regex non-match
(e.g. operand-style commit from an unmodeled printer); the bare
no-token form still parses as a real empty group. Previously a
non-match was silently swallowed as empty, corrupting commit-rank
accounting.
- A nested/conditional region inside the pipelined loop (scf.if, or a
"} else {" continuation) raises instead of mis-tracking the
loop/epilogue boundary via naive brace counting.
- A local_alloc after a local_dealloc (buffer reuse / allocation
aliasing, which v1 does not model) raises. A terminal dealloc (stock
epilogue cleanup) stays a clean proof.
- The compiled detector is STANDALONE: it skips the interpreted run via
pre_run_callback() == False, which is all()-combined and would suppress
a co-registered client's capture. ClientManager.add_clients now rejects
composing it with other clients, and the self-contradictory module
docstring is corrected.
- Documentation: the "byte-level proof over all inputs/grids" wording is
downgraded to what is actually solved — a wait-coverage proof for the
whole-tile cp.async RAW pipeline under the model boundary; the
byte_offset is a post-solve representative witness. Aliasing claims in
the plan are marked deferred/unsupported.
Tests: new coverage for the dropped-wait-token report, malformed commit
group, conditional region in loop, alloc-after-dealloc, and the
standalone-composition rejection. Stock proof and all existing mutation
tests unchanged.
…t bounds, stage-dim, report API, cache key) - Validate slots against stage geometry in build_pipeline_model: a ConstSlot must index an existing stage [0, stages) and a RotatingSlot must wrap at the stage count, for both copies and loads. Inconsistent geometry (e.g. a depth-1 buffer still indexed at slot 1) now fails closed as unsupported instead of being analyzed under a broken buffer model. - Allocation.has_stage_dim distinguishes a staged buffer (viewed by ttg.memdesc_index) from an un-staged one, even at stage depth 1. buffer_dims / stage_bytes key off it instead of stages > 1, so a single-buffered memdesc<1x...> no longer keeps a phantom leading dimension in the witness layout computation. - CompiledRaceReport.race_type reuses the dynamic detector's RaceType enum (RaceType.RAW) instead of a bare "RAW" string, so consumers can branch on race_type uniformly across both modes. - Cache analysis results under a stable SHA-256 of the TTGIR text rather than Python's process-randomized, non-collision-resistant built-in hash() — the cache holds proof/witness verdicts, a soundness boundary. - Document the standalone trade-off on the RaceDetector(compile=True) factory: the compiled backend cannot be composed with other clients. Tests: rewrite the shrunk-stage mutation as a well-formed single-buffer race (still reported) and add a separate inconsistent-geometry case that must be unsupported; add a disabled-flag factory test; assert race_type against RaceType.RAW. Stock proofs and all existing mutation tests unchanged.
…group ttg.async_commit_group always prints its !ttg.async.token result, and that result token is what an async_wait names (commit_by_result in hb.py keys on it). A commit group with no SSA result is therefore malformed and unreachable by any wait, so the reader now fails closed (UnsupportedTTGIR) instead of fabricating an empty token. Adds a unit test pinning the semantics.
… CLI entry point - CompiledRaceDetector is WARMUP_ONLY: TritonTrace.run skips the interpreter entirely (no language patching, no grid loop) and executes the REAL kernel after the warmup compile, so the host script keeps its true semantics (live outputs, asserts, autotuning). Load-bearing, not an optimization: the interpreter's in-place tl.core.tensor dunder patches leak past snapshot/restore and break a later real compile in the same process. - Real-compile windows run with TritonTrace module globals unwound to the underlying JITFunction (_unwrapped_jit_globals) so device-function callees wrapped by the CLI resolve in the code generator. - New triton-compiled-race-detector CLI entry point; prints a per-kernel verdict (race-free proof / RACE with reports / UNSUPPORTED reason). - Tests: trace-level tests run on CUDA and assert live outputs; add regressions for second-real-compile survival and for a kernel calling a wrapped device fn (exec'd into a synthetic module namespace to mirror the CLI's module-level globals — an in-function def would bind the callee as a closure freevar, which get_capture_scope() overlays on top of __globals__, outside the swap window's reach).
…-detector-compiled-mode # Conflicts: # triton_viz/wrapper.py
…-detector-compiled-mode
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a compiled-mode race detector: a static shared-memory race analysis over the TTGIR emitted by Triton's own compilation pipeline, selected via
RaceDetector(compile=True)or the newtriton-compiled-race-detectorCLI.post_warmup_callbackreceives the runtime's ownCompiledKernel) — never a hand-built ASTSource, so the analyzed IR is exactly the launched specialization, including the pipeliner stages. Analysis results are cached per TTGIR SHA-256.unsupported— never a silent ok.STANDALONE(ClientManager rejects composing it with other clients) andWARMUP_ONLY—TritonTrace.runskips the interpreter entirely and executes the REAL kernel, so host scripts keep their true semantics (live outputs, asserts, autotuning). This is load-bearing: the interpreter's in-placetl.core.tensordunder patches leak past snapshot/restore and break a later real compile in the same process.@triton.jitfunction — including device fns — into aTritonTrace, which the real code generator rejects as a callee. Real-compile windows run with trace module globals unwound to the underlyingJITFunction(_unwrapped_jit_globals) and restored afterwards.triton-compiled-race-detector <script.py>prints a per-kernel verdict — race-free (proof) / RACE with rendered reports / UNSUPPORTED with the reason.Stacked on #361 (
race-detector-z3-demo); base set accordingly. Supersedes #421 (closed as stale).Testing
tests/end_to_end/test_compiled_race_detector.py— 22 tests: golden-TTGIR proofs, pipeliner-bug mutations that must produce witnessed RAW reports, unsupported-degradation cases, and CUDA trace-level tests including two regressions: a second REAL compile after a traced launch, and a kernel calling a CLI-wrapped device function (exec'd into a synthetic module namespace to mirror the CLI's module-level globals; an in-function def would bind the callee as a closure freevar, whichget_capture_scope()overlays on top of__globals__, outside the swap window's reach).