Bug 1404443 - Baldr: avoid ARM push/pop in prologue (r=bbouvier)

MozReview-Commit-ID: 6yI9EtXNJMl

--HG--
extra : rebase_source : d0ff18df5b24d98977eb8d298d24817897fc6e6d

js/src/jit/MacroAssembler.cpp

@@ -3114,20 +3114,6 @@ MacroAssembler::wasmEmitTrapOutOfLineCode()
     trapSites().clear();
 }
 
-void
-MacroAssembler::wasmAssertNonExitInvariants(Register activation)
-{
-#ifdef DEBUG
-    // packedExitFP should not be tagged as wasm. Note this includes nullptr.
-    Label ok;
-    Address packedExitFP(activation, JitActivation::offsetOfPackedExitFP());
-    branchTestPtr(Assembler::Zero, packedExitFP, Imm32(uintptr_t(JitActivation::ExitFpWasmBit)),
-                  &ok);
-    breakpoint();
-    bind(&ok);
-#endif
-}
-
 void
 MacroAssembler::wasmEmitStackCheck(Register sp, Register scratch, Label* onOverflow)
 {

js/src/jit/MacroAssembler.h

@@ -1503,9 +1503,6 @@ class MacroAssembler : public MacroAssemblerSpecific
     // including "normal" OutOfLineCode.
     void wasmEmitTrapOutOfLineCode();
 
-    // Assert invariants that should be true within any non-exit-stub wasm code.
-    void wasmAssertNonExitInvariants(Register activation);
-
     // Perform a stack-overflow test, branching to the given Label on overflow.
     void wasmEmitStackCheck(Register sp, Register scratch, Label* onOverflow);
 

js/src/jit/arm/Assembler-arm.h

@@ -109,12 +109,20 @@ class ABIArgGenerator
 
 bool IsUnaligned(const wasm::MemoryAccessDesc& access);
 
+// These registers may be volatile or nonvolatile.
 static constexpr Register ABINonArgReg0 = r4;
 static constexpr Register ABINonArgReg1 = r5;
 static constexpr Register ABINonArgReg2 = r6;
+
+// These registers may be volatile or nonvolatile.
+// Note: these three registers are all guaranteed to be different
 static constexpr Register ABINonArgReturnReg0 = r4;
 static constexpr Register ABINonArgReturnReg1 = r5;
 
+// This register is guaranteed to be clobberable during the prologue of an ABI
+// call which must preserve both ABI argument and non-volatile registers.
+static constexpr Register NativeABIPrologueClobberable = ScratchRegister;
+
 // TLS pointer argument register for WebAssembly functions. This must not alias
 // any other register used for passing function arguments or return values.
 // Preserved by WebAssembly functions.

js/src/jit/arm64/Assembler-arm64.h

@@ -443,12 +443,20 @@ class ABIArgGenerator
     ABIArg current_;
 };
 
+// These registers may be volatile or nonvolatile.
 static constexpr Register ABINonArgReg0 = r8;
 static constexpr Register ABINonArgReg1 = r9;
 static constexpr Register ABINonArgReg2 = r10;
+
+// These registers may be volatile or nonvolatile.
+// Note: these three registers are all guaranteed to be different
 static constexpr Register ABINonArgReturnReg0 = r8;
 static constexpr Register ABINonArgReturnReg1 = r9;
 
+// This register is guaranteed to be clobberable during the prologue of an ABI
+// call which must preserve both ABI argument and non-volatile registers.
+static constexpr Register NativeABIPrologueClobberable = lr;
+
 // TLS pointer argument register for WebAssembly functions. This must not alias
 // any other register used for passing function arguments or return values.
 // Preserved by WebAssembly functions.

js/src/jit/mips32/Assembler-mips32.h

@@ -46,12 +46,20 @@ class ABIArgGenerator
     }
 };
 
+// These registers may be volatile or nonvolatile.
 static constexpr Register ABINonArgReg0 = t0;
 static constexpr Register ABINonArgReg1 = t1;
 static constexpr Register ABINonArgReg2 = t2;
+
+// These registers may be volatile or nonvolatile.
+// Note: these three registers are all guaranteed to be different
 static constexpr Register ABINonArgReturnReg0 = t0;
 static constexpr Register ABINonArgReturnReg1 = t1;
 
+// This register is guaranteed to be clobberable during the prologue of an ABI
+// call which must preserve both ABI argument and non-volatile registers.
+static constexpr Register NativeABIPrologueClobberable = t0;
+
 // TLS pointer argument register for WebAssembly functions. This must not alias
 // any other register used for passing function arguments or return values.
 // Preserved by WebAssembly functions.

js/src/jit/mips64/Assembler-mips64.h

@@ -39,12 +39,20 @@ class ABIArgGenerator
     }
 };
 
+// These registers may be volatile or nonvolatile.
 static constexpr Register ABINonArgReg0 = t0;
 static constexpr Register ABINonArgReg1 = t1;
 static constexpr Register ABINonArgReg2 = t2;
+
+// These registers may be volatile or nonvolatile.
+// Note: these three registers are all guaranteed to be different
 static constexpr Register ABINonArgReturnReg0 = t0;
 static constexpr Register ABINonArgReturnReg1 = t1;
 
+// This register is guaranteed to be clobberable during the prologue of an ABI
+// call which must preserve both ABI argument and non-volatile registers.
+static constexpr Register NativeABIPrologueClobberable = t0;
+
 // TLS pointer argument register for WebAssembly functions. This must not alias
 // any other register used for passing function arguments or return values.
 // Preserved by WebAssembly functions.

js/src/jit/x64/Assembler-x64.h

@@ -190,16 +190,22 @@ class ABIArgGenerator
     uint32_t stackBytesConsumedSoFar() const { return stackOffset_; }
 };
 
