[wasm][RyuJIT] Restore __stack_pointer in reverse P/Invoke epilogs

[lewing]

Problem

R2R Wasm code keeps the shadow SP in a local and leaves the __stack_pointer global stale. JIT_PInvokeBegin and the SuppressGCTransition publish (#130924) set __stack_pointer to the caller's shadow SP and leave it there. So when a reverse P/Invoke (UnmanagedCallersOnly) method's body, or any R2R code it calls, did an inlined P/Invoke, the method returned to its native caller with the global still lowered. That breaks the native ABI, which requires __stack_pointer to be restored on return.

In the observed case the P/Invoke was CoreLib's CastHelpers.ChkCastAny_NoCacheLookup under GetUserData<T>. When the UCO was reached through an R2R inlined P/Invoke (a delegate* unmanaged calli), the Debug check in JIT_PInvokeEndImpl fired: sp == stack_pointer_global_value at src/coreclr/vm/wasm/helpers.cpp:420. After that came GetFrame asserts and an out-of-bounds access. See #134681 for the full analysis.

The bug is reachable on main whenever an R2R method calls a UCO that itself does an inlined P/Invoke (see the test below). #134606 exposes it more broadly by making TestEntryPoint in readytorun/wasm/WasmInterpreterTransitions run as R2R instead of interpreted, which also routes the existing StreamLengthProxy case through R2R.

Fix

In CodeGen::genFnEpilog (codegenwasm.cpp), reverse P/Invoke methods now emit this sequence before return/end:

local.get <FP local, or SP local if there is no frame pointer>
i32.const genTotalFrameSize()
i32.add
global.set __stack_pointer

This restores the global to its entry value, the post-prolog SP plus the frame size. The sequence leaves the Wasm operand stack unchanged, so an already-pushed return value is preserved. It replaces the TODO-WASM: shadow stack maintenance comment. The emit pattern matches the SuppressGCTransition publish from #130924.

Test

WasmInterpreterTransitions gains R2RCallsNestingUco: an R2R method does a delegate* unmanaged calli into UcoWithInlinedPInvoke, which does a calli into a second UCO, UcoLeaf. Each calli is an inlined P/Invoke. The inner one lowers __stack_pointer, and the outer one's JIT_PInvokeEnd asserts that the global is back at the caller's SP.

Run against the Checked browser Core_Root from main CI build 1612746 (without #134606), with crossgen2 and the universal wasm JIT built from main:

test	main JIT	fixed JIT
main WasmInterpreterTransitions	pass	pass
with R2RCallsNestingUco	assert sp == stack_pointer_global_value (helpers.cpp:420), then GetFrame assert	pass (exit 100)

The test IL for this run was compiled locally with csc rather than through the repo's test build.

Validation

• Rebuilt clrjit_universal_wasm_arm64, swapped it into crossgen2, and recompiled WasmInterpreterTransitions against the CI Checked browser runtime from build 1612477, which includes [wasm][coreclr][R2R] Publish existing method bodies before interpreter fallback #134606's runtime changes.
  • Old JIT: reproduces the CI failure (the helpers.cpp:420 assert, then GetFrame asserts and OOB).
  • New JIT: passes with exit 100 and no asserts.
• Before the new test was added, a disassembly diff of composite-r2r.wasm showed exactly two changes: the two UCO epilogs (StreamLengthProxy and UnmanagedCallerCallsInterpreted), each with the 4-instruction restore added. Everything else was identical.
• The other 6 tests in the Helix readytorun work item pass: readytorun, fieldlayout, crossgen2smoke_donotalwaysusecrossgen2, Breadth1Test, Depth1Test, DynamicMethodGCStress.
• ./build.sh clr.alljits -c checked on osx-arm64, which includes the universal wasm JIT, builds cleanly, and jit-format reports no changes.

Related

• Unblocks the WasmInterpreterTransitions failure in [wasm][coreclr][R2R] Publish existing method bodies before interpreter fallback #134606. This PR is independent of [wasm][coreclr][R2R] Publish existing method bodies before interpreter fallback #134606 and has its own regression test on main, and [wasm][coreclr][R2R] Publish existing method bodies before interpreter fallback #134606 doesn't need to be stacked on it.
• Helps the open [wasm][coreclr] Dispatch R2R-compiled UnmanagedCallersOnly callbacks to their native entrypoint #134355, where native callers dispatch directly to R2R UCO bodies. There the UCO epilog is the only place that can restore __stack_pointer for the native caller.

Resolves #134681

Note

This PR description was generated with the help of GitHub Copilot.

[azure-pipelines]

Azure Pipelines:
Successfully started running 6 pipeline(s).
10 pipeline(s) were filtered out due to trigger conditions.
There may be pipelines that require an authorized user to comment /azp run to run.

[dotnet-policy-service]

Tagging subscribers to this area: @JulieLeeMSFT, @jakobbotsch
See info in area-owners.md if you want to be subscribed.

[dotnet-policy-service]

Tagging subscribers to 'arch-wasm': @lewing, @pavelsavara
See info in area-owners.md if you want to be subscribed.

[lewing]

@AndyAyersMS fixed the comment, reapprove when you get a chance