Bug 1599722 - Make Ion address stack arguments from FP r=lth,wingo

Make the WebAssembly ion compiler address incoming stack arguments from the frame pointer instead of from the stack pointer.

We have a plan to allow interstitial trampoline frames to be inserted between callers and callees, but only for cross-instance calls. This will mean offset to incoming stack arguments from SP is no longer a constant. The design has us instead address stack arguments from the frame pointer, as adding an interstitial trampoline frame won't modify the frame pointer.

Differential Revision: https://phabricator.services.mozilla.com/D73030

js/src/jit/CodeGenerator.cpp

@@ -10796,6 +10796,7 @@ bool CodeGenerator::generateWasm(wasm::FuncTypeIdDesc funcTypeId,
                                  wasm::FuncOffsets* offsets,
                                  wasm::StackMaps* stackMaps) {
   JitSpew(JitSpew_Codegen, "# Emitting wasm code");
+  setUseWasmStackArgumentAbi();
 
   size_t nInboundStackArgBytes = StackArgAreaSizeUnaligned(argTypes);
 

js/src/jit/MoveResolver.h

@@ -11,6 +11,7 @@
 #include "jit/JitAllocPolicy.h"
 #include "jit/Registers.h"
 #include "jit/RegisterSets.h"
+#include "jit/shared/Assembler-shared.h"
 
 namespace js {
 namespace jit {
@@ -58,6 +59,8 @@ class MoveOperand {
       kind_ = REG;
     }
   }
+  MoveOperand(const Address& addr, Kind kind = MEMORY)
+      : MoveOperand(AsRegister(addr.base), addr.offset, kind) {}
   MoveOperand(MacroAssembler& masm, const ABIArg& arg);
   MoveOperand(const MoveOperand& other) = default;
   bool isFloatReg() const { return kind_ == FLOAT_REG; }

js/src/jit/arm/CodeGenerator-arm.cpp

@@ -1065,11 +1065,11 @@ MoveOperand CodeGeneratorARM::toMoveOperand(LAllocation a) const {
   if (a.isFloatReg()) {
     return MoveOperand(ToFloatRegister(a));
   }
-  int32_t offset = ToStackOffset(a);
-  MOZ_ASSERT((offset & 3) == 0);
   MoveOperand::Kind kind =
       a.isStackArea() ? MoveOperand::EFFECTIVE_ADDRESS : MoveOperand::MEMORY;
-  return MoveOperand(StackPointer, offset, kind);
+  Address addr = ToAddress(a);
+  MOZ_ASSERT((addr.offset & 3) == 0);
+  return MoveOperand(addr, kind);
 }
 
 class js::jit::OutOfLineTableSwitch

js/src/jit/arm64/CodeGenerator-arm64.cpp

@@ -1045,9 +1045,7 @@ MoveOperand CodeGeneratorARM64::toMoveOperand(const LAllocation a) const {
   }
   MoveOperand::Kind kind =
       a.isStackArea() ? MoveOperand::EFFECTIVE_ADDRESS : MoveOperand::MEMORY;
-
-  return MoveOperand(AsRegister(masm.getStackPointer()), ToStackOffset(a),
-                     kind);
+  return MoveOperand(ToAddress(a), kind);
 }
 
 class js::jit::OutOfLineTableSwitch

js/src/jit/mips-shared/CodeGenerator-mips-shared.cpp

@@ -48,7 +48,7 @@ Operand CodeGeneratorMIPSShared::ToOperand(const LAllocation& a) {
   if (a.isFloatReg()) {
     return Operand(a.toFloatReg()->reg());
   }
-  return Operand(masm.getStackPointer(), ToStackOffset(&a));
+  return Operand(ToAddress(a));
 }
 
 Operand CodeGeneratorMIPSShared::ToOperand(const LAllocation* a) {
@@ -1091,12 +1091,11 @@ MoveOperand CodeGeneratorMIPSShared::toMoveOperand(LAllocation a) const {
   if (a.isFloatReg()) {
     return MoveOperand(ToFloatRegister(a));
   }
-  int32_t offset = ToStackOffset(a);
-  MOZ_ASSERT((offset & 3) == 0);
   MoveOperand::Kind kind =
       a.isStackArea() ? MoveOperand::EFFECTIVE_ADDRESS : MoveOperand::MEMORY;
-
-  return MoveOperand(StackPointer, offset, kind);
+  Address address = ToAddress(a);
+  MOZ_ASSERT((address.offset & 3) == 0);
+  return MoveOperand(address, kind);
 }
 
 void CodeGenerator::visitMathD(LMathD* math) {

js/src/jit/shared/CodeGenerator-shared-inl.h

@@ -242,15 +242,28 @@ int32_t CodeGeneratorShared::ToStackOffset(const LAllocation* a) const {
   return ToStackOffset(*a);
 }
 
-Address CodeGeneratorShared::ToAddress(const LAllocation& a) {
-  MOZ_ASSERT(a.isMemory());
+Address CodeGeneratorShared::ToAddress(const LAllocation& a) const {
+  MOZ_ASSERT(a.isMemory() || a.isStackArea());
+  if (useWasmStackArgumentAbi() && a.isArgument()) {
+    return Address(FramePointer, ToFramePointerOffset(a));
+  }
   return Address(masm.getStackPointer(), ToStackOffset(&a));
 }
 
-Address CodeGeneratorShared::ToAddress(const LAllocation* a) {
+Address CodeGeneratorShared::ToAddress(const LAllocation* a) const {
   return ToAddress(*a);
 }
 
+int32_t CodeGeneratorShared::ToFramePointerOffset(LAllocation a) const {
+  MOZ_ASSERT(useWasmStackArgumentAbi());
+  MOZ_ASSERT(a.isArgument());
+  return a.toArgument()->index() + sizeof(wasm::Frame);
+}
+
+int32_t CodeGeneratorShared::ToFramePointerOffset(const LAllocation* a) const {
+  return ToFramePointerOffset(*a);
+}
+
 void CodeGeneratorShared::saveLive(LInstruction* ins) {
   MOZ_ASSERT(!ins->isCall());
   LSafepoint* safepoint = ins->safepoint();

js/src/jit/shared/CodeGenerator-shared.cpp

@@ -45,6 +45,7 @@ MacroAssembler& CodeGeneratorShared::ensureMasm(MacroAssembler* masmArg) {
 CodeGeneratorShared::CodeGeneratorShared(MIRGenerator* gen, LIRGraph* graph,
                                          MacroAssembler* masmArg)
     : maybeMasm_(),
+      useWasmStackArgumentAbi_(false),
       masm(ensureMasm(masmArg)),
       gen(gen),
       graph(*graph),

js/src/jit/shared/CodeGenerator-shared.h

@@ -42,6 +42,8 @@ class CodeGeneratorShared : public LElementVisitor {
   MacroAssembler& ensureMasm(MacroAssembler* masm);
   mozilla::Maybe<IonHeapMacroAssembler> maybeMasm_;
 
+  bool useWasmStackArgumentAbi_;
+
  public:
   MacroAssembler& masm;
 
@@ -167,8 +169,13 @@ class CodeGeneratorShared : public LElementVisitor {
   inline int32_t ToStackOffset(LAllocation a) const;
   inline int32_t ToStackOffset(const LAllocation* a) const;
 
-  inline Address ToAddress(const LAllocation& a);
-  inline Address ToAddress(const LAllocation* a);
+  inline Address ToAddress(const LAllocation& a) const;
+  inline Address ToAddress(const LAllocation* a) const;
+
+  // Returns the offset from FP to address incoming stack arguments
+  // when we use wasm stack argument abi (useWasmStackArgumentAbi()).
+  inline int32_t ToFramePointerOffset(LAllocation a) const;
+  inline int32_t ToFramePointerOffset(const LAllocation* a) const;
 
   uint32_t frameSize() const {
     return frameClass_ == FrameSizeClass::None() ? frameDepth_
@@ -180,6 +187,10 @@ class CodeGeneratorShared : public LElementVisitor {
   void dumpNativeToBytecodeEntries();
   void dumpNativeToBytecodeEntry(uint32_t idx);
 
+  void setUseWasmStackArgumentAbi() { useWasmStackArgumentAbi_ = true; }
+
+  bool useWasmStackArgumentAbi() const { return useWasmStackArgumentAbi_; }
+
  public:
   MIRGenerator& mirGen() const { return *gen; }
 

js/src/jit/x64/CodeGenerator-x64.cpp

@@ -38,7 +38,7 @@ Operand CodeGeneratorX64::ToOperand64(const LInt64Allocation& a64) {
   if (a.isGeneralReg()) {
     return Operand(a.toGeneralReg()->reg());
   }
-  return Operand(masm.getStackPointer(), ToStackOffset(a));
+  return Operand(ToAddress(a));
 }
 
 FrameSizeClass FrameSizeClass::FromDepth(uint32_t frameDepth) {

js/src/jit/x86-shared/CodeGenerator-x86-shared.cpp

@@ -1808,7 +1808,7 @@ Operand CodeGeneratorX86Shared::ToOperand(const LAllocation& a) {
   if (a.isFloatReg()) {
     return Operand(a.toFloatReg()->reg());
   }
-  return Operand(masm.getStackPointer(), ToStackOffset(&a));
+  return Operand(ToAddress(a));
 }
 
 Operand CodeGeneratorX86Shared::ToOperand(const LAllocation* a) {
@@ -1828,7 +1828,7 @@ MoveOperand CodeGeneratorX86Shared::toMoveOperand(LAllocation a) const {
   }
   MoveOperand::Kind kind =
       a.isStackArea() ? MoveOperand::EFFECTIVE_ADDRESS : MoveOperand::MEMORY;
-  return MoveOperand(StackPointer, ToStackOffset(a), kind);
+  return MoveOperand(ToAddress(a), kind);
 }
 
 class OutOfLineTableSwitch : public OutOfLineCodeBase<CodeGeneratorX86Shared> {