jit: precompute symbol offsets, add x86-64 GOT relaxation, and include lightweight JIT benchmark tool

InternalDocs/README.md

@@ -42,6 +42,7 @@ Program Execution
 - [Stack references (_PyStackRef)](stackrefs.md)
 
 - [The JIT](jit.md)
+- [JIT performance roadmap](jit_performance.md)
 
 - [Garbage Collector Design](garbage_collector.md)
 

InternalDocs/jit_performance.md

@@ -0,0 +1,132 @@
+# JIT performance roadmap
+
+This note summarizes what would be required to improve Python execution speed
+when the JIT is enabled. It is intentionally focused on end-to-end Python speed,
+not just JIT compilation latency or isolated generated-code peepholes.
+
+## Current pipeline and why local codegen fixes are not enough
+
+The current JIT is a tracing JIT layered on the adaptive interpreter. Execution
+starts in tier 1, hot `JUMP_BACKWARD`/`RESUME` instructions trigger trace
+recording, the recorder translates bytecodes into micro-ops, the uop optimizer
+runs, and the full JIT copy-and-patches stencil code for the optimized executor.
+See [The JIT](jit.md) for the detailed control flow.
+
+Small native-code improvements can help individual uops, but large Python-level
+speedups need changes that increase the amount of time spent in optimized,
+stable tier-2 traces and reduce the amount of Python object work that remains
+inside those traces. The main limiting factors are:
+
+* **Warmup and trace selection.** Default loop and resume thresholds are high
+  (`JUMP_BACKWARD_INITIAL_VALUE`, `RESUME_INITIAL_VALUE`) to avoid tracing before
+  specialization has stabilized. This protects steady-state quality, but it also
+  means short-running code sees little or no JIT benefit.
+* **Trace coverage.** The tracer records linear traces and side exits. A Python
+  workload with frequent side exits, polymorphism, exceptions, generator state,
+  or megamorphic attribute access spends less time in hot compiled traces.
+* **Optimizer scope.** The uop optimizer has symbolic analysis and removes some
+  unneeded uops, but it does not yet perform the kind of broad scalar
+  replacement, allocation sinking, loop-invariant hoisting, or aggressive
+  inlining needed for large speedups on object-heavy Python code.
+* **Runtime helper boundaries.** Many uops still call C helpers or perform
+  generic object operations. Native dispatch is faster than interpreter dispatch,
+  but it is not enough if the trace still performs the same refcount, allocation,
+  dict, descriptor, and call machinery as tier 1.
+* **Invalidation and deoptimization costs.** Executors depend on recorded values
+  and version checks. Broadening optimization requires equally strong guard,
+  dependency, and invalidation machinery so optimized traces stay correct without
+  deopting too often.
+
+## Changes needed for large end-to-end speedups
+
+### 1. Measure coverage before optimizing codegen
+
+A full `--enable-experimental-jit` build should be benchmarked with `pyperformance`
+and pystats enabled. For every benchmark, track at least:
+
+* percentage of executed bytecodes/uops covered by tier-2 executors;
+* traces created, traces executed, side exits, deopts, and invalidations;
+* average trace run length and top exit reasons;
+* JIT memory/code size and time spent compiling;
+* speed with `PYTHON_JIT=0` versus `PYTHON_JIT=1`.
+
+This should drive decisions about thresholds, trace length, side-exit linking,
+and optimizer work. A generated-code peephole should be considered secondary if
+coverage or deopt rate is the dominant loss.
+
+### 2. Improve trace formation and side-exit linking
+
+For many real programs, the important question is not whether a single trace is
+fast, but whether execution remains in tier 2 after common branches. The next
+major changes should be:
+
+* use pystats to identify bytecodes that stop tracing or force cold exits most
+  often;
+* raise or reshape `FITNESS_INITIAL`/trace-cost accounting only after measuring
+  code-size and compile-time tradeoffs;
+* compile common side exits into linked traces sooner when they are stable;
+* avoid recording traces on atypical loop iterations, such as exhaustion or
+  error paths;
+* preserve tier-1 specialization quality before tracing, because the recorder
+  relies on specialized bytecodes and cache state.
+
+### 3. Add optimizer passes that remove Python object work
+
+Large speedups require fewer Python object operations inside hot loops. The uop
+optimizer should grow passes such as:
+
+* compact-int and exact-float arithmetic chains that keep unboxed values across
+  multiple uops;
+* allocation sinking for temporary tuples/lists/ranges when escape analysis can
+  prove they do not escape;
+* loop-invariant guard and load hoisting for stable globals, builtins, types,
+  and descriptor lookups;
+* redundant reference-count operation elimination where ownership is proven;
+* stronger call inlining for monomorphic Python functions, with robust guards on
+  code object, defaults, closure cells, globals, and callable versioning.
+
+These are semantic optimizer changes. They should be developed with targeted
+trace dumps plus end-to-end benchmarks, not as isolated stencil rewrites.
+
+### 4. Specialize high-impact runtime operations in tier 2
+
+After coverage data identifies hot uops and helper calls, add JIT-friendly fast
+paths for operations that dominate Python workloads:
+
+* exact `int`, `float`, `str`, `list`, `tuple`, and `dict` operations;
+* monomorphic attribute loads and method calls;
+* `FOR_ITER` over `range`, `list`, `tuple`, and dict views;
+* common Python-to-Python calls, including vectorcall setup;
+* common C builtin calls where arguments are exact and stable.
+
+The goal is to keep hot traces in straight-line native code with guards, not to
+bounce through generic C APIs for every operation.
+
+### 5. Make generated code quality a measured follow-up
+
+Once coverage and semantic optimization are improved, generated-code work should
+focus on the remaining hottest stencils:
+
+* remove unnecessary GOT/data loads for patched operands;
+* keep hot paths fall-through and move deopt paths cold;
+* reduce tail-call chain overhead between adjacent uops where safe;
+* add architecture-specific immediate and branch relaxations only when they show
+  up in disassembly and benchmarks;
+* track code size so better local code does not reduce instruction-cache
+  behavior or trace residency.
+
+## Suggested milestone plan
+
+1. **Measurement milestone:** add a reproducible full-JIT `pyperformance` job
+   that emits coverage, deopt, side-exit, trace-length, and memory stats.
+2. **Coverage milestone:** fix the top trace abort and side-exit causes from
+   the measurement data.
+3. **Semantic optimizer milestone:** implement one high-value unboxing or
+   allocation-sinking optimization and prove it on multiple benchmarks.
+4. **Call/attribute milestone:** inline monomorphic Python calls and hoist stable
+   attribute/global guards in tier 2.
+5. **Codegen milestone:** apply architecture peepholes to the remaining hottest
+   stencils after the semantic wins are visible.
+
+The expected large wins come from milestones 2--4. Codegen-only work is still
+useful, but it is unlikely to produce a broad, substantial improvement by itself.

