commit 30e7d973d80ca6ee2ac9d9fd2ff415865c76683e
parent 527c3ea6fc525402a16460a8d8e7e547e92bd05d
Author: Ryan Hunt <rhunt@eqrion.net>
Date: Thu, 18 Dec 2025 16:45:28 +0000
Bug 2002625 - wasm: Dynamically switch to main stack in builtin thunks. r=yury
Now that the sequence of operations performed in a CallOnMain
operation have been simplified, we can perform them relatively
easily in JIT code.
Extend the exit prologue and epilogues to optionally emit a
dynamic switch to the main stack when they are on a suspender.
A new SMDOC is added to explain this process.
Differential Revision: https://phabricator.services.mozilla.com/D274192
Diffstat:
16 files changed, 507 insertions(+), 106 deletions(-)
diff --git a/js/public/RootingAPI.h b/js/public/RootingAPI.h
@@ -1249,6 +1249,8 @@ class MOZ_RAII Rooted : public detail::RootedTraits<T>::StackBase,
T* address() { return &ptr; }
const T* address() const { return &ptr; }
+ static constexpr size_t offsetOfPtr() { return offsetof(Rooted, ptr); }
+
private:
T ptr;
diff --git a/js/src/jit/JitFrames.cpp b/js/src/jit/JitFrames.cpp
@@ -1416,8 +1416,7 @@ static void TraceJitActivation(JSTracer* trc, JitActivation* activation) {
}
#endif
- activation->traceRematerializedFrames(trc);
- activation->traceIonRecovery(trc);
+ activation->trace(trc);
// This is used for sanity checking continuity of the sequence of wasm stack
// maps as we unwind. It has no functional purpose.
diff --git a/js/src/jit/JitRuntime.h b/js/src/jit/JitRuntime.h
@@ -340,6 +340,9 @@ class JitRuntime {
const void* addressOfDisallowArbitraryCode() const {
return &disallowArbitraryCode_.refNoCheck();
}
+ static size_t offsetOfDisallowArbitraryCode() {
+ return offsetof(JitRuntime, disallowArbitraryCode_);
+ }
#endif
uint8_t* allocateIonOsrTempData(size_t size);
diff --git a/js/src/threading/ProtectedData.h b/js/src/threading/ProtectedData.h
@@ -156,7 +156,7 @@ class ProtectedData {
T& refNoCheck() { return value; }
const T& refNoCheck() const { return value; }
- static size_t offsetOfValue() { return offsetof(ThisType, value); }
+ static constexpr size_t offsetOfValue() { return offsetof(ThisType, value); }
private:
T value;
diff --git a/js/src/vm/JSContext.h b/js/src/vm/JSContext.h
@@ -396,6 +396,10 @@ struct JS_PUBLIC_API JSContext : public JS::RootingContext,
JSRuntime* runtime() { return runtime_; }
const JSRuntime* runtime() const { return runtime_; }
+ static size_t offsetOfRuntime() {
+ return offsetof(JSContext, runtime_) +
+ js::UnprotectedData<JSRuntime*>::offsetOfValue();
+ }
static size_t offsetOfRealm() { return offsetof(JSContext, realm_); }
friend class JS::AutoSaveExceptionState;
diff --git a/js/src/vm/JitActivation.cpp b/js/src/vm/JitActivation.cpp
@@ -34,6 +34,9 @@ js::jit::JitActivation::JitActivation(JSContext* cx)
: Activation(cx, Jit),
packedExitFP_(nullptr),
encodedWasmExitReason_(0),
+#ifdef ENABLE_WASM_JSPI
+ wasmExitSuspender_(cx, nullptr),
+#endif
prevJitActivation_(cx->jitActivation),
ionRecovery_(cx),
bailoutData_(nullptr),
@@ -64,6 +67,20 @@ js::jit::JitActivation::~JitActivation() {
MOZ_ASSERT_IF(rematerializedFrames_, rematerializedFrames_->empty());
}
+void js::jit::JitActivation::trace(JSTracer* trc) {
+ if (rematerializedFrames_) {
+ for (RematerializedFrameTable::Enum e(*rematerializedFrames_); !e.empty();
+ e.popFront()) {
+ e.front().value().trace(trc);
+ }
+ }
+
+ for (RInstructionResults* it = ionRecovery_.begin(); it != ionRecovery_.end();
+ it++) {
+ it->trace(trc);
+ }
+}
+
void js::jit::JitActivation::setBailoutData(
jit::BailoutFrameInfo* bailoutData) {
MOZ_ASSERT(!bailoutData_);
@@ -178,16 +195,6 @@ void js::jit::JitActivation::removeRematerializedFramesFromDebugger(
}
}
-void js::jit::JitActivation::traceRematerializedFrames(JSTracer* trc) {
- if (!rematerializedFrames_) {
- return;
- }
- for (RematerializedFrameTable::Enum e(*rematerializedFrames_); !e.empty();
- e.popFront()) {
- e.front().value().trace(trc);
- }
-}
-
bool js::jit::JitActivation::registerIonFrameRecovery(
RInstructionResults&& results) {
// Check that there is no entry in the vector yet.
@@ -220,13 +227,6 @@ void js::jit::JitActivation::removeIonFrameRecovery(JitFrameLayout* fp) {
ionRecovery_.erase(elem);
}
-void js::jit::JitActivation::traceIonRecovery(JSTracer* trc) {
- for (RInstructionResults* it = ionRecovery_.begin(); it != ionRecovery_.end();
- it++) {
- it->trace(trc);
- }
-}
-
void js::jit::JitActivation::startWasmTrap(wasm::Trap trap,
const wasm::TrapSite& trapSite,
const wasm::RegisterState& state) {
diff --git a/js/src/vm/JitActivation.h b/js/src/vm/JitActivation.h
@@ -30,6 +30,7 @@
#include "wasm/WasmConstants.h" // js::wasm::Trap
#include "wasm/WasmFrame.h" // js::wasm::Frame
#include "wasm/WasmFrameIter.h" // js::wasm::{ExitReason,RegisterState,WasmFrameIter}
+#include "wasm/WasmPI.h" // js::wasm::SuspenderObject
struct JS_PUBLIC_API JSContext;
class JS_PUBLIC_API JSTracer;
@@ -52,6 +53,12 @@ class JitActivation : public Activation {
// When hasWasmExitFP(), encodedWasmExitReason_ holds ExitReason.
uint32_t encodedWasmExitReason_;
+#ifdef ENABLE_WASM_JSPI
+ // This would be a 'Rooted', except that the 'Rooted' values in the super
+ // class `Activation` conflict with the LIFO ordering that 'Rooted' requires.
+ // So instead we manually trace it.
+ JS::Rooted<wasm::SuspenderObject*> wasmExitSuspender_;
+#endif
JitActivation* prevJitActivation_;
@@ -108,6 +115,8 @@ class JitActivation : public Activation {
explicit JitActivation(JSContext* cx);
~JitActivation();
+ void trace(JSTracer* trc);
+
bool isProfiling() const {
// All JitActivations can be profiled.
return true;
@@ -174,8 +183,6 @@ class JitActivation : public Activation {
// Remove a previous rematerialization by fp.
void removeRematerializedFrame(uint8_t* top);
- void traceRematerializedFrames(JSTracer* trc);
-
// Register the results of on Ion frame recovery.
bool registerIonFrameRecovery(RInstructionResults&& results);
@@ -186,8 +193,6 @@ class JitActivation : public Activation {
// from the activation.
void removeIonFrameRecovery(JitFrameLayout* fp);
- void traceIonRecovery(JSTracer* trc);
-
// Return the bailout information if it is registered.
const BailoutFrameInfo* bailoutData() const { return bailoutData_; }
@@ -235,6 +240,16 @@ class JitActivation : public Activation {
static size_t offsetOfEncodedWasmExitReason() {
return offsetof(JitActivation, encodedWasmExitReason_);
}
+#ifdef ENABLE_WASM_JSPI
+ wasm::SuspenderObject* wasmExitSuspender() const {
+ MOZ_ASSERT(hasWasmExitFP());
+ return wasmExitSuspender_;
+ }
+ static size_t offsetOfWasmExitSuspender() {
+ return offsetof(JitActivation, wasmExitSuspender_) +
+ Rooted<wasm::SuspenderObject*>::offsetOfPtr();
+ }
+#endif
void startWasmTrap(wasm::Trap trap, const wasm::TrapSite& trapSite,
const wasm::RegisterState& state);
diff --git a/js/src/vm/Runtime.h b/js/src/vm/Runtime.h
@@ -742,6 +742,10 @@ struct JSRuntime {
[[nodiscard]] bool createJitRuntime(JSContext* cx);
js::jit::JitRuntime* jitRuntime() const { return jitRuntime_.ref(); }
bool hasJitRuntime() const { return !!jitRuntime_; }
+ static constexpr size_t offsetOfJitRuntime() {
+ return offsetof(JSRuntime, jitRuntime_) +
+ js::UnprotectedData<js::jit::JitRuntime*>::offsetOfValue();
+ }
private:
// Used to generate random keys for hash tables.
diff --git a/js/src/wasm/WasmBuiltins.cpp b/js/src/wasm/WasmBuiltins.cpp
@@ -1930,6 +1930,23 @@ bool wasm::NeedsBuiltinThunk(SymbolicAddress sym) {
MOZ_CRASH("unexpected symbolic address");
}
+static bool NeedsDynamicSwitchToMainStack(SymbolicAddress sym) {
+ MOZ_ASSERT(NeedsBuiltinThunk(sym));
+ switch (sym) {
+#if ENABLE_WASM_JSPI
+ // These builtins must run on the suspendable so that they can access the
+ // wasm::Context::activeSuspender().
+ case SymbolicAddress::UpdateSuspenderState:
+ case SymbolicAddress::CurrentSuspender:
+ return false;
+#endif
+
+ // Nothing else should be running on a suspendable stack right now.
+ default:
+ return true;
+ }
+}
+
// ============================================================================
// [SMDOC] JS Fast Wasm Imports
//
@@ -2147,7 +2164,8 @@ bool wasm::EnsureBuiltinThunksInitialized(
MOZ_ASSERT(ABIForBuiltin(sym) == ABIKind::Wasm);
CallableOffsets offsets;
- if (!GenerateBuiltinThunk(masm, abiType, exitReason, funcPtr, &offsets)) {
+ if (!GenerateBuiltinThunk(masm, abiType, NeedsDynamicSwitchToMainStack(sym),
+ exitReason, funcPtr, &offsets)) {
return false;
}
if (!thunks->codeRanges.emplaceBack(CodeRange::BuiltinThunk, offsets)) {
@@ -2175,7 +2193,8 @@ bool wasm::EnsureBuiltinThunksInitialized(
ExitReason exitReason = ExitReason::Fixed::BuiltinNative;
CallableOffsets offsets;
- if (!GenerateBuiltinThunk(masm, abiType, exitReason, funcPtr, &offsets)) {
+ if (!GenerateBuiltinThunk(masm, abiType, /*dynamicSwitchToMainStack*/ true,
+ exitReason, funcPtr, &offsets)) {
return false;
}
if (!thunks->codeRanges.emplaceBack(CodeRange::BuiltinThunk, offsets)) {
diff --git a/js/src/wasm/WasmContext.cpp b/js/src/wasm/WasmContext.cpp
@@ -63,9 +63,9 @@ void Context::initStackLimit(JSContext* cx) {
// See the comment on wasm::Context for why we do this.
#ifdef ENABLE_WASM_JSPI
# if defined(_WIN32)
- _NT_TIB* tib = reinterpret_cast<_NT_TIB*>(::NtCurrentTeb());
- tibStackBase_ = tib->StackBase;
- tibStackLimit_ = tib->StackLimit;
+ tib_ = reinterpret_cast<_NT_TIB*>(::NtCurrentTeb());
+ tibStackBase_ = tib_->StackBase;
+ tibStackLimit_ = tib_->StackLimit;
# endif
#endif
}
@@ -101,11 +101,10 @@ void Context::enterSuspendableStack(JSContext* cx, SuspenderObject* suspender) {
// See the comment on wasm::Context for why we do this.
# if defined(_WIN32)
- _NT_TIB* tib = reinterpret_cast<_NT_TIB*>(::NtCurrentTeb());
- tibStackBase_ = tib->StackBase;
- tibStackLimit = tib->StackLimit;
- tib->StackBase = reinterpret_cast<void*>(suspender->stackMemoryBase());
- tib->StackLimit =
+ tibStackBase_ = tib_->StackBase;
+ tibStackLimit_ = tib_->StackLimit;
+ tib_->StackBase = reinterpret_cast<void*>(suspender->stackMemoryBase());
+ tib_->StackLimit =
reinterpret_cast<void*>(suspender->stackMemoryLimitForSystem());
# endif
@@ -121,9 +120,8 @@ void Context::leaveSuspendableStack(JSContext* cx) {
// See the comment on wasm::Context for why we do this.
# if defined(_WIN32)
- _NT_TIB* tib = reinterpret_cast<_NT_TIB*>(::NtCurrentTeb());
- tib->StackBase = tibStackBase_;
- tib->StackLimit = tibStackLimit_;
+ tib_->StackBase = static_cast<void*>(tibStackBase_);
+ tib_->StackLimit = static_cast<void*>(tibStackLimit_);
# endif
# ifdef DEBUG
diff --git a/js/src/wasm/WasmContext.h b/js/src/wasm/WasmContext.h
@@ -25,6 +25,10 @@
#include "js/NativeStackLimits.h"
+#ifdef _WIN32
+struct _NT_TIB;
+#endif
+
namespace js::wasm {
#ifdef ENABLE_WASM_JSPI
@@ -45,6 +49,9 @@ class Context {
static constexpr size_t offsetOfStackLimit() {
return offsetof(Context, stackLimit);
}
+ static constexpr size_t offsetOfMainStackLimit() {
+ return offsetof(Context, mainStackLimit);
+ }
void initStackLimit(JSContext* cx);
@@ -52,6 +59,15 @@ class Context {
static constexpr size_t offsetOfActiveSuspender() {
return offsetof(Context, activeSuspender_);
}
+# ifdef _WIN32
+ static constexpr size_t offsetOfTib() { return offsetof(Context, tib_); }
+ static constexpr size_t offsetOfTibStackBase() {
+ return offsetof(Context, tibStackBase_);
+ }
+ static constexpr size_t offsetOfTibStackLimit() {
+ return offsetof(Context, tibStackLimit_);
+ }
+# endif
SuspenderObject* activeSuspender() { return activeSuspender_; }
bool onSuspendableStack() const { return activeSuspender_ != nullptr; }
@@ -80,7 +96,9 @@ class Context {
# if defined(_WIN32)
// On WIN64, the Thread Information Block stack limits must be updated on
// stack switches to avoid failures on SP checks during vectored exeption
- // handling for traps. We store the original ones here for easy restoration.
+ // handling for traps. We store the original limits and the TIB here for
+ // easy restoration.
+ _NT_TIB* tib_ = nullptr;
void* tibStackBase_ = nullptr;
void* tibStackLimit_ = nullptr;
# endif
diff --git a/js/src/wasm/WasmFrameIter.cpp b/js/src/wasm/WasmFrameIter.cpp
@@ -19,6 +19,7 @@
#include "wasm/WasmFrameIter.h"
#include "jit/JitFrames.h"
+#include "jit/JitRuntime.h"
#include "jit/shared/IonAssemblerBuffer.h" // jit::BufferOffset
#include "js/ColumnNumber.h" // JS::WasmFunctionIndex, LimitedColumnNumberOneOrigin, JS::TaggedColumnNumberOneOrigin, JS::TaggedColumnNumberOneOrigin
#include "vm/JitActivation.h" // js::jit::JitActivation
@@ -34,6 +35,15 @@
#include "jit/MacroAssembler-inl.h"
#include "wasm/WasmInstance-inl.h"
+#ifdef XP_WIN
+// We only need the `windows.h` header, but this file can get unified built
+// with WasmSignalHandlers.cpp, which requires `winternal.h` to be included
+// before the `windows.h` header, and so we must include it here for that case.
+# include <winternl.h> // must include before util/WindowsWrapper.h's `#undef`s
+
+# include "util/WindowsWrapper.h"
+#endif
+
using namespace js;
using namespace js::jit;
using namespace js::wasm;
@@ -528,35 +538,38 @@ static const unsigned PoppedFPJitEntry = 4;
# error "Unknown architecture!"
#endif
-static void LoadActivation(MacroAssembler& masm, const Register& dest) {
+void wasm::LoadActivation(MacroAssembler& masm, Register instance,
+ Register dest) {
// WasmCall pushes a JitActivation.
- masm.loadPtr(Address(InstanceReg, wasm::Instance::offsetOfCx()), dest);
+ masm.loadPtr(Address(instance, wasm::Instance::offsetOfCx()), dest);
masm.loadPtr(Address(dest, JSContext::offsetOfActivation()), dest);
}
void wasm::SetExitFP(MacroAssembler& masm, ExitReason reason,
- Register scratch) {
+ Register activation, Register scratch) {
MOZ_ASSERT(!reason.isNone());
+ MOZ_ASSERT(activation != scratch);
- LoadActivation(masm, scratch);
-
+ // Write the encoded exit reason to the activation
masm.store32(
Imm32(reason.encode()),
- Address(scratch, JitActivation::offsetOfEncodedWasmExitReason()));
+ Address(activation, JitActivation::offsetOfEncodedWasmExitReason()));
+
+ // Tag the frame pointer in a different register so that we don't break
+ // async profiler unwinding.
+ masm.orPtr(Imm32(ExitFPTag), FramePointer, scratch);
- masm.orPtr(Imm32(ExitFPTag), FramePointer);
- masm.storePtr(FramePointer,
- Address(scratch, JitActivation::offsetOfPackedExitFP()));
- masm.andPtr(Imm32(int32_t(~ExitFPTag)), FramePointer);
+ // Write the tagged exitFP to the activation
+ masm.storePtr(scratch,
+ Address(activation, JitActivation::offsetOfPackedExitFP()));
}
-void wasm::ClearExitFP(MacroAssembler& masm, Register scratch) {
- LoadActivation(masm, scratch);
+void wasm::ClearExitFP(MacroAssembler& masm, Register activation) {
masm.storePtr(ImmWord(0x0),
- Address(scratch, JitActivation::offsetOfPackedExitFP()));
+ Address(activation, JitActivation::offsetOfPackedExitFP()));
masm.store32(
Imm32(0x0),
- Address(scratch, JitActivation::offsetOfEncodedWasmExitReason()));
+ Address(activation, JitActivation::offsetOfEncodedWasmExitReason()));
}
static void GenerateCallablePrologue(MacroAssembler& masm, uint32_t* entry) {
@@ -654,17 +667,13 @@ static void GenerateCallablePrologue(MacroAssembler& masm, uint32_t* entry) {
}
static void GenerateCallableEpilogue(MacroAssembler& masm, unsigned framePushed,
- ExitReason reason, uint32_t* ret) {
+ uint32_t* ret) {
AutoCreatedBy acb(masm, "GenerateCallableEpilogue");
if (framePushed) {
masm.freeStack(framePushed);
}
- if (!reason.isNone()) {
- ClearExitFP(masm, ABINonArgReturnVolatileReg);
- }
-
DebugOnly<uint32_t> poppedFP{};
#if defined(JS_CODEGEN_MIPS64)
@@ -759,7 +768,7 @@ void wasm::GenerateMinimalPrologue(MacroAssembler& masm, uint32_t* entry) {
// Generate the most minimal possible epilogue: `pop FP; return`.
void wasm::GenerateMinimalEpilogue(MacroAssembler& masm, uint32_t* ret) {
MOZ_ASSERT(masm.framePushed() == 0);
- GenerateCallableEpilogue(masm, /*framePushed=*/0, ExitReason::None(), ret);
+ GenerateCallableEpilogue(masm, /*framePushed=*/0, ret);
}
void wasm::GenerateFunctionPrologue(MacroAssembler& masm,
@@ -946,30 +955,249 @@ void wasm::GenerateFunctionEpilogue(MacroAssembler& masm, unsigned framePushed,
FuncOffsets* offsets) {
// Inverse of GenerateFunctionPrologue:
MOZ_ASSERT(masm.framePushed() == framePushed);
- GenerateCallableEpilogue(masm, framePushed, ExitReason::None(),
- &offsets->ret);
+ GenerateCallableEpilogue(masm, framePushed, &offsets->ret);
MOZ_ASSERT(masm.framePushed() == 0);
}
-void wasm::GenerateExitPrologue(MacroAssembler& masm, unsigned framePushed,
- ExitReason reason, CallableOffsets* offsets) {
- masm.haltingAlign(CodeAlignment);
+#ifdef ENABLE_WASM_JSPI
+void wasm::GenerateExitPrologueMainStackSwitch(MacroAssembler& masm,
+ Register instance,
+ Register scratch1,
+ Register scratch2,
+ Register scratch3) {
+ // Load the JSContext from the Instance into scratch1.
+ masm.loadPtr(Address(instance, wasm::Instance::offsetOfCx()), scratch1);
+
+ // Load wasm::Context::activeSuspender_ into scratch2.
+ masm.loadPtr(Address(scratch1, JSContext::offsetOfWasm() +
+ wasm::Context::offsetOfActiveSuspender()),
+ scratch2);
+
+ // If the activeSuspender_ is non-null, then we're on a suspendable stack
+ // and need to switch to the main stack.
+ Label alreadyOnSystemStack;
+ masm.branchTestPtr(Assembler::Zero, scratch2, scratch2,
+ &alreadyOnSystemStack);
+
+ // Reset the stack limit on wasm::Context to the main stack limit. We
+ // clobber scratch3 here.
+ masm.loadPtr(Address(scratch1, JSContext::offsetOfWasm() +
+ wasm::Context::offsetOfMainStackLimit()),
+ scratch3);
+ masm.storePtr(scratch3,
+ Address(scratch1, JSContext::offsetOfWasm() +
+ wasm::Context::offsetOfStackLimit()));
+
+ // Clear wasm::Context::activeSuspender_.
+ masm.storePtr(
+ ImmWord(0),
+ Address(scratch1, JSContext::offsetOfWasm() +
+ wasm::Context::offsetOfActiveSuspender()));
+
+ // Load the JitActivation from JSContext, and store the activeSuspender
+ // into wasmExitSuspender_. We clobber scratch3 here.
+ masm.loadPtr(Address(scratch1, JSContext::offsetOfActivation()), scratch3);
+ masm.storePtr(scratch2,
+ Address(scratch3, JitActivation::offsetOfWasmExitSuspender()));
+
+ // Switch the suspender's state to CalledOnMain.
+ masm.storeValue(JS::Int32Value(wasm::SuspenderState::CalledOnMain),
+ Address(scratch2, SuspenderObject::offsetOfState()));
+
+ // Switch the active SP to the Suspender's MainSP.
+ //
+ // NOTE: the FP is still pointing at our frame on the suspendable stack.
+ // This lets us address our incoming stack arguments using FP, and also
+ // switch back to the suspendable stack on return.
+ masm.loadStackPtrFromPrivateValue(
+ Address(scratch2, wasm::SuspenderObject::offsetOfMainSP()));
+
+ // Clear the disallow arbitrary code flag that is set when we enter a
+ // suspendable stack.
+# ifdef DEBUG
+ masm.loadPtr(Address(scratch1, JSContext::offsetOfRuntime()), scratch3);
+ masm.loadPtr(Address(scratch3, JSRuntime::offsetOfJitRuntime()), scratch3);
+ masm.store32(Imm32(0),
+ Address(scratch3, JitRuntime::offsetOfDisallowArbitraryCode()));
+# endif
+
+ // Update the Win32 TIB StackBase and StackLimit fields last. We clobber
+ // scratch2 and scratch3 here.
+# ifdef _WIN32
+ masm.loadPtr(Address(scratch1, JSContext::offsetOfWasm() +
+ wasm::Context::offsetOfTib()),
+ scratch2);
+ masm.loadPtr(Address(scratch1, JSContext::offsetOfWasm() +
+ wasm::Context::offsetOfTibStackBase()),
+ scratch3);
+ masm.storePtr(scratch3, Address(scratch2, offsetof(_NT_TIB, StackBase)));
+ masm.loadPtr(Address(scratch1, JSContext::offsetOfWasm() +
+ wasm::Context::offsetOfTibStackLimit()),
+ scratch3);
+ masm.storePtr(scratch3, Address(scratch2, offsetof(_NT_TIB, StackLimit)));
+# endif
+
+ masm.bind(&alreadyOnSystemStack);
+}
+
+void wasm::GenerateExitEpilogueMainStackReturn(MacroAssembler& masm,
+ Register instance,
+ Register activationAndScratch1,
+ Register scratch2) {
+ // scratch1 starts out with the JitActivation already loaded.
+ Register scratch1 = activationAndScratch1;
+
+ // Load JitActivation::wasmExitSuspender_ into scratch2.
+ masm.loadPtr(Address(scratch1, JitActivation::offsetOfWasmExitSuspender()),
+ scratch2);
+
+ // If wasmExitSuspender_ is null, then we were originally on the main stack
+ // and have no work to do here.
+ Label originallyOnSystemStack;
+ masm.branchTestPtr(Assembler::Zero, scratch2, scratch2,
+ &originallyOnSystemStack);
+
+ // Clear JitActivation::wasmExitSuspender.
+ masm.storePtr(ImmWord(0),
+ Address(scratch1, JitActivation::offsetOfWasmExitSuspender()));
+
+ // Restore the Suspender state back to Active.
+ masm.storeValue(JS::Int32Value(wasm::SuspenderState::Active),
+ Address(scratch2, SuspenderObject::offsetOfState()));
+
+ // We no longer need the JitActivation, reload the JSContext from
+ // instance into scratch1.
+ masm.loadPtr(Address(instance, wasm::Instance::offsetOfCx()), scratch1);
+
+ // Restore wasm::Context::activeSuspender_ using the wasmExitSuspender_.
+ masm.storePtr(
+ scratch2,
+ Address(scratch1, JSContext::offsetOfWasm() +
+ wasm::Context::offsetOfActiveSuspender()));
+
+ // Reset the stack limit to the suspender stack limit. This clobbers the
+ // suspender/scratch2, but it can now be reloaded from
+ // wasm::Context::activeSuspender_.
+ masm.loadPrivate(Address(scratch2, SuspenderObject::offsetOfStackMemory()),
+ scratch2);
+ masm.addPtr(Imm32(SuspendableRedZoneSize), scratch2);
+ masm.storePtr(scratch2,
+ Address(scratch1, JSContext::offsetOfWasm() +
+ wasm::Context::offsetOfStackLimit()));
+
+ // Update the Win32 TIB StackBase and StackLimit fields. This code is
+ // really register constrained and would benefit if we could use the Win32
+ // TIB directly through its segment register in masm.
+# ifdef _WIN32
+ // Load the TIB into scratch2.
+ masm.loadPtr(Address(scratch1, JSContext::offsetOfWasm() +
+ wasm::Context::offsetOfTib()),
+ scratch2);
+
+ // Load the sytem stack limit for this suspender and store to
+ // TIB->StackLimit. This clobbers scratch1.
+ masm.loadPtr(Address(scratch1, JSContext::offsetOfWasm() +
+ wasm::Context::offsetOfActiveSuspender()),
+ scratch1);
+ masm.loadPtr(Address(scratch1, wasm::SuspenderObject::offsetOfStackMemory()),
+ scratch1);
+ masm.storePtr(scratch1, Address(scratch2, offsetof(_NT_TIB, StackBase)));
+
+ // Reload JSContext into scratch1.
+ masm.loadPtr(Address(instance, wasm::Instance::offsetOfCx()), scratch1);
+
+ // Compute the stack base for this suspender and store to TIB->StackBase.
+ // This clobbers scratch1.
+ masm.loadPtr(Address(scratch1, JSContext::offsetOfWasm() +
+ wasm::Context::offsetOfActiveSuspender()),
+ scratch1);
+ masm.loadPtr(Address(scratch1, wasm::SuspenderObject::offsetOfStackMemory()),
+ scratch1);
+ masm.addPtr(Imm32(SuspendableStackPlusRedZoneSize), scratch1);
+ masm.storePtr(scratch1, Address(scratch2, offsetof(_NT_TIB, StackBase)));
+
+ // Reload JSContext into scratch1.
+ masm.loadPtr(Address(instance, wasm::Instance::offsetOfCx()), scratch1);
+# endif
+
+ // Set the disallow arbitrary code flag now that we're going back to a
+ // suspendable stack.
+# ifdef DEBUG
+ masm.loadPtr(Address(scratch1, JSContext::offsetOfRuntime()), scratch1);
+ masm.loadPtr(Address(scratch1, JSRuntime::offsetOfJitRuntime()), scratch1);
+ masm.store32(Imm32(1),
+ Address(scratch1, JitRuntime::offsetOfDisallowArbitraryCode()));
+# endif
+
+ masm.bind(&originallyOnSystemStack);
+}
+#endif // ENABLE_WASM_JSPI
+
+void wasm::GenerateExitPrologue(MacroAssembler& masm, ExitReason reason,
+ bool switchToMainStack,
+ unsigned framePushedPreSwitch,
+ unsigned framePushedPostSwitch,
+ CallableOffsets* offsets) {
+ MOZ_ASSERT(masm.framePushed() == 0);
+ masm.haltingAlign(CodeAlignment);
GenerateCallablePrologue(masm, &offsets->begin);
+ Register scratch1 = ABINonArgReg0;
+ Register scratch2 = ABINonArgReg1;
+#ifdef ENABLE_WASM_JSPI
+ Register scratch3 = ABINonArgReg2;
+#endif
+
// This frame will be exiting compiled code to C++ so record the fp and
// reason in the JitActivation so the frame iterators can unwind.
- SetExitFP(masm, reason, ABINonArgReturnVolatileReg);
+ LoadActivation(masm, InstanceReg, scratch1);
+ SetExitFP(masm, reason, scratch1, scratch2);
- MOZ_ASSERT(masm.framePushed() == 0);
- masm.reserveStack(framePushed);
+#ifdef ENABLE_WASM_JSPI
+ if (switchToMainStack) {
+ masm.reserveStack(framePushedPreSwitch);
+
+ GenerateExitPrologueMainStackSwitch(masm, InstanceReg, scratch1, scratch2,
+ scratch3);
+
+ // We may be on another stack now, reset the framePushed.
+ masm.setFramePushed(0);
+ masm.reserveStack(framePushedPostSwitch);
+ } else {
+ masm.reserveStack(framePushedPreSwitch + framePushedPostSwitch);
+ }
+#else
+ masm.reserveStack(framePushedPreSwitch + framePushedPostSwitch);
+#endif // ENABLE_WASM_JSPI
}
-void wasm::GenerateExitEpilogue(MacroAssembler& masm, unsigned framePushed,
- ExitReason reason, CallableOffsets* offsets) {
- // Inverse of GenerateExitPrologue:
- MOZ_ASSERT(masm.framePushed() == framePushed);
- GenerateCallableEpilogue(masm, framePushed, reason, &offsets->ret);
+void wasm::GenerateExitEpilogue(MacroAssembler& masm, ExitReason reason,
+ bool switchToMainStack,
+ CallableOffsets* offsets) {
+ Register scratch1 = ABINonArgReturnReg0;
+#if ENABLE_WASM_JSPI
+ Register scratch2 = ABINonArgReturnReg1;
+#endif
+
+ LoadActivation(masm, InstanceReg, scratch1);
+ ClearExitFP(masm, scratch1);
+
+#ifdef ENABLE_WASM_JSPI
+ // The exit prologue may have switched from a suspender's stack to the main
+ // stack, and we need to detect this and revert back to the suspender's
+ // stack. See GenerateExitPrologue for more information.
+ if (switchToMainStack) {
+ GenerateExitEpilogueMainStackReturn(masm, InstanceReg, scratch1, scratch2);
+ }
+#endif // ENABLE_WASM_JSPI
+
+ // Reset our stack pointer back to the frame pointer. This may switch the
+ // stack pointer back to our original stack.
+ masm.moveToStackPtr(FramePointer);
+ masm.setFramePushed(0);
+
+ GenerateCallableEpilogue(masm, /*framePushed*/ 0, &offsets->ret);
MOZ_ASSERT(masm.framePushed() == 0);
}
@@ -980,7 +1208,7 @@ static void AssertNoWasmExitFPInJitExit(MacroAssembler& masm) {
// JIT exit stub from/to normal wasm code, packedExitFP is not tagged wasm.
#ifdef DEBUG
Register scratch = ABINonArgReturnReg0;
- LoadActivation(masm, scratch);
+ LoadActivation(masm, InstanceReg, scratch);
Label ok;
masm.branchTestPtr(Assembler::Zero,
@@ -1032,8 +1260,7 @@ void wasm::GenerateJitExitEpilogue(MacroAssembler& masm,
// Inverse of GenerateJitExitPrologue:
MOZ_ASSERT(masm.framePushed() == 0);
AssertNoWasmExitFPInJitExit(masm);
- GenerateCallableEpilogue(masm, /*framePushed*/ 0, ExitReason::None(),
- &offsets->ret);
+ GenerateCallableEpilogue(masm, /*framePushed*/ 0, &offsets->ret);
MOZ_ASSERT(masm.framePushed() == 0);
}
diff --git a/js/src/wasm/WasmFrameIter.h b/js/src/wasm/WasmFrameIter.h
@@ -350,14 +350,111 @@ class ProfilingFrameIterator {
// Prologue/epilogue code generation
+void LoadActivation(jit::MacroAssembler& masm, jit::Register instance,
+ jit::Register dest);
void SetExitFP(jit::MacroAssembler& masm, ExitReason reason,
- jit::Register scratch);
-void ClearExitFP(jit::MacroAssembler& masm, jit::Register scratch);
+ jit::Register activation, jit::Register scratch);
+void ClearExitFP(jit::MacroAssembler& masm, jit::Register activation);
-void GenerateExitPrologue(jit::MacroAssembler& masm, unsigned framePushed,
- ExitReason reason, CallableOffsets* offsets);
-void GenerateExitEpilogue(jit::MacroAssembler& masm, unsigned framePushed,
- ExitReason reason, CallableOffsets* offsets);
+#ifdef ENABLE_WASM_JSPI
+// [SMDOC] Wasm dynamic stack switches on 'exit'
+//
+// The SpiderMonkey codebase and embedders, shouldn't run on wasm suspendable
+// stacks. Some code theoretically could work okay on an alternative stack, but
+// we want to be conservative and not assume that. This gives us flexibility to
+// use smaller stacks than the main stack and not worry about stack overflow.
+//
+// To ensure this, all wasm 'exits' from JIT to the VM are instrumented to
+// perform a dynamic check and switch to the main stack if they are currently
+// running on a wasm stack.
+//
+// This is done in the prologue of the exit, and reversed in the epilogue.
+//
+// If we're running on a suspendable stack, we switch SP to the main stack's SP,
+// but keep the FP pointing at the original FP on the incoming stack:
+//
+// Suspendable Stack
+// ┌────────────────┐
+// │ Caller Args │
+// ├────────────────┤
+// │ wasm::Frame │
+// └────────────────┘◄───── FP
+//
+// SP
+// Main Stack │
+// ┌────────────────┐ │
+// │ Previous │ │
+// │ Frames │ │
+// ├────────────────┤ │
+// │ │ │
+// │ framePushed() │ │
+// │ for Exit Stub │ │
+// │ │ │
+// └────────────────┘◄───────┘
+//
+// If we're not running on a suspendable stack, nothing is done at all and
+// SP/FP are unchanged:
+//
+// Main Stack
+// ┌────────────────┐
+// │ Caller Args │
+// ├────────────────┤
+// │ wasm::Frame │
+// ├────────────────┤◄───── FP
+// │ │ SP
+// │ framePushed() │ │
+// │ for exit stub │ │
+// │ │ │
+// └────────────────┘◄───────┘
+//
+// This 'split' function body lets the function still address all the incoming
+// arguments through FP, and it's own 'framePushed' through SP.
+//
+// However this means the SP/FP are no longer guaranteed to be contiguous (they
+// are in the main stack case, but we don't know that statically). So the
+// function body must not access the original frame or incoming arguments
+// through SP, or the 'framePushed' area through FP.
+void GenerateExitPrologueMainStackSwitch(jit::MacroAssembler& masm,
+ jit::Register instance,
+ jit::Register scratch1,
+ jit::Register scratch2,
+ jit::Register scratch3);
+
+// Generate the dynamic switch back to the wasm suspendable stack we originally
+// were on. See "Wasm dynamic stack switches on 'exit'" for more information.
+//
+// NOTE: this doesn't actually switch SP back to the original SP. The caller
+// must do that through some method, such as setting SP := FP.
+void GenerateExitEpilogueMainStackReturn(jit::MacroAssembler& masm,
+ jit::Register instance,
+ jit::Register activationAndScratch1,
+ jit::Register scratch2);
+#endif
+
+// Generate an 'exit' prologue.
+//
+// This will exit the JitActivation, allowing arbitrary code to run. The
+// `reason` will be noted on the JitActivation for any future stack iteration.
+//
+// If `switchToMainStack` is true, the prologue will check if a suspendable
+// stack is active, and if so switch the stack to the main stack.
+//
+// In this case, the body of the exit function will have a 'split' sp/fp where
+// the fp points at the wasm::Frame on the suspendable stack and the sp points
+// to the main stack. See "Wasm dynamic stack switches on 'exit'" above for more
+// information and a diagram.
+//
+// `framePushedPreSwitch` will be reserved on the original stack, and
+// `framePushedPostSwitch` will be reserved on the final stack (either the
+// original stack, or the main stack if there is a switch).
+void GenerateExitPrologue(jit::MacroAssembler& masm, ExitReason reason,
+ bool switchToMainStack, unsigned framePushedPreSwitch,
+ unsigned framePushedPostSwitch,
+ CallableOffsets* offsets);
+// Generate an 'exit' epilogue that is the inverse of
+// wasm::GenerateExitPrologue.
+void GenerateExitEpilogue(jit::MacroAssembler& masm, ExitReason reason,
+ bool switchToMainStack, CallableOffsets* offsets);
// Generate the most minimal possible prologue/epilogue: `push FP; FP := SP`
// and `pop FP; return` respectively.
diff --git a/js/src/wasm/WasmPI.h b/js/src/wasm/WasmPI.h
@@ -268,6 +268,12 @@ class SuspenderObject : public NativeObject {
return getFixedSlot(SuspendableExitFPSlot).toPrivate();
}
+ static constexpr size_t offsetOfState() {
+ return getFixedSlotOffset(StateSlot);
+ }
+ static constexpr size_t offsetOfStackMemory() {
+ return getFixedSlotOffset(StackMemorySlot);
+ }
static constexpr size_t offsetOfMainFP() {
return getFixedSlotOffset(MainFPSlot);
}
diff --git a/js/src/wasm/WasmStubs.cpp b/js/src/wasm/WasmStubs.cpp
@@ -23,12 +23,14 @@
#include "jit/ABIArgGenerator.h"
#include "jit/JitFrames.h"
+#include "jit/JitRuntime.h"
#include "jit/RegisterAllocator.h"
#include "js/Printf.h"
#include "util/Memory.h"
#include "wasm/WasmCode.h"
#include "wasm/WasmGenerator.h"
#include "wasm/WasmInstance.h"
+#include "wasm/WasmPI.h"
#include "jit/MacroAssembler-inl.h"
#include "wasm/WasmInstance-inl.h"
@@ -1997,13 +1999,13 @@ static bool GenerateImportInterpExit(MacroAssembler& masm, const FuncImport& fi,
// The abiArgCount includes a stack result pointer argument if needed.
unsigned abiArgCount = ArgTypeVector(funcType).lengthWithStackResults();
unsigned argBytes = std::max<size_t>(1, abiArgCount) * sizeof(Value);
- unsigned framePushed =
- StackDecrementForCall(ABIStackAlignment,
- sizeof(Frame), // pushed by prologue
- argOffset + argBytes);
-
- GenerateExitPrologue(masm, framePushed, ExitReason::Fixed::ImportInterp,
- offsets);
+ unsigned frameAlignment =
+ ComputeByteAlignment(sizeof(Frame), ABIStackAlignment);
+ unsigned framePushed = AlignBytes(argOffset + argBytes, ABIStackAlignment);
+ GenerateExitPrologue(masm, ExitReason::Fixed::ImportInterp,
+ /*switchToMainStack*/ false,
+ /*framePushedPreSwitch*/ frameAlignment,
+ /*framePushedPostSwitch*/ framePushed, offsets);
// Fill the argument array.
Register scratch = ABINonArgReturnReg0;
@@ -2122,8 +2124,8 @@ static bool GenerateImportInterpExit(MacroAssembler& masm, const FuncImport& fi,
MOZ_ASSERT(NonVolatileRegs.has(HeapReg));
#endif
- GenerateExitEpilogue(masm, framePushed, ExitReason::Fixed::ImportInterp,
- offsets);
+ GenerateExitEpilogue(masm, ExitReason::Fixed::ImportInterp,
+ /*switchToMainStack*/ false, offsets);
return FinishOffsets(masm, offsets);
}
@@ -2412,7 +2414,8 @@ static bool GenerateImportJitExit(MacroAssembler& masm,
// The JIT might have clobbered exitFP at this point. Since there's
// going to be a CoerceInPlace call, pretend we're still doing the JIT
// call by restoring our tagged exitFP.
- SetExitFP(masm, ExitReason::Fixed::ImportJit, scratch);
+ LoadActivation(masm, InstanceReg, scratch);
+ SetExitFP(masm, ExitReason::Fixed::ImportJit, scratch, scratch2);
// argument 0: argv
ABIArgMIRTypeIter i(coerceArgTypes, ABIKind::System);
@@ -2476,6 +2479,7 @@ static bool GenerateImportJitExit(MacroAssembler& masm,
// Maintain the invariant that exitFP is either unset or not set to a
// wasm tagged exitFP, per the jit exit contract.
+ LoadActivation(masm, InstanceReg, scratch);
ClearExitFP(masm, scratch);
masm.jump(&done);
@@ -2513,18 +2517,19 @@ struct ABIFunctionArgs {
};
bool wasm::GenerateBuiltinThunk(MacroAssembler& masm, ABIFunctionType abiType,
- ExitReason exitReason, void* funcPtr,
- CallableOffsets* offsets) {
+ bool switchToMainStack, ExitReason exitReason,
+ void* funcPtr, CallableOffsets* offsets) {
AssertExpectedSP(masm);
masm.setFramePushed(0);
ABIFunctionArgs args(abiType);
- uint32_t framePushed =
- StackDecrementForCall(ABIStackAlignment,
- sizeof(Frame), // pushed by prologue
- StackArgBytesForNativeABI(args));
-
- GenerateExitPrologue(masm, framePushed, exitReason, offsets);
+ unsigned frameAlignment =
+ ComputeByteAlignment(sizeof(Frame), ABIStackAlignment);
+ unsigned framePushed =
+ AlignBytes(StackArgBytesForNativeABI(args), ABIStackAlignment);
+ GenerateExitPrologue(masm, exitReason, switchToMainStack,
+ /*framePushedPreSwitch*/ frameAlignment,
+ /*framePushedPostSwitch*/ framePushed, offsets);
// Copy out and convert caller arguments, if needed. We are translating from
// the wasm ABI to the system ABI.
@@ -2575,7 +2580,7 @@ bool wasm::GenerateBuiltinThunk(MacroAssembler& masm, ABIFunctionType abiType,
MOZ_ASSERT(callArgs.done());
// Call into the native builtin function
- AssertStackAlignment(masm, ABIStackAlignment);
+ masm.assertStackAlignment(ABIStackAlignment);
MoveSPForJitABI(masm);
masm.call(ImmPtr(funcPtr, ImmPtr::NoCheckToken()));
@@ -2603,7 +2608,7 @@ bool wasm::GenerateBuiltinThunk(MacroAssembler& masm, ABIFunctionType abiType,
}
#endif
- GenerateExitEpilogue(masm, framePushed, exitReason, offsets);
+ GenerateExitEpilogue(masm, exitReason, switchToMainStack, offsets);
return FinishOffsets(masm, offsets);
}
@@ -2883,7 +2888,8 @@ static bool GenerateDebugStub(MacroAssembler& masm, Label* throwLabel,
masm.haltingAlign(CodeAlignment);
masm.setFramePushed(0);
- GenerateExitPrologue(masm, 0, ExitReason::Fixed::DebugStub, offsets);
+ GenerateExitPrologue(masm, ExitReason::Fixed::DebugStub,
+ /*switchToMainStack*/ false, 0, 0, offsets);
// Save all registers used between baseline compiler operations.
masm.PushRegsInMask(AllAllocatableRegs);
@@ -2922,7 +2928,8 @@ static bool GenerateDebugStub(MacroAssembler& masm, Label* throwLabel,
masm.setFramePushed(framePushed);
masm.PopRegsInMask(AllAllocatableRegs);
- GenerateExitEpilogue(masm, 0, ExitReason::Fixed::DebugStub, offsets);
+ GenerateExitEpilogue(masm, ExitReason::Fixed::DebugStub,
+ /*switchToMainStack*/ false, offsets);
return FinishOffsets(masm, offsets);
}
@@ -2941,7 +2948,8 @@ static bool GenerateRequestTierUpStub(MacroAssembler& masm,
masm.haltingAlign(CodeAlignment);
masm.setFramePushed(0);
- GenerateExitPrologue(masm, 0, ExitReason::Fixed::RequestTierUp, offsets);
+ GenerateExitPrologue(masm, ExitReason::Fixed::RequestTierUp,
+ /*switchToMainStack*/ false, 0, 0, offsets);
// Save all registers used between baseline compiler operations.
masm.PushRegsInMask(AllAllocatableRegs);
@@ -3006,7 +3014,8 @@ static bool GenerateRequestTierUpStub(MacroAssembler& masm,
masm.setFramePushed(framePushed);
masm.PopRegsInMask(AllAllocatableRegs);
- GenerateExitEpilogue(masm, 0, ExitReason::Fixed::RequestTierUp, offsets);
+ GenerateExitEpilogue(masm, ExitReason::Fixed::RequestTierUp,
+ /*switchToMainStack*/ false, offsets);
return FinishOffsets(masm, offsets);
}
diff --git a/js/src/wasm/WasmStubs.h b/js/src/wasm/WasmStubs.h
@@ -244,8 +244,8 @@ class ABIResultIter {
extern bool GenerateBuiltinThunk(jit::MacroAssembler& masm,
jit::ABIFunctionType abiType,
- ExitReason exitReason, void* funcPtr,
- CallableOffsets* offsets);
+ bool switchToMainStack, ExitReason exitReason,
+ void* funcPtr, CallableOffsets* offsets);
extern bool GenerateStubs(const CodeMetadata& codeMeta,
const FuncImportVector& imports,
