diff --git a/InternalDocs/README.md b/InternalDocs/README.md index 3e8ab4423157534..ec297e1941b7d88 100644 --- a/InternalDocs/README.md +++ b/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) diff --git a/InternalDocs/jit_performance.md b/InternalDocs/jit_performance.md new file mode 100644 index 000000000000000..0ec95ee813402be --- /dev/null +++ b/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. diff --git a/Python/jit.c b/Python/jit.c index 50e606a7a0963f6..378bcc4ce129115 100644 --- a/Python/jit.c +++ b/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; diff --git a/Tools/jit/README.md b/Tools/jit/README.md index 2a687bb9e899e77..ca67d5a9f9737be 100644 --- a/Tools/jit/README.md +++ b/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. \ No newline at end of file +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. diff --git a/Tools/jit/benchmark.py b/Tools/jit/benchmark.py new file mode 100644 index 000000000000000..926aa8d9eb6c851 --- /dev/null +++ b/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()