Python/jit.c

@@ -197,6 +197,7 @@ typedef uint32_t symbol_mask[SYMBOL_MASK_WORDS];
 typedef struct {
     unsigned char *mem;
     symbol_mask mask;
+    uint8_t offsets[SYMBOL_MASK_WORDS * 32];
     size_t size;
 } symbol_state;
 
@@ -276,17 +277,29 @@ get_symbol_slot(int ordinal, symbol_state *state, int size)
     const uint32_t state_mask = state->mask[ordinal / 32];
     assert(symbol_mask & state_mask);
 
-     // Count the number of set bits in the symbol mask lower than ordinal
-    size_t index = _Py_popcount32(state_mask & (symbol_mask - 1));
-    for (int i = 0; i < ordinal / 32; i++) {
-        index += _Py_popcount32(state->mask[i]);
-    }
-
+    size_t index = state->offsets[ordinal];
     unsigned char *slot = state->mem + index * size;
     assert((size_t)(index + 1) * size <= state->size);
     return slot;
 }
 
+static void
+initialize_symbol_state(symbol_state *state, size_t slot_size)
+{
+    size_t index = 0;
+    for (size_t word = 0; word < Py_ARRAY_LENGTH(state->mask); word++) {
+        uint32_t mask = state->mask[word];
+        while (mask) {
+            int bit = _Py_bit_length(mask & -mask) - 1;
+            size_t ordinal = word * 32 + bit;
+            assert(index < Py_ARRAY_LENGTH(state->offsets));
+            state->offsets[ordinal] = index++;
+            mask &= mask - 1;
+        }
+    }
+    state->size = index * slot_size;
+}
+
 // Return the address of the GOT slot for the requested symbol ordinal.
 static uintptr_t
 got_symbol_address(int ordinal, jit_state *state)