+// These registers may be volatile or nonvolatile.
 // Avoid r11, which is the MacroAssembler's ScratchReg.
 static constexpr Register ABINonArgReg0 = rax;
 static constexpr Register ABINonArgReg1 = rbx;
 static constexpr Register ABINonArgReg2 = r10;
 
+// These registers may be volatile or nonvolatile.
 // Note: these three registers are all guaranteed to be different
 static constexpr Register ABINonArgReturnReg0 = r10;
 static constexpr Register ABINonArgReturnReg1 = r12;
 static constexpr Register ABINonVolatileReg = r13;
 
+// This register is guaranteed to be clobberable during the prologue of an ABI
+// call which must preserve both ABI argument and non-volatile registers.
+static constexpr Register NativeABIPrologueClobberable = rax;
+
 // TLS pointer argument register for WebAssembly functions. This must not alias
 // any other register used for passing function arguments or return values.
 // Preserved by WebAssembly functions.

js/src/jit/x86/Assembler-x86.h

@@ -81,15 +81,21 @@ class ABIArgGenerator
 
 };
 
+// These registers may be volatile or nonvolatile.
 static constexpr Register ABINonArgReg0 = eax;
 static constexpr Register ABINonArgReg1 = ebx;
 static constexpr Register ABINonArgReg2 = ecx;
 
+// These registers may be volatile or nonvolatile.
 // Note: these three registers are all guaranteed to be different
 static constexpr Register ABINonArgReturnReg0 = ecx;
 static constexpr Register ABINonArgReturnReg1 = edx;
 static constexpr Register ABINonVolatileReg = ebx;
 
+// This register is guaranteed to be clobberable during the prologue of an ABI
+// call which must preserve both ABI argument and non-volatile registers.
+static constexpr Register NativeABIPrologueClobberable = eax;
+
 // TLS pointer argument register for WebAssembly functions. This must not alias
 // any other register used for passing function arguments or return values.
 // Preserved by WebAssembly functions.

js/src/wasm/WasmFrameIter.cpp

@@ -383,22 +383,12 @@ GenerateCallablePrologue(MacroAssembler& masm, unsigned framePushed, ExitReason
             *tierEntry = masm.currentOffset();
     }
 
-    // To make the jit exit faster, we don't set exitFP although we set the
-    // exit reason: the profiling iteration knows how to recover FP.
-    if (!reason.isNone() && !(reason.isFixed() && reason.fixed() == ExitReason::Fixed::ImportJit)) {
-        Register act = ABINonArgReg0;
-
-        // Native callers expect the native ABI, which assume that non-saved
-        // registers are preserved. Explicitly preserve the act register
-        // in that case.
-        if (reason.isNative() && !act.volatile_())
-            masm.Push(act);
-
-        SetExitFP(masm, act);
-
-        if (reason.isNative() && !act.volatile_())
-            masm.Pop(act);
-    }
+    // Set exitFP to tell the wasm frame iterators to know where to start
+    // unwinding. In the special case of the optimized import-JIT stub, we're
+    // just calling into more JIT code which can be unwound, so exitFP isn't
+    // needed.
+    if (!reason.isNone() && !(reason.isFixed() && reason.fixed() == ExitReason::Fixed::ImportJit))
+        SetExitFP(masm, NativeABIPrologueClobberable);
 
     if (framePushed)
         masm.subFromStackPtr(Imm32(framePushed));
@@ -534,12 +524,34 @@ wasm::GenerateExitEpilogue(MacroAssembler& masm, unsigned framePushed, ExitReaso
     masm.setFramePushed(0);
 }
 
+static void
+AssertNoWasmExitFPInJitExit(MacroAssembler& masm)
+{
+    // As a general stack invariant, if Activation::packedExitFP is tagged as
+    // wasm, it must point to a valid wasm::Frame. The JIT exit stub calls into
+    // JIT code and thus does not really exit, thus, when entering/leaving the
+    // JIT exit stub from/to normal wasm code, packedExitFP is not tagged wasm.
+#ifdef DEBUG
+    Register scratch = ABINonArgReturnReg0;
+    LoadActivation(masm, scratch);
+
+    Label ok;
+    masm.branchTestPtr(Assembler::Zero,
+                       Address(scratch, JitActivation::offsetOfPackedExitFP()),
+                       Imm32(uintptr_t(JitActivation::ExitFpWasmBit)),
+                       &ok);
+    masm.breakpoint();
+    masm.bind(&ok);
+#endif
+}
+
 void
 wasm::GenerateJitExitPrologue(MacroAssembler& masm, unsigned framePushed, CallableOffsets* offsets)
 {
     masm.haltingAlign(CodeAlignment);
     GenerateCallablePrologue(masm, framePushed, ExitReason(ExitReason::Fixed::ImportJit),
                              &offsets->begin, nullptr, CompileMode::Once, 0);
+    AssertNoWasmExitFPInJitExit(masm);
     masm.setFramePushed(framePushed);
 }
 
@@ -548,13 +560,7 @@ wasm::GenerateJitExitEpilogue(MacroAssembler& masm, unsigned framePushed, Callab
 {
     // Inverse of GenerateJitExitPrologue:
     MOZ_ASSERT(masm.framePushed() == framePushed);
-
-#ifdef DEBUG
-    Register scratch = ABINonArgReturnReg0;
-    LoadActivation(masm, scratch);
-    masm.wasmAssertNonExitInvariants(scratch);
-#endif
-
+    AssertNoWasmExitFPInJitExit(masm);
     GenerateCallableEpilogue(masm, framePushed, ExitReason::None(), &offsets->ret);
     masm.setFramePushed(0);
 }