jit(wasm): general CALL_ASSEMBLER + float residual/local perf gaps (nbody/fannkuch/float_loop/spectral)#564
jit(wasm): general CALL_ASSEMBLER + float residual/local perf gaps (nbody/fannkuch/float_loop/spectral)#564youknowone wants to merge 5 commits into
Conversation
WalkthroughThe PR adds float-aware WASM local and residual-call code generation, restructures CALL_ASSEMBLER target selection and lifecycle handling, wires CA bridge and baseline resume paths, and exposes guest-side JIT execution statistics. ChangesWASM JIT execution changes
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant WASMTrace
participant CALL_ASSEMBLER
participant wasm_ca_resume_deopt
participant try_compile_ca_bridge
participant BaselineHelper
WASMTrace->>CALL_ASSEMBLER: enter compiled target
CALL_ASSEMBLER->>wasm_ca_resume_deopt: deopt with guard descriptor
wasm_ca_resume_deopt->>try_compile_ca_bridge: attempt CA bridge compilation
try_compile_ca_bridge-->>wasm_ca_resume_deopt: return terminal decline status
wasm_ca_resume_deopt->>BaselineHelper: resume baseline evaluation when needed
Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
🤖 Codex parity reviewStatic analysis of this diff vs the local RPython/PyPy sources (commit f02eda5). Files in the reviewed diff1. Regressions to PyPy parity introduced by this patch
2. Other mismatches introduced by this patch
3. Pre-existing mismatches (already present before this patch)
4. Structural adaptations
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8e2a773af7
ℹ️ 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".
| value_types[ia.index as usize] = ValType::F64; | ||
| } | ||
| } | ||
| for op in ops { |
There was a problem hiding this comment.
When a LABEL introduces an independent float OpRef that is not an inputarg or an op result, this census leaves that local as I64 because it only marks inputargs/results. Later the local backedge and resume loader choose conversions from value_types[label_arg.raw()], while float users call emit_resolve_f64 expecting the same local to be F64, so float loop-carried label args can generate invalid Wasm local types. The Cranelift backend has an explicit LABEL-arg declaration path for this case (majit-backend-cranelift/src/compiler.rs:9077-9099); this needs the same OpRef::ty()-based marking.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3290979f23
ℹ️ 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".
| if !matches!( | ||
| op.opcode, | ||
| majit_ir::OpCode::CallAssemblerI | majit_ir::OpCode::CallAssemblerR | ||
| ) { |
There was a problem hiding this comment.
When a general CALL_ASSEMBLER target is later marked terminally declined, the generated terminal branch calls wasm_ca_baseline_call(op.arg(0), compiled_ptr), whose implementation casts that first word to *mut PyFrame and returns eval_with_jit's PyObjectRef bits. This admission path now accepts CallAssemblerI and any all-Int/Ref descriptor inputs, so a red-only integer-result or non-frame-first CALL_ASSEMBLER will either dereference the wrong value as a PyFrame or store a boxed object pointer into an integer result once the terminal-decline bit flips; the direct path needs to be limited to the ABI that the baseline helper actually implements, or the fallback needs per-descr marshalling/result handling.
Useful? React with 👍 / 👎.
| Err(err) => { | ||
| set_pending_ca_exception(err); | ||
| 0 |
There was a problem hiding this comment.
For a terminal-declined CA callee that raises, this only stores the error in LAST_CA_EXCEPTION and returns 0, but the compiled caller continues through the recorded GuardNoException/GuardException checks, which read the backend jit_exc_type/jit_exc_value slots rather than this TLS stash. That makes the caller miss the exception and consume 0 as the call result; the helper should set the backend JIT exception state before returning to compiled code.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@majit/majit-backend-wasm/src/failguard.rs`:
- Around line 236-239: Update the Drop implementation for CompiledWasmLoop to
invalidate every registered self.ca_callers entry before the loop’s storage can
be freed, clearing the corresponding ca.terminal_declined_ptr and
ca.callee_compiled_ptr raw pointers; preserve the existing
CALL_ASSEMBLER_TARGETS removal while ensuring active callers cannot retain
dangling references.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: ad648ff3-7069-49d0-abd4-d83c7a9b0516
📒 Files selected for processing (9)
majit/majit-backend-wasm/src/codegen.rsmajit/majit-backend-wasm/src/failguard.rsmajit/majit-backend-wasm/src/glue.rsmajit/majit-backend-wasm/src/lib.rsmajit/majit-metainterp/src/pyjitpl.rspyre/pyre-jit/src/call_jit.rspyre/pyre-jit/src/eval.rspyre/pyre-wasm-runner/src/main.rspyre/pyre-wasm/src/lib.rs
| pub struct CompiledWasmLoop { | ||
| /// Owning `JitCellToken` number, used to retract this loop's | ||
| /// CALL_ASSEMBLER target metadata on drop. | ||
| pub token_number: u64, |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm Drop never walks ca_callers, and check whether eviction elsewhere
# accounts for active CA callers before removing a compiled loop.
rg -n 'ca_callers' majit/majit-backend-wasm/src/*.rs
rg -n 'fn drop' -A 40 majit/majit-backend-wasm/src/failguard.rs | rg -n 'ca_callers|CompiledWasmLoop'
rg -n 'fn try_to_free_some_loops' -A 60 majit/majit-metainterp/src/pyjitpl.rs | rg -n 'ca_callers|CompiledWasmLoop|wasm'Repository: youknowone/pyre
Length of output: 805
Dropping a CompiledWasmLoop can leave active CA callers with dangling pointers.
Drop removes this loop from CALL_ASSEMBLER_TARGETS, but it never invalidates self.ca_callers. That means ordinary eviction can free a callee while other compiled wasm code still has raw pointers into this struct (ca.terminal_declined_ptr / ca.callee_compiled_ptr). Either walk the callers here and invalidate them, or keep the callee alive while callers remain registered.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@majit/majit-backend-wasm/src/failguard.rs` around lines 236 - 239, Update the
Drop implementation for CompiledWasmLoop to invalidate every registered
self.ca_callers entry before the loop’s storage can be freed, clearing the
corresponding ca.terminal_declined_ptr and ca.callee_compiled_ptr raw pointers;
preserve the existing CALL_ASSEMBLER_TARGETS removal while ensuring active
callers cannot retain dangling references.
Route the float residual families (CallF, CallPureF, CallLoopinvariantF, CallMayForceF) with a Float result through a direct in-module call_indirect instead of the env.jit_call host trampoline. residual_call_float_sig derives the callee's real wasm signature from the call descr (Int/Ref -> i64, Float -> f64, result f64); the distinct signatures are declared as function types after the existing i64 and CA type families so no prior type index shifts, and Float SSA values are reinterpreted to/from their i64 bit carrier around the call. Assembler and release-GIL float calls stay on the trampoline. On nbody, jit_calls fall from 5.28M to 292K with unchanged energy output. Assisted-by: Claude
Type each value local by its SSA value's type (f64 for Float, i64 for Int/Ref) instead of homing every value in a uniform i64 local. Float producers store the f64 result directly and float consumers read it directly; a reinterpret is inserted only where a float value crosses into the i64 world (frame stores, guards, residual i64-ABI args, loop seed loads). Int/Ref-only traces coalesce to the prior single i64 local run so their emitted module is byte-identical. Removes the per-float-op i64<->f64 reinterprets that Cranelift lowered to gp<->fp fmov moves. float_loop wasm-vs-dynasm work-ratio drops from ~4.5x to ~0.9x, spectral_norm from ~1.8x to ~1.3x; nbody float arithmetic is likewise unboxed. Assisted-by: Claude
…atures The typed float call_indirect derived its wasm signature from the call descr, mapping Int/Ref args to i64. That is unsound for a float-result residual target whose pointer argument is i32 on wasm32 (e.g. jit_bigint_to_f64_or_inf(&BigInt) -> f64): the emitted (i64)->f64 type mismatches the (i32)->f64 table entry and call_indirect traps. Unlike the audited int residual family, float-result targets are not guaranteed to use a uniform word ABI, so the descr's Int/Ref types are not a reliable oracle for the callee's wasm parameter widths. Accept a float residual for the direct path only when every argument is Float (an unambiguous f64); Int/Ref/other args keep the reflecting jit_call trampoline. Float ** (pow), the dominant win, is [Float,Float]->Float and stays direct. Assisted-by: Claude
…f the host jit_execute round-trip `glue::execute` (host-import path) transmutes the trace's shared-table slot (`func_id`) to an `extern "C" fn(u32) -> u32` and calls it, running the trace inside the guest instead of crossing to the host `jit_execute_wasm` import. The slot is the same index the host dispatched by (`table.get(func_id)`) and the CA self-recursion arm already re-enters through (`call_indirect(0, self_slot)`). The host import is retained on a `black_box(false)` branch so `wasm-ld` keeps it in the import table (dropping it would shift import indices and break baked JIT indices). A guest `JIT_EXECUTE_COUNT` atomic, exported as `pyre_jit_execute_count` and read by the runner, keeps the `executes` stat now that the host no longer observes each entry. Assisted-by: Claude
… terminal-decline back-off The wasm backend previously declined every trace containing a CALL_ASSEMBLER except the self-recursive-int fib shape (bridge_is_self_recursive_int_ca), so an outer loop that calls an inner-loop-bearing callee never compiled and round-tripped through the interpreter each iteration. - general_int_call_assembler_target admits a loop/bridge whose CALL_ASSEMBLER targets any compiled loop when every arg type and the result are Int|Ref. - emit_ca marshals the resolved callee geometry (frame bytes, gcmap, shared table slot) and call_indirect(0, callee_slot)s it, reading the result or blackhole-resuming the callee on a non-finish exit. - try_compile_ca_bridge feeds a CA callee guard through the normal must-compile/bridge path; when the callee bridge terminally declines on wasm, mark_call_assembler_terminal_decline invalidates the callers so the next trace refuses the target and returns to the baseline path (wasm_ca_baseline_call), avoiding a deopt-to-blackhole-per-call regression. - The frozen chain frame geometry reserves 64 value slots / 64 Ref homes so a later CA-callee bridge fits the loop's frame layout. nbody 2.44s -> 0.46s (timeout gone); ca_int/ca_innerloop 14.8s -> ~0.8s; byte-exact vs dynasm within the known cross-ISA pow ULP tolerance. Assisted-by: Claude
youknowone
changed the title
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6464e1f2eb
ℹ️ 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".
| if is_float { | ||
| emit_resolve_f64(&mut sink, constants, value_types, arg0); | ||
| emit_resolve_f64(&mut sink, constants, value_types, op.arg(1).to_opref()); | ||
| sink.f64_ne(); |
There was a problem hiding this comment.
When GUARD_VALUE specializes a Float, this branch uses numeric f64_ne; for a promoted NaN the value is not equal to itself, so the trace guard will fail every time even when the runtime bits match the recorded value. Guard-value hotness/status records float raw bits and the frame/failarg ABI stores floats as i64 bit carriers, so the guard should compare the reinterpreted bits rather than f64 equality.
Useful? React with 👍 / 👎.
1 participant
Closes the largest wasm-JIT-vs-native(aarch64 dynasm) steady-state perf gaps: float-heavy code (residual crossings + f64 homing) and compiled→compiled loop entry (general
CALL_ASSEMBLER). Every lever is the direct wasm analogue of what the native backends already do.Commits
lower float residual calls through typed direct call_indirect— float residual calls (CallF/CallPureF/CallLoopinvariantF/CallMayForceF, e.g.x ** 0.5→ libmpow) went through theenv.jit_callhost trampoline (guest→host→guest signature reflection + arg marshalling). They now use a direct in-modulecall_indirectwith a per-descr(f64…)->f64wasm signature declared after the existing i64/CA type families. Onnbodythis drops host crossings from 5.28M → 292K.home Float-typed SSA values in f64 wasm locals— every SSA value was homed in a uniformi64local, so each float op was bracketed byi64<->f64reinterprets that Cranelift lowers to realfmov gp<->fpmoves. Float-typed values now usef64locals; a reinterpret is inserted only where a float crosses into the i64 world (frame stores, guards, residual i64-ABI args, loop-seed loads). Int/Ref-only traces coalesce to the prior single i64 local run, so their emitted module is byte-identical.restrict direct float residual calls to all-float-arg signatures— review follow-up. A float-result residual target whose pointer argument isi32on wasm32 (e.g.jit_bigint_to_f64_or_inf(&BigInt) -> f64) would trap on a(i64)->f64call_indirecttype mismatch. The direct path now accepts a float residual only when every argument isFloat(unambiguousf64);Int/Refargs keep the reflecting trampoline.pow([Float,Float]->Float) — the dominant win — stays direct.enter compiled traces via in-guest call_indirect instead of the host jit_execute round-trip— trace entry (glue::execute) now transmutes the trace's shared-table slot toextern "C" fn(u32)->u32and runs the trace inside the guest instead of crossing to thejit_execute_wasmhost import (the same slot the host dispatched by, and the one the CA self-recursion arm already re-enters through). The host import is retained on a deadblack_box(false)branch sowasm-ldkeeps it in the import table — dropping it would shift import indices and break baked JIT indices. A guestJIT_EXECUTE_COUNTatomic, exported aspyre_jit_execute_countand read by the runner, preserves theexecutesstat.lower general CALL_ASSEMBLER to any compiled-loop callee with a terminal-decline back-off— the wasm backend previously declined every trace containing aCALL_ASSEMBLERexcept the self-recursive-int fib shape (bridge_is_self_recursive_int_ca), so an outer loop that calls an inner-loop-bearing callee never compiled and round-tripped through the interpreter each iteration.general_int_call_assembler_targetnow admits a loop/bridge whoseCALL_ASSEMBLERtargets any compiled loop when every arg type and the result areInt|Ref.emit_camarshals the resolved callee geometry (frame bytes, gcmap, shared-table slot) andcall_indirect(0, callee_slot)s it, reading the result or blackhole-resuming the callee on a non-finish exit.try_compile_ca_bridgefeeds the CA callee guard through the normal must-compile/bridge path; when the callee bridge terminally declines on wasm,mark_call_assembler_terminal_declineinvalidates the callers so the next trace refuses the target and returns to the baseline path (wasm_ca_baseline_call), avoiding a deopt-to-blackhole-per-call regression.nbody 2.44s → 0.46s (timeout gone);
ca_int/ca_innerloop14.8s → ~0.8s; byte-exact vs dynasm within the known cross-ISA pow ULP tolerance.Measured (wasm-vs-dynasm work-ratio, startup-corrected)
nbody: the float levers (host crossings + fmov) took 17.6x → 12.5x; generalCALL_ASSEMBLERthen took 12.5x → 2.4x (the module-global dict-call outer loop now compiles instead of round-tripping the interpreter every iteration). fannkuch's remaining 3.16x is an irreducible wasm array-access instruction-density constant factor (wasm has no scaled-index addressing —emit_array_addr≈ 9 wasm instrs vs one aarch64ldr [base,idx,lsl#3,#off]), verified not a loop-peeling/optimizer gap: the peel-abort → retry-without-unroll is PyPy-orthodox (virtualstate.py:372-375pins constant jump-values,intutils.py:1336-42widen_updatedoesn't loosen small bounds,pyjitpl.py:3046-50has the same non-unrolled fallback).Verification
pyre/check.py --backend dynasm,wasmgreen except the pre-existingsynth/list_insert_pop_indexBASEFAIL (both backends): dynasm 187/1, wasm 186/1.float_loop/spectral_norm/ca_*byte-exact dynasm-vs-wasm;nbodybyte-exact within the known cross-ISApowULP tolerance.majit-backend-wasmunit tests pass;int_loop/fib_loop(non-float) unchanged.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Diagnostics