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

Author: markshannon

Tools/jit/README.md

@@ -86,3 +86,8 @@ If you're looking for information on how to update the JIT build dependencies, s
 [^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

@@ -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()