ttsem 0.1.0__py3-none-any.whl

ttsem/__init__.py

@@ -0,0 +1,137 @@
+"""ttsem: an executable semantics of Triton's TTIR and TTGIR, and a per-pass validator.
+
+The package has three layers, each usable on its own:
+
+* **Level 1, values.** :func:`ttsem.mlir.parse` reads a generic-form MLIR module and
+  :class:`ttsem.interp.Interp` executes it over ``numpy`` arrays. No GPU, no Triton install,
+  no ``triton-opt``: the module text is all that is needed.
+* **Level 2, layouts.** :class:`ttsem.interp2.LayoutInterp` and
+  :func:`ttsem.interp2.check_module` add the linear-layout model, so a TTGIR module can be
+  checked for layout violations without running it.
+* **Level 3, concurrency.** :func:`ttsem.races.detect` replays one execution and reports
+  shared-memory races between the warps of a CTA.
+
+On top of those, :func:`ttsem.validate.validate_stages` takes the modules a compilation
+produced before every pass and names the first pass whose output the semantics reads
+differently from the device, and :func:`ttsem.minimize.reduce` shrinks that module to the
+smallest one on which the pass still changes the meaning.
+
+Everything above level 3 needs a Triton install, because a launch has to be recorded and
+compiled to get the stages in the first place; those names are imported lazily, so importing
+``ttsem`` with nothing but ``numpy`` installed works and gives you the interpreter.
+
+The whole-program entry point keeps its module name rather than being re-exported here, so
+that ``from ttsem import validate`` is unambiguously the submodule: call
+``ttsem.validate.validate(program, record)`` for it.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from ttsem.interp import Interp
+from ttsem.interp2 import LayoutInterp, check_module, launch_config, module_attrs
+from ttsem.ir_types import Block, Module, Op, Region, Type
+from ttsem.memory import Memory, MemoryFault
+from ttsem.mlir import NotGeneric, ParseError, parse, to_generic, unparse
+from ttsem.races import detect
+from ttsem.values import Poison, Unsupported, Value
+
+__version__ = "0.1.0"
+
+# The names below pull in `triton`, so they are resolved on first use rather than at import.
+_LAZY = {
+    "DumpTooLarge": ("ttsem.harness", "DumpTooLarge"),
+    "LaunchRecord": ("ttsem.harness", "LaunchRecord"),
+    "capture_launches": ("ttsem.harness", "capture_launches"),
+    "compare": ("ttsem.harness", "compare"),
+    "dump_for_launch": ("ttsem.harness", "dump_for_launch"),
+    "execute": ("ttsem.harness", "execute"),
+    "ir_for_launch": ("ttsem.harness", "ir_for_launch"),
+    "run_launch": ("ttsem.harness", "run_launch"),
+    "PassResult": ("ttsem.validate", "PassResult"),
+    "Report": ("ttsem.validate", "Report"),
+    "capture_dump": ("ttsem.validate", "capture_dump"),
+    "culprit_of": ("ttsem.validate", "culprit_of"),
+    "split_dump": ("ttsem.validate", "split_dump"),
+    "validate_stages": ("ttsem.validate", "validate_stages"),
+    "reduce": ("ttsem.minimize", "reduce"),
+}
+
+if TYPE_CHECKING:  # so that type checkers and editors see the same API
+    from ttsem.harness import (
+        DumpTooLarge,
+        LaunchRecord,
+        capture_launches,
+        compare,
+        dump_for_launch,
+        execute,
+        ir_for_launch,
+        run_launch,
+    )
+    from ttsem.minimize import reduce
+    from ttsem.validate import (
+        PassResult,
+        Report,
+        capture_dump,
+        culprit_of,
+        split_dump,
+        validate_stages,
+    )
+
+
+def __getattr__(name: str) -> Any:
+    try:
+        module_name, attr = _LAZY[name]
+    except KeyError:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
+    import importlib
+
+    value = getattr(importlib.import_module(module_name), attr)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(_LAZY))
+
+
+__all__ = [
+    "Block",
+    "DumpTooLarge",
+    "Interp",
+    "LaunchRecord",
+    "LayoutInterp",
+    "Memory",
+    "MemoryFault",
+    "Module",
+    "NotGeneric",
+    "Op",
+    "ParseError",
+    "PassResult",
+    "Poison",
+    "Region",
+    "Report",
+    "Type",
+    "Unsupported",
+    "Value",
+    "__version__",
+    "capture_dump",
+    "capture_launches",
+    "check_module",
+    "compare",
+    "culprit_of",
+    "detect",
+    "dump_for_launch",
+    "execute",
+    "ir_for_launch",
+    "launch_config",
+    "module_attrs",
+    "parse",
+    "reduce",
+    "run_launch",
+    "split_dump",
+    "to_generic",
+    "unparse",
+    "validate_stages",
+]

ttsem/__main__.py

@@ -0,0 +1,92 @@
+"""``python -m ttsem``: what the package can do, and which module does it."""
+
+from __future__ import annotations
+
+import sys
+
+from ttsem import __version__
+
+USAGE = f"""ttsem {__version__} - an executable semantics of Triton's TTIR and TTGIR.
+
+The package is a library first; every entry point below is a module you run with `-m`.
+After a pip install, the `ttsem` command is the same as `python -m ttsem`.
+
+  python -m ttsem demo [--live]
+      The whole idea on one real bug (triton#11519): replay the committed recording of a
+      per-pass validation, or, with --live and the Triton wheel, redo it here. No GPU.
+
+  python -m ttsem demo --torch
+      Races in the kernels torch.compile generates (pytorch#197829, #198033), found by
+      running Inductor's Triton code under the semantics. Needs torch; no GPU.
+
+  python -m ttsem judge TASK.py ANSWER.py [--json] [--rule kernelbench|v3]
+  python -m ttsem judge --manifest M.jsonl --out DIR [--jobs N]
+      A KernelBench answer written with Triton, judged by KernelBench's own correctness
+      rule with every launch executed by the semantics: verified, wrong, unsafe (with the
+      kernel's line), no_kernel (no Triton kernel launched), error, too_slow or not_judged.
+      Each answer runs in its own bwrap sandbox. Needs torch and Triton; no GPU.
+      With --tritonbench, a TritonBench answer against its task's reference, by value.
+
+  python -m ttsem.validate PROGRAM.py [--device cpu|cuda] [--json OUT]
+      Compile PROGRAM.py's first launch with the dump on and run every pass's input
+      through the semantics. Names the first pass whose output the semantics reads
+      differently from the device. Needs Triton.
+
+  python -m ttsem.minimize PROGRAM.py [--stage PASS] [--out small.mlir]
+      Shrink that module to the smallest one on which the pass still changes the meaning.
+
+  python -m ttsem.sanitize pytest [PYTEST ARGS] | SCRIPT.py
+      Run a test suite or a script written for a GPU on a machine without one: every
+      launch is executed by the semantics, and an out-of-bounds access, a race between
+      program instances or an operation with an undefined result fails the test with the
+      kernel's own line.
+
+  python -m ttsem.races MODULE.generic.mlir [--strip-barriers]
+      Shared-memory races between the warps of a CTA, from one execution of the module.
+      Needs no Triton and no GPU.
+
+  python -m ttsem.check_fixtures [--dir DIR] [--pattern GLOB]
+      Level-2 layout checks over a directory of generic-form TTGIR dumps. No Triton.
+
+  python -m ttsem.bot summary RUN/ | diff PREVIOUS/ CURRENT/
+      Bookkeeping over the JSONL a validated test run writes.
+
+  pytest -p ttsem.pytest_ttsem ...            validate every launch of a test run
+  python -m ttsem.inductor_trace --out DIR SCRIPT.py
+                                              the same for torch.compile's kernels
+  python -m ttsem.validate_corpus --programs DIR --results OUT
+  python -m ttsem.pass_stress --programs DIR  which passes a corpus actually wakes
+  python -m ttsem.dump_corpus --programs DIR --out DIR
+
+In Python:
+
+  from ttsem import parse, Interp, Memory        # level 1, numpy only
+  from ttsem import check_module                 # level 2, layouts
+  from ttsem import detect                       # level 3, races
+  from ttsem import validate_stages, culprit_of  # validate_stages needs Triton, culprit_of does not
+"""
+
+
+def main(argv: list[str] | None = None) -> int:
+    argv = list(sys.argv[1:] if argv is None else argv)
+    if argv and argv[0] == "demo":
+        from ttsem.demo import main as demo_main
+
+        return demo_main(argv[1:])
+    if argv and argv[0] == "judge":
+        from ttsem.judge import main as judge_main
+
+        return judge_main(argv[1:])
+    if argv and argv[0] in ("-h", "--help"):
+        sys.stdout.write(USAGE)
+        return 0
+    if argv:
+        sys.stderr.write(f"unknown command {argv[0]!r}\n\n")
+        sys.stderr.write(USAGE)
+        return 2
+    sys.stdout.write(USAGE)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

ttsem/_compile_runner.py

@@ -0,0 +1,65 @@
+"""Subprocess entry point for ``validate.py``'s ``MLIR_ENABLE_DUMP`` capture without a GPU.
+
+Installs the fake CUDA driver (:func:`ttsem.harness._install_fake_driver`) and
+replaces ``JITFunction.run`` so every launch is compiled for that fake target -- which is what
+actually produces the ``MLIR_ENABLE_DUMP`` pass-by-pass text on stderr, the whole point of
+running this -- and never executed: there is no GPU here to launch a real kernel on.
+
+Not meant to be imported: run as ``python -m ttsem._compile_runner PROGRAM.py --cc CC``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+from typing import Any
+
+from triton.backends.compiler import GPUTarget
+from triton.runtime.jit import JITFunction
+
+from ttsem import harness
+
+
+def _install(target: GPUTarget) -> None:
+    def run(self: JITFunction, *args: Any, grid: Any, warmup: bool, **kwargs: Any) -> Any:
+        if not warmup:
+            try:
+                harness._compile_asm(self, args, kwargs, target)
+            except Exception as e:  # the dump on stderr is what matters, not this return value
+                print(f"[_compile_runner] compile failed: {e!r}", file=sys.stderr)
+        return None
+
+    JITFunction.run = run
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("program", type=Path)
+    ap.add_argument("--cc", type=int, default=harness.DEFAULT_CC)
+    ap.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
+    ap.add_argument(
+        "--generic",
+        action="store_true",
+        help="switch the wheel's printer to the generic form, so the dump needs no triton-opt",
+    )
+    args = ap.parse_args()
+
+    if args.generic:
+        from ttsem import mlir
+
+        try:
+            mlir.enable_generic_printing()
+        except RuntimeError as exc:  # a source build hides the symbol: the dump comes pretty
+            print(f"ttsem: {exc}", file=sys.stderr)
+    if args.device == "cuda":
+        harness._run_program(args.program, "cuda")
+        return 0
+    harness._install_fake_driver(args.cc)
+    _install(harness.target_for("cpu", args.cc))
+    harness._run_program(args.program, "cpu")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())