wasm: pass a failing guard's fail args to its bridge as call parameters - #1274
Conversation
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (13)
WalkthroughThe PR refactors wasm module generation around retained inputs, adds parameterized and inlined bridge dispatch, supports loop re-emission and live trace replacement, exposes trace-entry and inline diagnostics, and changes bridge target lookup to use invalidation-filtered procedure tokens. ChangesWasm JIT bridge and trace updates
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
Possibly related PRs
Poem
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
…local A guard arm wrote its fail index into frame[0] and the shared epilogue read it back, subtracted the trace's fail-index base, scaled it by four and added the cell array base to reach the guard's bridge slot. The fail index is a constant at the arm, so the whole cell address is one too: the arm now computes it at emit time and leaves it in `bridge_slot_local`, and the epilogue loads the slot from there. `frame[0]` is still written, for the host round-trip that reads it. Finish and GuardAlwaysFails reach the same epilogue, so they set the local too; GuardAlwaysFails previously left an index there, which nothing read. Executed wasm ops, measured with `wasm_ops`: global_quasiimmut_invalidation -2.88%, the nested-loop microbenchmark -1.56%, fannkuch -0.63%, short_circuit_value_kept_stack -0.52%, raise_catch_loop -0.51%, fib_recursive -0.29%. A flat loop, whose steady state takes no guard exit, moves +0.00%. `compile_ms` on fannkuch reads 62.3 against 91.4. An earlier form of this inlined the whole dispatch, tail call included, into every guard arm. It cut more executed ops (gqi -6.71%) and lost to its own compile time: fannkuch's `compile_ms` went 75.2 to 181.9, against a +9.1% wall clock regression on a 1.6s fixture. pyre/check.py --backend wasm: 428/428. Assisted-by: Claude
…ridge into it `CompiledWasmLoop` retains the post-intern `ModuleBuildInputs` its module was built from, so `WasmBackend::reemit_loop` can rebuild it. The rebuild takes a fresh global fail-index base and a fresh guard-cell array, replays the bridge slots recorded for still-standalone bridges, re-bases the recorded descr ranges, and installs the result through a new `jit_replace_wasm` host binding that `table.set`s the loop's ORIGINAL slot — the slot is a baked `i32.const` immediate inside entry bridges and sibling loops' bridges, so a new slot would strand them. `intern_ref_constants` is never re-run: a second pass over already-rewritten ops finds no `ConstPtr`, yielding a zero GC-table base that reads linear memory at 0. The snapshot is kept only when a re-emission is armed, since it is live for the token's whole lifetime. The guard-cell array moves out of `build_wasm_module` to its callers, and `bridge_cells_base` / `num_guard_cells` become `Cell`s, because a rebuild's array has a different address and `JitCellToken.compiled` is a write-once `OnceLock`. `alloc_bridge_cells(0)` returns no array, keeping dispatch omitted for a guardless trace rather than handing out a zero-length allocation's non-null address. `ModuleBuildInputs` also carries inlined bridge regions. A region is emitted inside the loop as a `block` opened after `loop` and closed before the region's ops, so the failing guard reaches it with a `br` and its terminal JUMP takes the existing local-LABEL lowering — a parallel move between wasm locals. The guard arm moves its fail args straight into the region's inputarg locals instead of spilling them to the frame. A trial build decides eligibility, so `build_wasm_module` declines rather than asserting on a shape it cannot emit, such as a trace with no LABEL to branch back to. Both switches are read on the host and armed through guest exports; the guest has no environment, so `std::env::var` inside the backend never fires. Bridge diagnostics gain the re-emit and inline outcome counters. Two synthetic fixtures record one more `guard_failures` than their baseline with both switches off, and the cause is not yet identified. Sweeping `FROZEN_CHAIN_VALUE_SLOTS` over 64/96/128/192 moves which fixture is off by one without ever putting all three at their baseline, so the floor stays at 64. Assisted-by: Claude
A bridge module's entry was `(i32) -> i32` and read its inputargs back out of positional frame slots that the failing guard had just written. Declare one entry type per arity, `(i32, i64 x n) -> i32`, bind param `k+1` to `inputargs[k]`, and carry f64 values as their raw bits so one type serves each arity. Ref-typed params are stored into the callee's own Ref homes at entry, since wasm locals are not scanned. A bridge's inputarg list is its guard's fail-arg list, so the arity and the wasm type index are constants where the guard is emitted. Each armed guard arm therefore loads its own bridge cell and, when it is nonzero, resolves its fail args straight onto the operand stack and tail-calls with its own constant type index; a zero cell keeps the spill and the return to the host. The shared epilogue no longer carries a bridge tail call. `PYRE_WASM_BRIDGE_PARAMS` now only turns this off (`0`, `false`, `off`), read on the host and applied through a guest export. Measured with `PYRE_WASM_JIT_STATS`, off vs on, executed ops: global_quasiimmut_invalidation 6,100,873,004 -> 5,573,151,872 short_circuit_value_kept_stack 10,425,576,281 -> 9,972,315,185 fannkuch 30,139,493,121 -> 27,755,382,974 fib_recursive 38,583,936,147 -> 38,140,657,358 `compile_ms` rises on all four (6.0 -> 8.6, 20.5 -> 27.6, 57.4 -> 77.6, 38.1 -> 53.2): each armed arm grows by `6 + k` instructions, and every arm carries the cold spill path whether or not a bridge ever attaches. `pyre/check.py --backend wasm --synthetic-only` is 415/415 both with the switch defaulted on and with `PYRE_WASM_BRIDGE_PARAMS=0`. Assisted-by: Claude
…meters `compile_bridge` publishes an accepted bridge's own LABELs keyed on that bridge's table slot. A later trace whose cross-module terminal JUMP resolves to one of those targets emits `return_call_indirect(0, 0)` — type 0, which is `(i32) -> i32`. Since a bridge entry can declare `(i32, i64 x n) -> i32`, the two can disagree, and an indirect call whose declared type does not match the callee's traps at run time. Suppress publication when `bridge_entry_arity` is `Some(n)` with `n > 0`. `Some(0)` keeps publishing: `call_indirect` compares signatures structurally, and a zero-arity parameter entry is `(i32) -> i32` even though codegen mints it a separate type index. The emitted-module test asserts that equality. `bridge_param_label_suppressed` counts the suppression. Over the four graded fixtures plus `const_arg_call_resume` and `retrace_accumulator_type_flip` it fires once, on the last of those. `pyre/check.py --backend wasm --synthetic-only` is 415/415. Assisted-by: Claude
`executes` counts host entries only, so it is blind to the in-guest `return_call_indirect` traffic between a loop and its bridges. Nothing said how often a trace module is entered, or through which `br_table` key, which is what an executed-op budget has to be divided by. Add a per-trace, per-key counter in guest memory, armed from the host: `PYRE_WASM_TRACE_ENTRY_CENSUS` calls `pyre_jit_trace_entry_census_enable`, and `pyre_jit_trace_entry_census` reads the table back. The counter instructions, their locals and the census globals are emitted only when the census is armed at compile time; the `local.tee`/`local.get` pair that keeps the dispatch key for them is inside the same condition, because otherwise it costs two fuel on every entry into a peeled module. Disarmed `global_quasiimmut_invalidation` reads 5,575,086,653 against 5,575,060,318 for the same tree without the census. That fixture, N=10,193,192, reads per outer iteration in steady state: the loop entered 3x at key 2, one bridge 3x at key 0, another 1x at key 0. Its executed-op slope, differenced at N=200,000 and 400,000, is 497.0002, which those counts and the taken-path lengths account for exactly. `pyre/check.py --backend wasm --synthetic-only` is 415/415. Assisted-by: Claude
`gate_triage_complete` fails when a `PYRE_*` or `MAJIT_*` name is read from the environment with no entry in the matching triage file. `PYRE_WASM_BRIDGE_PARAMS` goes to §6a as a live default-ON brake; `PYRE_WASM_INLINE_BRIDGE` and `PYRE_WASM_REEMIT` get a §6a2 for default-OFF experiments; `PYRE_WASM_TRACE_ENTRY_CENSUS` joins the §6c diagnostics list and `MAJIT_TRACE_ENTRY_CENSUS` gets a majit entry. Assisted-by: Claude
`compile_trace_inner` took the JUMP descr from the `compiled_loops` side table while `has_compiled_targets`, consulted for the same decision, answered from the procedure token. pyjitpl.py:3005-3007 reads the token once and hands that same object to `compile_trace`, so both halves now come from `warm_state.get_procedure_token`, which applies the warmstate.py:191-196 invalidation filter. `first_target_token` is the token-side accessor for the head of `target_tokens`. Measured on `synth/global_quasiimmut_invalidation`: the per-trace, per-resume-key entry census is byte-identical before and after, and `check.py --backend wasm --synthetic-only` is 415/415. Assisted-by: Claude
🤖 Codex parity reviewStatic analysis of this diff vs the local RPython/PyPy sources (commit bdf4fe5). Files in the reviewed diff1. Regressions to PyPy parity introduced by this patchNone. 2. Other mismatches introduced by this patch
3. Pre-existing mismatches (already present before this patch)
4. Structural adaptations
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bdf4fe5dd1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| merged_inputargs.extend(bridge.inputargs.iter().map(InputArg::fresh_value_copy)); | ||
| for op in &bridge.ops { | ||
| if op.opcode == OpCode::LoadFromGcTable { | ||
| gc_table_bases.insert(op.pos.get().raw(), bridge.gc_table_base); |
There was a problem hiding this comment.
Remap bridge value IDs before merging traces
When PYRE_WASM_INLINE_BRIDGE is enabled, each bridge has been normalized independently, so its input and result IDs normally restart in the same range as the owner trace. Merging those operations unchanged makes ValueLocals alias unrelated values, and this insertion can also overwrite the GC-table base for an owner or earlier bridge's LoadFromGcTable, causing the merged trace to read the wrong rooted object. Remap every inlined region into a unique value-ID namespace before building the merged stream rather than keying this side table by trace-local IDs.
AGENTS.md reference: AGENTS.md:L126-L132
Useful? React with 👍 / 👎.
| pub _bridge_owned_cells: RefCell<Vec<Box<[u32]>>>, | ||
| /// Direct-loop guard index to bridge table slot. A re-emission replays | ||
| /// these slots into its fresh loop cell array. | ||
| pub bridge_slots: RefCell<std::collections::HashMap<u32, u32>>, |
There was a problem hiding this comment.
Use dense storage for per-guard bridge slots
When re-emission is enabled, this map stores values keyed by a bounded per-trace guard index whose extent is already tracked by num_guard_cells; a Vec<Option<u32>> or equivalent dense sequence expresses that shape directly. Introducing a HashMap here without demonstrating a dict-shaped upstream owner violates the repository's required RPython container-parity workflow and creates another parallel backend store that future ports must synchronize.
AGENTS.md reference: AGENTS.md:L115-L124
Useful? React with 👍 / 👎.
Seven commits on the wasm backend's trace-crossing path.
What changed
21d7cc5722f2faPYRE_WASM_REEMIT/PYRE_WASM_INLINE_BRIDGE, both default OFFd0cdd8bPYRE_WASM_BRIDGE_PARAMS, default ON0efae53d270418PYRE_WASM_TRACE_ENTRY_CENSUS/MAJIT_TRACE_ENTRY_CENSUS, default OFFf41ec9bbdf4fe5compiled_loopsside tableThe measurement
Guard→bridge parameter passing, one binary,
PYRE_WASM_BRIDGE_PARAMS=0against the default:synth/global_quasiimmut_invalidationsynth/short_circuit_value_kept_stackAt the landing measurement
fannkuchwas −7.91% andfib_recursive−1.15%.compile_msrises on every fixture. Each guard arm grows6 + kinstructions because thecold cell-is-zero host-resume path cannot be dropped, so every arm pays the cost while only
the guards that actually attach a bridge collect the win, and no arity threshold separates
them — win and cost both scale with
k.Why this shape
Upstream binds a bridge's inputargs to the locations the failing guard already froze and
patches the guard's jump straight into the bridge:
llsupport/assembler.py:201-225rebuild_faillocs_from_descr,x86/regalloc.py:212-217, 291-316prepare_bridge/_update_bindings,x86/assembler.py:690, 965-988. The jitframe round trip exists onlybefore a bridge is attached (
generate_quick_failure,x86/assembler.py:2001-2020). Sospilling every fail arg to the frame on each crossing is the deviation, and declaring
(i32, i64 × n) -> i32so a guard hands its fail args over as call arguments is the closestwasm-expressible analogue. Arities are small (2–16).
The first build of this put the
return_call_indirectin the one shared function epilogue,which cannot know which guard branched to it, so it recovered at runtime what each arm knew
statically: a tag bit, an arity local and a linear arity chain. Every fixture regressed
(
global_quasiimmut_invalidation+8.39%). Moving the tail call into each guard's own arm,where the type index is a constant, is the whole difference between the two results — a
guard's arity is a property of the guard, not of its bridge, so nothing had to be deferred to
runtime.
0efae53closes the trap the parameter entry opened:call_indirecttype-checksstructurally, so a bridge whose entry takes parameters must not publish LABEL targets that a
later closing jump would reach through a signature it does not have.
What did not land
+11.5% on
global_quasiimmut_invalidationand +4.5% onshort_circuit_value_kept_stack.Re-emitting into the same table slot makes every re-entry through that slot pay the merged
module's longer resume loader, and on this fixture that is ~3 re-entries per outer iteration
against one improved edge. Both halves stay behind default-OFF gates as the switched-off arm
of a one-binary comparison.
bdf4fe5puts the JUMP descr andhas_compiled_targetson the same object so thewarmstate.py:191-196invalidation filtercovers both. The per-key entry census is byte-identical before and after on
global_quasiimmut_invalidation, so it changes no measured behaviour there.Where the remaining budget is
The entry census attributes
global_quasiimmut_invalidation's steady state exactly: 497executed wasm ops per outer iteration, of which
3 × 71 = 213are entries into the loopmodule at resume key 2. The emitted wasm shows that key's resume loader is 52 ops — 4 frame
loads plus 10 GC-table constants rematerialised — and that the first thing in the loop body is
the invalidation-flag test, which after the store always exits to the attached bridge. Those
52 ops therefore feed a path that consumes two of them.
Instrumenting the metainterp and the backend's label registry says those entries land in the
invalidated module, and names why. After the store a fresh loop is compiled and the cell
adopts it (procedure token 1 → 3), so the close gate is answering about a live token — but that
token's target list is
[old preamble, old header, its own label], and virtual-state matchingresolves the post-store closes to the old header. The backend registry confirms the
executed edge: the first loop publishes two label targets (
fh=7813 k=1,fh=7813 k=2), thereplacement republishes only the preamble descr (
fh=7816 k=1, overprior_fh=Some(7813)—last-write-wins, as
failguard.rsdocuments), and all three post-store bridges bakefh=7813 k=2. The replacement module then collects a few hundred entries against ~30.5M intothe old one.
compile.py:290is the law this breaks: everycompile_loopdoesjitcell_token.target_tokens = [start_descr]— a fresh single-element list. A new token nevercarries a previous compile's targets, so upstream's invalidation filter
(
warmstate.py:191-196) covers the whole list, because the list lives on the token that wasinvalidated. Pyre mints a new token and seeds it with the prior loop's targets
(
seed_prior_target_tokens, pyjitpl.rs:6474 / :8171, feedingunroll_opt.target_tokens— notethis is not the
prior_front_target_tokensfallback at :7046/:8585, which a probe showsnever fires here), so the dead loop's labels travel past the filter on a live token.
That is a frontend fix of its own scope and is not attempted in this PR. The wasm entry cost
above is real and additive to it — upstream's
consider_jump(x86/regalloc.py:1303-1340) isa parallel move of the jump's args into the label's precomputed locations plus one
JMP, withconstants as immediates at their use sites and no entry loader of any kind — but it should be
measured after the target is corrected, since that changes which paths are hot.
Verification
python3 pyre/check.py --backend wasm --synthetic-only— 415/415cargo test --workspace— green, 138 test targetscargo build --release -p pyrex --bin pyre-dynasm --no-default-features --features dynasm— greenSummary by CodeRabbit
New Features
Bug Fixes
Documentation