GH-146073: Add example script for dumping JIT traces #148840

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

markshannon merged 2 commits into python:main from markshannon:executor-visualization

Apr 22, 2026

+196 −0

Tools/jit/README.md

-Original file line number
+Diff line change
@@ Expand Up @@
     [^pep-744]: [PEP 744](https://peps.python.org/pep-0744/)
     [^why-llvm]: Clang is specifically needed because it's the only C compiler with support for guaranteed tail calls (`musttail`), which are required by CPython's continuation-passing-style approach to JIT compilation. Since LLVM also includes other functionalities we need (namely, object file parsing and disassembly), it's convenient to only support one toolchain at this time.
+    ### Understanding JIT behavior
+    The [example_trace_dump.py](./example_trace_dump.py) script will (when configured as described in the script) dump out the
+    executors for a range of tiny programs to show the behavior of the JIT front-end.

Tools/jit/example_trace_dump.py

-Original file line number
+Diff line change
@@ -0,0 +1,191 @@
+    # This script is best run with pystats enabled to help visualize the shape of the traces.
+    # ./configure --enable-experimental-jit=interpreter  -C --with-pydebug --enable-pystats
+    # The resulting images can be visualize on linux as follows:
+    # $ cd folder_with_gv_files
+    # $ dot -Tsvg -Osvg *.gv
+    # $ firefox *.gv.svg
+    # type: ignore
+    import sys
+    import os.path
+    from types import FunctionType
+    # All functions declared in this module will be run to generate
+    # a .gv file of the executors, unless the name starts with an underscore.
+    def _gen(n):
+        for _ in range(n):
+            yield n
+    def gen_in_loop(n):
+        t = 0
+        for n in _gen(n):
+            t += n
+        return n
+    def short_loop(n):
+        t = 0
+        for _ in range(n):
+            t += 1
+            t += 1
+            t += 1
+            t += 1
+            t += 1
+        return t
+    exec(
+        "\n".join(
+            ["def mid_loop(n):"]
+            + ["    t = 0"]
+            + ["    for _ in range(n):"]
+            + ["        t += 1"] * 20
+            + ["    return t"]
+        ),
+        globals(),
+    )
+    exec(
+        "\n".join(
+            ["def long_loop(n):"]
+            + ["    t = 0"]
+            + ["    for _ in range(n):"]
+            + ["        t += 1"] * 100
+            + ["    return t"]
+        ),
+        globals(),
+    )
+    def _add(a, b):
+        return a + b
+    def short_loop_with_calls(n):
+        t = 0
+        for _ in range(n):
+            t = _add(t, 1)
+            t = _add(t, 1)
+            t = _add(t, 1)
+            t = _add(t, 1)
+            t = _add(t, 1)
+        return t
+    exec(
+        "\n".join(
+            ["def mid_loop_with_calls(n):"]
+            + ["    t = 0"]
+            + ["    for _ in range(n):"]
+            + ["        t = _add(t, 1)"] * 20
+            + ["    return t"]
+        ),
+        globals(),
+    )
+    exec(
+        "\n".join(
+            ["def long_loop_with_calls(n):"]
+            + ["    t = 0"]
+            + ["    for _ in range(n):"]
+            + ["        t = _add(t, 1)"] * 100
+            + ["    return t"]
+        ),
+        globals(),
+    )
+    def short_loop_with_side_exits(n):
+        t = 0
+        for i in range(n):
+            if t < 0:
+                break
+            t += 1
+            if t < 0:
+                break
+            t += 1
+            if t < 0:
+                break
+            t += 1
+            if t < 0:
+                break
+            t += 1
+            if t < 0:
+                break
+            t += 1
+        return t
+    exec(
+        "\n".join(
+            ["def mid_loop_with_side_exits(n):"]
+            + ["    t = 0"]
+            + ["    for _ in range(n):"]
+            + ["        if t < 0:", "            break", "        t += 1"] * 20
+            + ["    return t"]
+        ),
+        globals(),
+    )
+    exec(
+        "\n".join(
+            ["def long_loop_with_side_exits(n):"]
+            + ["    t = 0"]
+            + ["    for _ in range(n):"]
+            + ["        if t < 0:", "            break", "        t += 1"] * 100
+            + ["    return t"]
+        ),
+        globals(),
+    )
+    def short_branchy_loop(n):
+        # Branches are correlated and exit 1 time in 4.
+        t = 0
+        for i in range(n):
+            # Start with a few operations to form a viable trace
+            t += 1
+            t += 1
+            t += 1
+            if not t & 6:
+                continue
+            t += 1
+            if not t & 12:
+                continue
+            t += 1
+            if not t & 24:
+                continue
+            t += 1
+            if not t & 48:
+                continue
+            t += 1
+        return t
+    def _run_and_dump(func, n, outdir):
+        sys._clear_internal_caches()
+        func(n)
+        sys._dump_tracelets(os.path.join(outdir, f"{func.__name__}.gv"))
+    def _main():
+        if len(sys.argv) < 2 or len(sys.argv) > 3:
+            print(f"Usage: {sys.argv[0] if sys.argv else " "} OUTDIR [loops]")
+        outdir = sys.argv[1]
+        n = int(sys.argv[2]) if len(sys.argv) > 2 else 5000
+        functions = [
+            func
+            for func in globals().values()
+            if isinstance(func, FunctionType) and not func.__name__.startswith("_")
+        ]
+        for func in functions:
+            _run_and_dump(func, n, outdir)
+    if __name__ == "__main__":
+        _main()

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

GH-146073: Add example script for dumping JIT traces #148840

Diff view

Diff view

There are no files selected for viewing

Uh oh!