@@ -514,6 +527,39 @@ patch_x86_64_32rx(unsigned char *location, uint64_t value)
     memcpy(&relaxed, (void *)(value + 4), sizeof(relaxed));
     relaxed -= 4;
 
+    if (loc8[-2] == 0x8B &&
+        (loc8[-1] & 0xC7) == 0x05 &&
+        relaxed < (1ULL << 32))
+    {
+        // mov reg, [rip + AAA] -> mov reg, XXX
+        //
+        // This is especially helpful for JIT 32-bit operands on x86-64:
+        // many are not valid RIP-relative addresses, but still fit directly
+        // in a 32-bit immediate. Keep the original instruction length by using
+        // the C7 /0 encoding, whose immediate starts at the existing
+        // displacement location.
+        uint8_t modrm = loc8[-1];
+        uint8_t reg = (modrm >> 3) & 7;
+        if (loc8[-3] == 0x48 || loc8[-3] == 0x4C) {
+            // Preserve REX.W. If the old destination used REX.R, move that
+            // high register bit to REX.B because C7 encodes the destination in
+            // the r/m field.
+            if (relaxed < (1ULL << 31)) {
+                loc8[-3] = 0x48 | ((loc8[-3] & 0x04) >> 2);
+                loc8[-2] = 0xC7;
+                loc8[-1] = 0xC0 | reg;
+                patch_32(location, relaxed);
+                return;
+            }
+        }
+        else {
+            loc8[-2] = 0xC7;
+            loc8[-1] = 0xC0 | reg;
+            patch_32(location, relaxed);
+            return;
+        }
+    }
+
     if ((int64_t)relaxed - (int64_t)location >= -(1LL << 31) &&
         (int64_t)relaxed - (int64_t)location + 1 < (1LL << 31))
     {
@@ -662,19 +708,21 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
     data_size += group->data_size;
     combine_symbol_mask(group->trampoline_mask, state.trampolines.mask);
     combine_symbol_mask(group->got_mask, state.got_symbols.mask);
-    // Calculate the size of the trampolines required by the whole trace
-    for (size_t i = 0; i < Py_ARRAY_LENGTH(state.trampolines.mask); i++) {
-        state.trampolines.size += _Py_popcount32(state.trampolines.mask[i]) * TRAMPOLINE_SIZE;
-    }
-    for (size_t i = 0; i < Py_ARRAY_LENGTH(state.got_symbols.mask); i++) {
-        state.got_symbols.size += _Py_popcount32(state.got_symbols.mask[i]) * GOT_SLOT_SIZE;
-    }
-    // Round up to the nearest page:
+    // Calculate the size and ordinal-to-slot mapping for the symbols needed by
+    // the whole trace. Emit-time patching can then find each symbol slot with a
+    // direct indexed load instead of recounting all earlier mask bits.
+    initialize_symbol_state(&state.trampolines, TRAMPOLINE_SIZE);
+    initialize_symbol_state(&state.got_symbols, GOT_SLOT_SIZE);
+    // Round up to the nearest required alignment. Keep the padding at zero
+    // when the preceding region is already aligned; otherwise every aligned
+    // trace pays for an unnecessary extra alignment unit (or page).
     size_t page_size = get_page_size();
     assert((page_size & (page_size - 1)) == 0);
-    size_t code_padding = DATA_ALIGN - ((code_size + state.trampolines.size) & (DATA_ALIGN - 1));
-    size_t padding = page_size - ((code_size + state.trampolines.size + code_padding + data_size + state.got_symbols.size) & (page_size - 1));
-    size_t total_size = code_size + state.trampolines.size + code_padding + data_size + state.got_symbols.size + padding;
+    size_t code_and_trampolines_size = code_size + state.trampolines.size;
+    size_t code_padding = (0 - code_and_trampolines_size) & (DATA_ALIGN - 1);
+    size_t used_size = code_and_trampolines_size + code_padding + data_size + state.got_symbols.size;
+    size_t padding = (0 - used_size) & (page_size - 1);
+    size_t total_size = used_size + padding;
     unsigned char *memory = jit_alloc(total_size);
     if (memory == NULL) {
         return -1;

Tools/jit/README.md

@@ -90,4 +90,20 @@ If you're looking for information on how to update the JIT build dependencies, s
 ### 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.
+executors for a range of tiny programs to show the behavior of the JIT front-end.
+
+### Measuring JIT performance
+
+Use [benchmark.py](./benchmark.py) for a quick local smoke test before and after
+JIT changes. It runs a few hot-loop microbenchmarks twice, once with
+`PYTHON_JIT=0` and once with `PYTHON_JIT=1`, and prints the best timings and
+speedup ratios. This is not a replacement for `pyperformance`, but it is useful
+for checking whether a focused change affects tier-2 execution before running a
+larger benchmark suite.
+
+When configured with `--enable-experimental-jit=interpreter`, the benchmark
+measures the tier-2 micro-op interpreter described in
+[InternalDocs/jit.md](../../InternalDocs/jit.md). That mode is useful for
+understanding trace recording and optimizer behavior, but it can be slower than
+tier 1 because it does not exercise the copy-and-patch native-code backend. Use
+a full `--enable-experimental-jit` build when evaluating end-user JIT speedups.

Tools/jit/benchmark.py

@@ -0,0 +1,113 @@
+"""Small JIT-focused benchmark driver.
+
+This is intentionally lightweight: it runs a few hot-loop microbenchmarks under
+both ``PYTHON_JIT=0`` and ``PYTHON_JIT=1`` for a supplied Python executable.
+Use it after building with ``--enable-experimental-jit`` or
+``--enable-experimental-jit=interpreter`` to confirm that a proposed JIT change
+moves execution time in the expected direction before running larger suites.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+
+
+BENCHMARK = r'''
+import json
+import statistics
+import sys
+import time
+
+
+def loop_add(n):
+    total = 0
+    for i in range(n):
+        total += i
+    return total
+
+
+def loop_attr(n):
+    class C:
+        def __init__(self):
+            self.x = 1
+    obj = C()
+    total = 0
+    for _ in range(n):
+        total += obj.x
+    return total
+
+
+def loop_call(n):
+    def f(x):
+        return x + 1
+    total = 0
+    for i in range(n):
+        total += f(i)
+    return total
+
+
+def run(fn, n, warmups, samples):
+    for _ in range(warmups):
+        fn(10_000)
+    timings = []
+    for _ in range(samples):
+        start = time.perf_counter()
+        fn(n)
+        timings.append(time.perf_counter() - start)
+    return {
+        "best": min(timings),
+        "mean": statistics.mean(timings),
+        "stdev": statistics.pstdev(timings),
+    }
+
+
+jit = getattr(sys, "_jit", None)
+print(json.dumps({
+    "jit_available": bool(jit and jit.is_available()),
+    "jit_enabled": bool(jit and jit.is_enabled()),
+    "benchmarks": {
+        "loop_add": run(loop_add, 3_000_000, 5, 7),
+        "loop_attr": run(loop_attr, 3_000_000, 5, 7),
+        "loop_call": run(loop_call, 2_000_000, 5, 7),
+    },
+}, sort_keys=True))
+'''
+
+
+def run_once(python: str, jit: bool) -> dict[str, object]:
+    env = os.environ.copy()
+    env["PYTHON_JIT"] = "1" if jit else "0"
+    proc = subprocess.run(
+        [python, "-c", BENCHMARK],
+        check=True,
+        env=env,
+        stdout=subprocess.PIPE,
+        text=True,
+    )
+    return json.loads(proc.stdout)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Compare hot-loop timings with PYTHON_JIT disabled/enabled."
+    )
+    parser.add_argument("python", nargs="?", default=sys.executable)
+    args = parser.parse_args()
+
+    disabled = run_once(args.python, False)
+    enabled = run_once(args.python, True)
+    print(f"Python: {args.python}")
+    print(f"JIT available: {enabled['jit_available']}")
+    print("benchmark          jit-off best   jit-on best    speedup")
+    for name, off in disabled["benchmarks"].items():
+        on = enabled["benchmarks"][name]
+        speedup = off["best"] / on["best"]
+        print(f"{name:16} {off['best']:11.6f} {on['best']:11.6f} {speedup:9.3f}x")
+
+
+if __name__ == "__main__":
+    main()