commit 1ab1eb6ce51b15cc04140d297bfb664985a28683
parent f8b8182ee6960e6ebc5d883ca30cc27a56a12fba
Author: Ryan Hunt <rhunt@eqrion.net>
Date: Thu, 18 Dec 2025 16:45:28 +0000
Bug 2002625 - wasm: Simplify codegen after merging data and suspender. r=yury
Now that they are actually the same thing, simplify their codegen.
This uncovered some latent issues:
1. SuspenderReg could sometimes be the ReturnReg, when we start using
SuspenderReg more often in the JIT functions this lead to a segfault.
2. InstanceReg is not pinned inside of Ion and must be manually provided.
InstanceReg is now manually provided to all the MIR functions.
ContinueOnSuspendable's register allocation is reworked to avoid the
segfault hazard. The other's will be improved in another bug, but don't
seem to have the hazard right now.
Differential Revision: https://phabricator.services.mozilla.com/D274191
Diffstat:
4 files changed, 68 insertions(+), 95 deletions(-)
diff --git a/js/src/jit/CodeGenerator.cpp b/js/src/jit/CodeGenerator.cpp
@@ -10569,7 +10569,6 @@ void CodeGenerator::visitWasmStackSwitchToSuspendable(
const Register SuspenderReg = lir->suspender()->toGeneralReg()->reg();
const Register FnReg = lir->fn()->toGeneralReg()->reg();
const Register DataReg = lir->data()->toGeneralReg()->reg();
- const Register SuspenderDataReg = ABINonArgReg3;
# ifdef JS_CODEGEN_ARM64
vixl::UseScratchRegisterScope temps(&masm);
@@ -10610,18 +10609,16 @@ void CodeGenerator::visitWasmStackSwitchToSuspendable(
masm.framePushed() - sizeof(wasm::Frame), WasmStackAlignment);
masm.reserveStack(reserve);
- masm.movePtr(SuspenderReg, SuspenderDataReg);
-
// Switch stacks to suspendable, keep original FP to maintain
// frames chain between main and suspendable stack segments.
masm.storeStackPtrToPrivateValue(
- Address(SuspenderDataReg, wasm::SuspenderObject::offsetOfMainSP()));
+ Address(SuspenderReg, wasm::SuspenderObject::offsetOfMainSP()));
masm.storePrivateValue(
FramePointer,
- Address(SuspenderDataReg, wasm::SuspenderObject::offsetOfMainFP()));
+ Address(SuspenderReg, wasm::SuspenderObject::offsetOfMainFP()));
- masm.loadStackPtrFromPrivateValue(Address(
- SuspenderDataReg, wasm::SuspenderObject::offsetOfSuspendableSP()));
+ masm.loadStackPtrFromPrivateValue(
+ Address(SuspenderReg, wasm::SuspenderObject::offsetOfSuspendableSP()));
masm.assertStackAlignment(WasmStackAlignment);
@@ -10654,7 +10651,7 @@ void CodeGenerator::visitWasmStackSwitchToSuspendable(
Address(masm.getStackPointer(), -int32_t(sizeof(wasm::Frame))),
ScratchReg2);
masm.storePrivateValue(
- ScratchReg2, Address(SuspenderDataReg,
+ ScratchReg2, Address(SuspenderReg,
wasm::SuspenderObject::offsetOfSuspendableExitFP()));
masm.mov(&returnCallsite, ReturnAddressReg);
@@ -10717,7 +10714,6 @@ void CodeGenerator::visitWasmStackSwitchToMain(LWasmStackSwitchToMain* lir) {
const Register SuspenderReg = lir->suspender()->toGeneralReg()->reg();
const Register FnReg = lir->fn()->toGeneralReg()->reg();
const Register DataReg = lir->data()->toGeneralReg()->reg();
- const Register SuspenderDataReg = ABINonArgReg3;
# ifdef JS_CODEGEN_ARM64
vixl::UseScratchRegisterScope temps(&masm);
@@ -10759,40 +10755,27 @@ void CodeGenerator::visitWasmStackSwitchToMain(LWasmStackSwitchToMain* lir) {
masm.framePushed() - sizeof(wasm::Frame), WasmStackAlignment);
masm.reserveStack(reserve);
- masm.movePtr(SuspenderReg, SuspenderDataReg);
+ masm.movePtr(SuspenderReg, SuspenderReg);
// Switch stacks to main.
- masm.storeStackPtrToPrivateValue(Address(
- SuspenderDataReg, wasm::SuspenderObject::offsetOfSuspendableSP()));
+ masm.storeStackPtrToPrivateValue(
+ Address(SuspenderReg, wasm::SuspenderObject::offsetOfSuspendableSP()));
masm.storePrivateValue(
- FramePointer, Address(SuspenderDataReg,
- wasm::SuspenderObject::offsetOfSuspendableFP()));
+ FramePointer,
+ Address(SuspenderReg, wasm::SuspenderObject::offsetOfSuspendableFP()));
masm.loadStackPtrFromPrivateValue(
- Address(SuspenderDataReg, wasm::SuspenderObject::offsetOfMainSP()));
+ Address(SuspenderReg, wasm::SuspenderObject::offsetOfMainSP()));
masm.loadPrivate(
- Address(SuspenderDataReg, wasm::SuspenderObject::offsetOfMainFP()),
+ Address(SuspenderReg, wasm::SuspenderObject::offsetOfMainFP()),
FramePointer);
// Set main_ra field to returnCallsite.
-# ifdef JS_CODEGEN_X86
- // SuspenderDataReg is also ScratchReg1, use DataReg as a scratch register.
- MOZ_ASSERT(ScratchReg1 == SuspenderDataReg);
- masm.push(DataReg);
- masm.mov(&returnCallsite, DataReg);
- masm.storePrivateValue(
- DataReg,
- Address(SuspenderDataReg,
- wasm::SuspenderObject::offsetOfSuspendedReturnAddress()));
- masm.pop(DataReg);
-# else
- MOZ_ASSERT(ScratchReg1 != SuspenderDataReg);
masm.mov(&returnCallsite, ScratchReg1);
masm.storePrivateValue(
ScratchReg1,
- Address(SuspenderDataReg,
+ Address(SuspenderReg,
wasm::SuspenderObject::offsetOfSuspendedReturnAddress()));
-# endif
masm.assertStackAlignment(WasmStackAlignment);
@@ -10826,12 +10809,12 @@ void CodeGenerator::visitWasmStackSwitchToMain(LWasmStackSwitchToMain* lir) {
ScratchReg2);
masm.storePrivateValue(
ScratchReg2,
- Address(SuspenderDataReg, wasm::SuspenderObject::offsetOfMainExitFP()));
+ Address(SuspenderReg, wasm::SuspenderObject::offsetOfMainExitFP()));
// Load InstanceReg from suspendable stack exit frame.
- masm.loadPrivate(Address(SuspenderDataReg,
- wasm::SuspenderObject::offsetOfSuspendableExitFP()),
- ScratchReg2);
+ masm.loadPrivate(
+ Address(SuspenderReg, wasm::SuspenderObject::offsetOfSuspendableExitFP()),
+ ScratchReg2);
masm.loadPtr(
Address(ScratchReg2, wasm::FrameWithInstances::callerInstanceOffset()),
ScratchReg2);
@@ -10839,9 +10822,9 @@ void CodeGenerator::visitWasmStackSwitchToMain(LWasmStackSwitchToMain* lir) {
WasmCallerInstanceOffsetBeforeCall));
// Load RA from suspendable stack exit frame.
- masm.loadPrivate(Address(SuspenderDataReg,
- wasm::SuspenderObject::offsetOfSuspendableExitFP()),
- ScratchReg1);
+ masm.loadPrivate(
+ Address(SuspenderReg, wasm::SuspenderObject::offsetOfSuspendableExitFP()),
+ ScratchReg1);
masm.loadPtr(Address(ScratchReg1, wasm::Frame::returnAddressOffset()),
ReturnAddressReg);
@@ -10911,29 +10894,13 @@ void CodeGenerator::visitWasmStackSwitchToMain(LWasmStackSwitchToMain* lir) {
void CodeGenerator::visitWasmStackContinueOnSuspendable(
LWasmStackContinueOnSuspendable* lir) {
#ifdef ENABLE_WASM_JSPI
- const Register SuspenderReg = lir->suspender()->toGeneralReg()->reg();
- const Register ResultReg = lir->result()->toGeneralReg()->reg();
- const Register SuspenderDataReg = ABINonArgReg3;
-
-# ifdef JS_CODEGEN_ARM64
- vixl::UseScratchRegisterScope temps(&masm);
- const Register ScratchReg1 = temps.AcquireX().asUnsized();
-# elif defined(JS_CODEGEN_X86)
- const Register ScratchReg1 = ABINonArgReturnReg1;
-# elif defined(JS_CODEGEN_X64)
- const Register ScratchReg1 = ScratchReg;
-# elif defined(JS_CODEGEN_ARM)
- const Register ScratchReg1 = ABINonArgReturnVolatileReg;
-# elif defined(JS_CODEGEN_LOONG64) || defined(JS_CODEGEN_RISCV64) || \
- defined(JS_CODEGEN_MIPS64)
- UseScratchRegisterScope temps(masm);
- const Register ScratchReg1 = temps.Acquire();
-# else
-# error "NYI: scratch register"
-# endif
- const Register ScratchReg2 = ABINonArgReg1;
+ MOZ_ASSERT(ToRegister(lir->instance()) == InstanceReg);
+ Register suspender = ToRegister(lir->suspender());
+ Register result = ToRegister(lir->result());
+ Register temp1 = ToRegister(lir->temp0());
+ Register temp2 = ToRegister(lir->temp1());
- masm.Push(SuspenderReg);
+ masm.Push(suspender);
int32_t framePushedAtSuspender = masm.framePushed();
masm.Push(InstanceReg);
@@ -10945,35 +10912,32 @@ void CodeGenerator::visitWasmStackContinueOnSuspendable(
masm.framePushed() - sizeof(wasm::Frame), WasmStackAlignment);
masm.reserveStack(reserve);
- masm.movePtr(SuspenderReg, SuspenderDataReg);
masm.storeStackPtrToPrivateValue(
- Address(SuspenderDataReg, wasm::SuspenderObject::offsetOfMainSP()));
+ Address(suspender, wasm::SuspenderObject::offsetOfMainSP()));
masm.storePrivateValue(
FramePointer,
- Address(SuspenderDataReg, wasm::SuspenderObject::offsetOfMainFP()));
+ Address(suspender, wasm::SuspenderObject::offsetOfMainFP()));
// Adjust exit frame FP.
- masm.loadPrivate(Address(SuspenderDataReg,
- wasm::SuspenderObject::offsetOfSuspendableExitFP()),
- ScratchReg1);
- masm.storePtr(FramePointer,
- Address(ScratchReg1, wasm::Frame::callerFPOffset()));
+ masm.loadPrivate(
+ Address(suspender, wasm::SuspenderObject::offsetOfSuspendableExitFP()),
+ temp1);
+ masm.storePtr(FramePointer, Address(temp1, wasm::Frame::callerFPOffset()));
// Adjust exit frame RA.
- masm.mov(&returnCallsite, ScratchReg2);
+ masm.mov(&returnCallsite, temp2);
- masm.storePtr(ScratchReg2,
- Address(ScratchReg1, wasm::Frame::returnAddressOffset()));
+ masm.storePtr(temp2, Address(temp1, wasm::Frame::returnAddressOffset()));
// Adjust exit frame caller instance slot.
masm.storePtr(
InstanceReg,
- Address(ScratchReg1, wasm::FrameWithInstances::callerInstanceOffset()));
+ Address(temp1, wasm::FrameWithInstances::callerInstanceOffset()));
// Switch stacks to suspendable.
- masm.loadStackPtrFromPrivateValue(Address(
- SuspenderDataReg, wasm::SuspenderObject::offsetOfSuspendableSP()));
+ masm.loadStackPtrFromPrivateValue(
+ Address(suspender, wasm::SuspenderObject::offsetOfSuspendableSP()));
masm.loadPrivate(
- Address(SuspenderDataReg, wasm::SuspenderObject::offsetOfSuspendableFP()),
+ Address(suspender, wasm::SuspenderObject::offsetOfSuspendableFP()),
FramePointer);
masm.assertStackAlignment(WasmStackAlignment);
@@ -11001,16 +10965,16 @@ void CodeGenerator::visitWasmStackContinueOnSuspendable(
masm.assertStackAlignment(WasmStackAlignment);
// Transfer results to ReturnReg so it will appear at SwitchToMain return.
- masm.mov(ResultReg, ReturnReg);
-
- const Register ReturnAddressReg = ScratchReg1;
+ // temp2 is fixed to be the ReturnReg, and so use it here.
+ MOZ_ASSERT(temp2 == ReturnReg);
+ masm.mov(result, temp2);
// Pretend we just returned from the function.
masm.loadPrivate(
- Address(SuspenderDataReg,
+ Address(suspender,
wasm::SuspenderObject::offsetOfSuspendedReturnAddress()),
- ReturnAddressReg);
- masm.jump(ReturnAddressReg);
+ temp1);
+ masm.jump(temp1);
// About to use valid FramePointer -- restore framePushed.
masm.setFramePushed(framePushed);
@@ -11036,13 +11000,12 @@ void CodeGenerator::visitWasmStackContinueOnSuspendable(
masm.freeStack(reserve);
masm.Pop(InstanceReg);
- masm.Pop(SuspenderReg);
+ masm.Pop(suspender);
- // Using SuspenderDataReg and ABINonArgReg2 as temps.
- masm.switchToWasmInstanceRealm(SuspenderDataReg, ABINonArgReg2);
+ masm.switchToWasmInstanceRealm(temp1, temp2);
callWasmUpdateSuspenderState(wasm::UpdateSuspenderStateAction::Leave,
- SuspenderReg, ScratchReg1);
+ suspender, temp1);
#else
MOZ_CRASH("NYI");
#endif // ENABLE_WASM_JSPI
diff --git a/js/src/jit/Lowering.cpp b/js/src/jit/Lowering.cpp
@@ -7335,6 +7335,7 @@ void LIRGenerator::visitWasmStackSwitchToSuspendable(
MWasmStackSwitchToSuspendable* ins) {
#ifdef ENABLE_WASM_JSPI
auto* lir = new (alloc()) LWasmStackSwitchToSuspendable(
+ useFixedAtStart(ins->instance(), InstanceReg),
useFixedAtStart(ins->suspender(), ABINonArgReg0),
useFixedAtStart(ins->fn(), ABINonArgReg1),
useFixedAtStart(ins->data(), ABINonArgReg2));
@@ -7349,7 +7350,8 @@ void LIRGenerator::visitWasmStackSwitchToSuspendable(
void LIRGenerator::visitWasmStackSwitchToMain(MWasmStackSwitchToMain* ins) {
#ifdef ENABLE_WASM_JSPI
auto* lir = new (alloc())
- LWasmStackSwitchToMain(useFixedAtStart(ins->suspender(), ABINonArgReg0),
+ LWasmStackSwitchToMain(useFixedAtStart(ins->instance(), InstanceReg),
+ useFixedAtStart(ins->suspender(), ABINonArgReg0),
useFixedAtStart(ins->fn(), ABINonArgReg1),
useFixedAtStart(ins->data(), ABINonArgReg2));
@@ -7364,8 +7366,9 @@ void LIRGenerator::visitWasmStackContinueOnSuspendable(
MWasmStackContinueOnSuspendable* ins) {
#ifdef ENABLE_WASM_JSPI
auto* lir = new (alloc()) LWasmStackContinueOnSuspendable(
- useFixedAtStart(ins->suspender(), ABINonArgReg0),
- useFixedAtStart(ins->result(), ABINonArgReg2));
+ useFixedAtStart(ins->instance(), InstanceReg),
+ useRegisterAtStart(ins->suspender()), useRegisterAtStart(ins->result()),
+ tempFixed(ABINonArgReturnReg0), tempFixed(ReturnReg));
add(lir, ins);
assignWasmSafepoint(lir);
diff --git a/js/src/jit/MIROps.yaml b/js/src/jit/MIROps.yaml
@@ -4386,8 +4386,9 @@
- name: WasmStackSwitchToMain
operands:
- suspender: Object
- fn: Object
+ instance: WasmAnyRef
+ suspender: WasmAnyRef
+ fn: WasmAnyRef
data: WasmAnyRef
result_type: WasmAnyRef
type_policy: none
@@ -4396,8 +4397,9 @@
- name: WasmStackSwitchToSuspendable
operands:
- suspender: Object
- fn: Object
+ instance: WasmAnyRef
+ suspender: WasmAnyRef
+ fn: WasmAnyRef
data: WasmAnyRef
type_policy: none
possibly_calls: true
@@ -4405,11 +4407,13 @@
- name: WasmStackContinueOnSuspendable
operands:
- suspender: Object
+ instance: WasmAnyRef
+ suspender: WasmAnyRef
result: WasmAnyRef
type_policy: none
possibly_calls: true
generate_lir: true
+ lir_temps: 2
- name: WasmBinarySimd128
gen_boilerplate: false
diff --git a/js/src/wasm/WasmIonCompile.cpp b/js/src/wasm/WasmIonCompile.cpp
@@ -3291,13 +3291,16 @@ class FunctionCompiler {
MInstruction* ins;
switch (kind) {
case StackSwitchKind::SwitchToMain:
- ins = MWasmStackSwitchToMain::New(alloc(), suspender, fn, data);
+ ins = MWasmStackSwitchToMain::New(alloc(), instancePointer_, suspender,
+ fn, data);
break;
case StackSwitchKind::SwitchToSuspendable:
- ins = MWasmStackSwitchToSuspendable::New(alloc(), suspender, fn, data);
+ ins = MWasmStackSwitchToSuspendable::New(alloc(), instancePointer_,
+ suspender, fn, data);
break;
case StackSwitchKind::ContinueOnSuspendable:
- ins = MWasmStackContinueOnSuspendable::New(alloc(), suspender, data);
+ ins = MWasmStackContinueOnSuspendable::New(alloc(), instancePointer_,
+ suspender, data);
break;
}
if (!ins) {
