Bug 1661861: Generate arguments rectifier for trial inlining r=jandem

We don't have to store a separate return offset for trial inlining, because the rectifier code after returning from the call is exactly the same. Post-bailout, inline frames will return into the normal non-inline code.

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

js/src/jit/BaselineCacheIRCompiler.cpp

@@ -3261,21 +3261,7 @@ bool BaselineCacheIRCompiler::emitCallInlinedFunction(ObjOperandId calleeId,
     return false;
   }
 
-  // Load JitScript
-  masm.loadPtr(Address(calleeReg, JSFunction::offsetOfScript()), codeReg);
-  masm.branchIfScriptHasNoJitScript(codeReg, failure->label());
-
-  masm.loadJitScript(codeReg, codeReg);
-
-  // Load BaselineScript
-  masm.loadPtr(Address(codeReg, JitScript::offsetOfBaselineScript()), codeReg);
-  static_assert(BaselineDisabledScript == 0x1);
-  masm.branchPtr(Assembler::BelowOrEqual, codeReg,
-                 ImmWord(BaselineDisabledScript), failure->label());
-
-  // Load Baseline jitcode
-  masm.loadPtr(Address(codeReg, BaselineScript::offsetOfMethod()), codeReg);
-  masm.loadPtr(Address(codeReg, JitCode::offsetOfCode()), codeReg);
+  masm.loadBaselineJitCodeRaw(calleeReg, codeReg, failure->label());
 
   if (!updateArgc(flags, argcReg, scratch)) {
     return false;
@@ -3322,7 +3308,12 @@ bool BaselineCacheIRCompiler::emitCallInlinedFunction(ObjOperandId calleeId,
   masm.load16ZeroExtend(Address(calleeReg, JSFunction::offsetOfNargs()),
                         calleeReg);
   masm.branch32(Assembler::AboveOrEqual, argcReg, calleeReg, &noUnderflow);
-  masm.assumeUnreachable("Arguments rectifier not yet supported.");
+
+  // Call the trial-inlining arguments rectifier.
+  ArgumentsRectifierKind kind = ArgumentsRectifierKind::TrialInlining;
+  TrampolinePtr argumentsRectifier =
+      cx_->runtime()->jitRuntime()->getArgumentsRectifier(kind);
+  masm.movePtr(argumentsRectifier, codeReg);
 
   masm.bind(&noUnderflow);
   masm.callJit(codeReg);

js/src/jit/Ion.cpp

@@ -174,8 +174,11 @@ bool JitRuntime::generateTrampolines(JSContext* cx) {
   static_assert(sizeof(JitFrameLayout) == sizeof(WasmToJSJitFrameLayout),
                 "thus a rectifier frame can be used with a wasm frame");
 
-  JitSpew(JitSpew_Codegen, "# Emitting sequential arguments rectifier");
-  generateArgumentsRectifier(masm);
+  JitSpew(JitSpew_Codegen, "# Emitting arguments rectifier");
+  generateArgumentsRectifier(masm, ArgumentsRectifierKind::Normal);
+
+  JitSpew(JitSpew_Codegen, "# Emitting trial inlining arguments rectifier");
+  generateArgumentsRectifier(masm, ArgumentsRectifierKind::TrialInlining);
 
   JitSpew(JitSpew_Codegen, "# Emitting EnterJIT sequence");
   generateEnterJIT(cx, masm);

js/src/jit/JitRealm.h

@@ -118,6 +118,8 @@ class BaselineICFallbackCode {
   }
 };
 
+enum class ArgumentsRectifierKind { Normal, TrialInlining };
+
 enum class DebugTrapHandlerKind { Interpreter, Compiler, Count };
 
 using EnterJitCode = void (*)(void*, unsigned int, Value*, InterpreterFrame*,
@@ -161,9 +163,11 @@ class JitRuntime {
   // Generic bailout table; used if the bailout table overflows.
   WriteOnceData<uint32_t> bailoutHandlerOffset_{0};
 
-  // Argument-rectifying thunk, in the case of insufficient arguments passed
-  // to a function call site.
+  // Argument-rectifying thunks, in the case of insufficient arguments passed
+  // to a function call site. The return offset is used to rebuild stack frames
+  // when bailing out.
   WriteOnceData<uint32_t> argumentsRectifierOffset_{0};
+  WriteOnceData<uint32_t> trialInliningArgumentsRectifierOffset_{0};
   WriteOnceData<uint32_t> argumentsRectifierReturnOffset_{0};
 
   // Thunk that invalides an (Ion compiled) caller on the Ion stack.
@@ -258,7 +262,8 @@ class JitRuntime {
                                  Label* profilerExitTail);
   void generateBailoutTailStub(MacroAssembler& masm, Label* bailoutTail);
   void generateEnterJIT(JSContext* cx, MacroAssembler& masm);
-  void generateArgumentsRectifier(MacroAssembler& masm);
+  void generateArgumentsRectifier(MacroAssembler& masm,
+                                  ArgumentsRectifierKind kind);
   BailoutTable generateBailoutTable(MacroAssembler& masm, Label* bailoutTail,
                                     uint32_t frameClass);
   void generateBailoutHandler(MacroAssembler& masm, Label* bailoutTail);
@@ -355,7 +360,11 @@ class JitRuntime {
   TrampolinePtr getBailoutTable(const FrameSizeClass& frameClass) const;
   uint32_t getBailoutTableSize(const FrameSizeClass& frameClass) const;
 
-  TrampolinePtr getArgumentsRectifier() const {
+  TrampolinePtr getArgumentsRectifier(
+      ArgumentsRectifierKind kind = ArgumentsRectifierKind::Normal) const {
+    if (kind == ArgumentsRectifierKind::TrialInlining) {
+      return trampolineCode(trialInliningArgumentsRectifierOffset_);
+    }
     return trampolineCode(argumentsRectifierOffset_);
   }
 

js/src/jit/MacroAssembler.cpp

@@ -2113,6 +2113,28 @@ void MacroAssembler::loadJitCodeNoArgCheck(Register func, Register dest) {
   loadPtr(Address(dest, JitScript::offsetOfJitCodeSkipArgCheck()), dest);
 }
 
+void MacroAssembler::loadBaselineJitCodeRaw(Register func, Register dest,
+                                            Label* failure) {
+  // Load JitScript
+  loadPtr(Address(func, JSFunction::offsetOfScript()), dest);
+  if (failure) {
+    branchIfScriptHasNoJitScript(dest, failure);
+  }
+  loadJitScript(dest, dest);
+
+  // Load BaselineScript
+  loadPtr(Address(dest, JitScript::offsetOfBaselineScript()), dest);
+  if (failure) {
+    static_assert(BaselineDisabledScript == 0x1);
+    branchPtr(Assembler::BelowOrEqual, dest, ImmWord(BaselineDisabledScript),
+              failure);
+  }
+
+  // Load Baseline jitcode
+  loadPtr(Address(dest, BaselineScript::offsetOfMethod()), dest);
+  loadPtr(Address(dest, JitCode::offsetOfCode()), dest);
+}
+
 void MacroAssembler::loadBaselineFramePtr(Register framePtr, Register dest) {
   if (framePtr != dest) {
     movePtr(framePtr, dest);

js/src/jit/MacroAssembler.h

@@ -3871,6 +3871,8 @@ class MacroAssembler : public MacroAssemblerSpecific {
  public:
   void loadJitCodeRaw(Register func, Register dest);
   void loadJitCodeNoArgCheck(Register func, Register dest);
+  void loadBaselineJitCodeRaw(Register func, Register dest,
+                              Label* failure = nullptr);
 
   void loadBaselineFramePtr(Register framePtr, Register dest);
 

js/src/jit/TrialInlining.cpp

@@ -300,13 +300,6 @@ bool TrialInliner::maybeInlineCall(const ICEntry& entry, BytecodeLocation loc) {
   MOZ_ASSERT_IF(loc.getOp() == JSOp::FunCall,
                 data->callFlags.getArgFormat() == CallFlags::FunCall);
 
-  // TODO: The arguments rectifier is not yet supported.
-  uint32_t argc =
-      loc.getOp() == JSOp::FunCall ? loc.getCallArgc() - 1 : loc.getCallArgc();
-  if (argc < data->target->nargs()) {
-    return true;
-  }
-
   ICScript* newICScript = createInlinedICScript(data->target, loc);
   if (!newICScript) {
     return false;

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

@@ -436,8 +436,16 @@ void JitRuntime::generateInvalidator(MacroAssembler& masm, Label* bailoutTail) {
   masm.jump(bailoutTail);
 }
 
-void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm) {
-  argumentsRectifierOffset_ = startTrampolineCode(masm);
+void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm,
+                                            ArgumentsRectifierKind kind) {
+  switch (kind) {
+    case ArgumentsRectifierKind::Normal:
+      argumentsRectifierOffset_ = startTrampolineCode(masm);
+      break;
+    case ArgumentsRectifierKind::TrialInlining:
+      trialInliningArgumentsRectifierOffset_ = startTrampolineCode(masm);
+      break;
+  }
   masm.pushReturnAddress();
 
   // Copy number of actual arguments into r0 and r8.
@@ -522,8 +530,16 @@ void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm) {
 
   // Call the target function.
   masm.andPtr(Imm32(CalleeTokenMask), r1);
-  masm.loadJitCodeRaw(r1, r3);
-  argumentsRectifierReturnOffset_ = masm.callJitNoProfiler(r3);
+  switch (kind) {
+    case ArgumentsRectifierKind::Normal:
+      masm.loadJitCodeRaw(r1, r3);
+      argumentsRectifierReturnOffset_ = masm.callJitNoProfiler(r3);
+      break;
+    case ArgumentsRectifierKind::TrialInlining:
+      masm.loadBaselineJitCodeRaw(r1, r3);
+      masm.callJitNoProfiler(r3);
+      break;
+  }
 
   // arg1
   //  ...

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

@@ -360,8 +360,16 @@ void JitRuntime::generateInvalidator(MacroAssembler& masm, Label* bailoutTail) {
   masm.jump(bailoutTail);
 }
 
-void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm) {
-  argumentsRectifierOffset_ = startTrampolineCode(masm);
+void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm,
+                                            ArgumentsRectifierKind kind) {
+  switch (kind) {
+    case ArgumentsRectifierKind::Normal:
+      argumentsRectifierOffset_ = startTrampolineCode(masm);
+      break;
+    case ArgumentsRectifierKind::TrialInlining:
+      trialInliningArgumentsRectifierOffset_ = startTrampolineCode(masm);
+      break;
+  }
 
   // Save the return address for later.
   masm.push(lr);
@@ -449,9 +457,17 @@ void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm) {
             r1,   // Callee token.
             r6);  // Frame descriptor.
 
-  // Load the address of the code that is getting called.
-  masm.loadJitCodeRaw(r5, r3);
-  argumentsRectifierReturnOffset_ = masm.callJitNoProfiler(r3);
+  // Call the target function.
+  switch (kind) {
+    case ArgumentsRectifierKind::Normal:
+      masm.loadJitCodeRaw(r5, r3);
+      argumentsRectifierReturnOffset_ = masm.callJitNoProfiler(r3);
+      break;
+    case ArgumentsRectifierKind::TrialInlining:
+      masm.loadBaselineJitCodeRaw(r5, r3);
+      masm.callJitNoProfiler(r3);
+      break;
+  }
 
   // Clean up!
   // Get the size of the stack frame, and clean up the later fixed frame.

js/src/jit/mips32/Trampoline-mips32.cpp

@@ -403,8 +403,16 @@ void JitRuntime::generateInvalidator(MacroAssembler& masm, Label* bailoutTail) {
   masm.jump(bailoutTail);
 }
 
-void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm) {
-  argumentsRectifierOffset_ = startTrampolineCode(masm);
+void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm,
+                                            ArgumentsRectifierKind kind) {
+  switch (kind) {
+    case ArgumentsRectifierKind::Normal:
+      argumentsRectifierOffset_ = startTrampolineCode(masm);
+      break;
+    case ArgumentsRectifierKind::TrialInlining:
+      trialInliningArgumentsRectifierOffset_ = startTrampolineCode(masm);
+      break;
+  }
   masm.pushReturnAddress();
 
   Register numActArgsReg = t6;
@@ -511,8 +519,16 @@ void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm) {
 
   // Call the target function.
   masm.andPtr(Imm32(CalleeTokenMask), calleeTokenReg);
-  masm.loadJitCodeRaw(calleeTokenReg, t1);
-  argumentsRectifierReturnOffset_ = masm.callJitNoProfiler(t1);
+  switch (kind) {
+    case ArgumentsRectifierKind::Normal:
+      masm.loadJitCodeRaw(calleeTokenReg, t1);
+      argumentsRectifierReturnOffset_ = masm.callJitNoProfiler(t1);
+      break;
+    case ArgumentsRectifierKind::TrialInlining:
+      masm.loadBaselineJitCodeRaw(calleeTokenReg, t1);
+      masm.callJitNoProfiler(t1);
+      break;
+  }
 
   // arg1
   //  ...

js/src/jit/mips64/Trampoline-mips64.cpp

@@ -409,11 +409,20 @@ void JitRuntime::generateInvalidator(MacroAssembler& masm, Label* bailoutTail) {
   masm.jump(bailoutTail);
 }
 
-void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm) {
+void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm,
+                                            ArgumentsRectifierKind kind) {
   // Do not erase the frame pointer in this function.
 
-  argumentsRectifierOffset_ = startTrampolineCode(masm);
+  switch (kind) {
+    case ArgumentsRectifierKind::Normal:
+      argumentsRectifierOffset_ = startTrampolineCode(masm);
+      break;
+    case ArgumentsRectifierKind::TrialInlining:
+      trialInliningArgumentsRectifierOffset_ = startTrampolineCode(masm);
+      break;
+  }
   masm.pushReturnAddress();
+
   // Caller:
   // [arg2] [arg1] [this] [[argc] [callee] [descr] [raddr]] <- sp
 
@@ -565,8 +574,16 @@ void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm) {
 
   // Call the target function.
   masm.andPtr(Imm32(uint32_t(CalleeTokenMask)), calleeTokenReg);
-  masm.loadJitCodeRaw(calleeTokenReg, t1);
-  argumentsRectifierReturnOffset_ = masm.callJitNoProfiler(t1);
+  switch (kind) {
+    case ArgumentsRectifierKind::Normal:
+      masm.loadJitCodeRaw(calleeTokenReg, t1);
+      argumentsRectifierReturnOffset_ = masm.callJitNoProfiler(t1);
+      break;
+    case ArgumentsRectifierKind::TrialInlining:
+      masm.loadBaselineJitCodeRaw(calleeTokenReg, t1);
+      masm.callJitNoProfiler(t1);
+      break;
+  }
 
   // Remove the rectifier frame.
   // t2 <- descriptor with FrameType.

js/src/jit/none/Trampoline-none.cpp

@@ -17,7 +17,10 @@ using namespace js::jit;
 
 void JitRuntime::generateEnterJIT(JSContext*, MacroAssembler&) { MOZ_CRASH(); }
 void JitRuntime::generateInvalidator(MacroAssembler&, Label*) { MOZ_CRASH(); }
-void JitRuntime::generateArgumentsRectifier(MacroAssembler&) { MOZ_CRASH(); }
+void JitRuntime::generateArgumentsRectifier(MacroAssembler&,
+                                            ArgumentsRectifierKind kind) {
+  MOZ_CRASH();
+}
 JitRuntime::BailoutTable JitRuntime::generateBailoutTable(MacroAssembler&,
                                                           Label*, uint32_t) {
   MOZ_CRASH();

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

@@ -390,10 +390,18 @@ void JitRuntime::generateInvalidator(MacroAssembler& masm, Label* bailoutTail) {
   masm.jmp(bailoutTail);
 }
 
-void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm) {
+void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm,
+                                            ArgumentsRectifierKind kind) {
   // Do not erase the frame pointer in this function.
 
-  argumentsRectifierOffset_ = startTrampolineCode(masm);
+  switch (kind) {
+    case ArgumentsRectifierKind::Normal:
+      argumentsRectifierOffset_ = startTrampolineCode(masm);
+      break;
+    case ArgumentsRectifierKind::TrialInlining:
+      trialInliningArgumentsRectifierOffset_ = startTrampolineCode(masm);
+      break;
+  }
 
   // Caller:
   // [arg2] [arg1] [this] [[argc] [callee] [descr] [raddr]] <- rsp
@@ -529,8 +537,16 @@ void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm) {
 
   // Call the target function.
   masm.andq(Imm32(uint32_t(CalleeTokenMask)), rax);
-  masm.loadJitCodeRaw(rax, rax);
-  argumentsRectifierReturnOffset_ = masm.callJitNoProfiler(rax);
+  switch (kind) {
+    case ArgumentsRectifierKind::Normal:
+      masm.loadJitCodeRaw(rax, rax);
+      argumentsRectifierReturnOffset_ = masm.callJitNoProfiler(rax);
+      break;
+    case ArgumentsRectifierKind::TrialInlining:
+      masm.loadBaselineJitCodeRaw(rax, rax);
+      masm.callJitNoProfiler(rax);
+      break;
+  }
 
   // Remove the rectifier frame.
   masm.pop(r9);  // r9 <- descriptor with FrameType.

js/src/jit/x86/Trampoline-x86.cpp

@@ -379,8 +379,16 @@ void JitRuntime::generateInvalidator(MacroAssembler& masm, Label* bailoutTail) {
   masm.jmp(bailoutTail);
 }
 
-void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm) {
-  argumentsRectifierOffset_ = startTrampolineCode(masm);
+void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm,
+                                            ArgumentsRectifierKind kind) {
+  switch (kind) {
+    case ArgumentsRectifierKind::Normal:
+      argumentsRectifierOffset_ = startTrampolineCode(masm);
+      break;
+    case ArgumentsRectifierKind::TrialInlining:
+      trialInliningArgumentsRectifierOffset_ = startTrampolineCode(masm);
+      break;
+  }
 
   // Caller:
   // [arg2] [arg1] [this] [ [argc] [callee] [descr] [raddr] ] <- esp
@@ -515,8 +523,16 @@ void JitRuntime::generateArgumentsRectifier(MacroAssembler& masm) {
 
   // Call the target function.
   masm.andl(Imm32(CalleeTokenMask), eax);
-  masm.loadJitCodeRaw(eax, eax);
-  argumentsRectifierReturnOffset_ = masm.callJitNoProfiler(eax);
+  switch (kind) {
+    case ArgumentsRectifierKind::Normal:
+      masm.loadJitCodeRaw(eax, eax);
+      argumentsRectifierReturnOffset_ = masm.callJitNoProfiler(eax);
+      break;
+    case ArgumentsRectifierKind::TrialInlining:
+      masm.loadBaselineJitCodeRaw(eax, eax);
+      masm.callJitNoProfiler(eax);
+      break;
+  }
 
   // Remove the rectifier frame.
   masm.pop(ebx);                           // ebx <- descriptor with FrameType.