Move bail-in implementation to common LinearScan

Author: nhat-nguyen

lib/Backend/LinearScan.cpp

@@ -306,7 +306,7 @@ LinearScan::RegAlloc()
 
         if (insertBailInAfter == instr)
         {
-            instrNext = linearScanMD.GenerateBailInForGeneratorYield(instr, bailOutInfoForBailIn);
+            instrNext = this->bailIn.GenerateBailIn(instr, bailOutInfoForBailIn);
             insertBailInAfter = nullptr;
             bailOutInfoForBailIn = nullptr;
         }
@@ -3361,18 +3361,7 @@ LinearScan::KillImplicitRegs(IR::Instr *instr)
 
     if (instr->m_opcode == Js::OpCode::Yield)
     {
-#if defined(_M_X64)
-        RegNum regs[] = { RegRAX, RegRCX };
-#else
-        RegNum regs[] = { RegEAX, RegECX };
-#endif
-        for (int i = 0; i < 2; i++)
-        {
-            this->SpillReg(regs[i]);
-            this->tempRegs.Clear(regs[i]);
-            this->RecordLoopUse(nullptr, regs[i]);
-        }
-
+        this->bailIn.SpillRegsForBailIn();
         return;
     }
 #endif
@@ -4953,3 +4942,246 @@ LinearScan::ProcessLazyBailOut(IR::Instr *instr)
         this->FillBailOutRecord(instr);
     }
 }
+
+LinearScan::GeneratorBailIn::GeneratorBailIn(Func* func, LinearScan* linearScan) :
+    func { func },
+    linearScan { linearScan },
+    jitFnBody { func->GetJITFunctionBody() },
+    initializedRegs { func->m_alloc },
+    regs {
+#if defined(_M_X64)
+    RegRAX, RegRCX
+#elif defined(_M_IX86)
+    RegEAX, RegECX
+#endif
+    },
+    interpreterFrameRegOpnd { IR::RegOpnd::New(nullptr, regs[0], TyMachPtr, func) },
+    tempRegOpnd { IR::RegOpnd::New(nullptr, regs[1], TyVar, func) }
+{
+    // The yield register holds the evaluated value of the expression passed as
+    // the parameter to .next(), this can be obtained from the generator object itself,
+    // so no need to restore.
+    this->initializedRegs.Set(this->jitFnBody->GetYieldReg());
+
+    // The environment is loaded before the resume jump table, no need to restore either.
+    this->initializedRegs.Set(this->jitFnBody->GetEnvReg());
+}
+
+void LinearScan::GeneratorBailIn::SpillRegsForBailIn()
+{
+    for (int i = 0; i < GeneratorBailIn::regNum; i++)
+    {
+        this->linearScan->SpillReg(this->regs[i]);
+        this->linearScan->tempRegs.Clear(this->regs[i]);
+        this->linearScan->RecordLoopUse(nullptr, this->regs[i]);
+    }
+}
+
+// Restores the live stack locations followed by the live registers from
+// the interpreter's register slots.
+// RecordDefs each live register that is restored.
+//
+// Generates the following code:
+// 
+// PUSH rax ; if needed
+// PUSH rcx ; if needed
+//
+// MOV rax, param0
+// MOV rax, [rax + JavascriptGenerator::GetFrameOffset()]
+//
+// for each live stack location, sym
+//
+//   MOV rcx, [rax + regslot offset]
+//   MOV sym(stack location), rcx
+//
+// for each live register, sym (rax is restore last if it is live)
+//
+//   MOV sym(register), [rax + regslot offset]
+//
+// POP rax; if needed
+// POP rcx; if needed
+IR::Instr* LinearScan::GeneratorBailIn::GenerateBailIn(IR::Instr* resumeLabelInstr, BailOutInfo* bailOutInfo)
+{
+    Assert(!bailOutInfo->capturedValues || bailOutInfo->capturedValues->constantValues.Empty());
+    Assert(!bailOutInfo->capturedValues || bailOutInfo->capturedValues->copyPropSyms.Empty());
+    Assert(!bailOutInfo->liveLosslessInt32Syms || bailOutInfo->liveLosslessInt32Syms->IsEmpty());
+    Assert(!bailOutInfo->liveFloat64Syms || bailOutInfo->liveFloat64Syms->IsEmpty());
+
+    IR::Instr* instrAfter = resumeLabelInstr->m_next;
+
+    // 1) Load the generator object that was passed as one of the arguments to the jitted frame
+    LinearScan::InsertMove(this->interpreterFrameRegOpnd, this->CreateGeneratorObjectOpnd(), instrAfter);
+
+    // 2) Gets the InterpreterStackFrame pointer into rax
+    IR::IndirOpnd* generatorFrameOpnd = IR::IndirOpnd::New(this->interpreterFrameRegOpnd, Js::JavascriptGenerator::GetFrameOffset(), TyMachPtr, this->func);
+    LinearScan::InsertMove(this->interpreterFrameRegOpnd, generatorFrameOpnd, instrAfter);
+
+    // 3) Put the Javascript's `arguments` object, which is stored in the interpreter frame, to the jit's stack slot if needed
+    //    See BailOutRecord::RestoreValues
+    if (this->func->HasArgumentSlot())
+    {
+        IR::IndirOpnd* generatorArgumentsOpnd = IR::IndirOpnd::New(this->interpreterFrameRegOpnd, Js::InterpreterStackFrame::GetOffsetOfArguments(), TyMachPtr, this->func);
+        LinearScan::InsertMove(this->tempRegOpnd, generatorArgumentsOpnd, instrAfter);
+        LinearScan::InsertMove(LowererMD::CreateStackArgumentsSlotOpnd(this->func), this->tempRegOpnd, instrAfter);
+    }
+
+    BailInInsertionPoint insertionPoint
+    {
+        nullptr,    /* raxRestoreInstr */
+        instrAfter, /* instrInsertStackSym */
+        instrAfter  /* instrInsertRegSym */
+    };
+
+    // 4) Restore symbols
+    // - We don't need to restore argObjSyms because StackArgs is currently not enabled
+    //   Commented out here in case we do want to enable it in the future:
+    // this->InsertRestoreSymbols(bailOutInfo->capturedValues->argObjSyms, insertionPoint, saveInitializedReg);
+    // 
+    // - We move all argout symbols right before the call so we don't need to restore argouts either
+    this->InsertRestoreSymbols(bailOutInfo->byteCodeUpwardExposedUsed, insertionPoint);
+    Assert(!this->func->IsStackArgsEnabled());
+
+    return instrAfter;
+}
+
+void LinearScan::GeneratorBailIn::InsertRestoreSymbols(BVSparse<JitArenaAllocator>* symbols, BailInInsertionPoint& insertionPoint)
+{
+    if (symbols == nullptr)
+    {
+        return;
+    }
+
+    FOREACH_BITSET_IN_SPARSEBV(symId, symbols)
+    {
+        StackSym* stackSym = this->func->m_symTable->FindStackSym(symId);
+        Lifetime* lifetime = stackSym->scratch.linearScan.lifetime;
+
+        if (!this->NeedsReloadingValueWhenBailIn(stackSym, lifetime))
+        {
+            continue;
+        }
+
+        Js::RegSlot regSlot = stackSym->GetByteCodeRegSlot();
+        IR::Opnd* srcOpnd = IR::IndirOpnd::New(
+            this->interpreterFrameRegOpnd,
+            this->GetOffsetFromInterpreterStackFrame(regSlot),
+            stackSym->GetType(),
+            this->func
+        );
+
+        if (lifetime->isSpilled)
+        {
+            Assert(!stackSym->IsConst());
+            // Stack restores require an extra register since we can't move an indir directly to an indir on amd64
+            IR::SymOpnd* dstOpnd = IR::SymOpnd::New(stackSym, stackSym->GetType(), this->func);
+            LinearScan::InsertMove(this->tempRegOpnd, srcOpnd, insertionPoint.instrInsertStackSym);
+            LinearScan::InsertMove(dstOpnd, this->tempRegOpnd, insertionPoint.instrInsertStackSym);
+        }
+        else
+        {
+            // Register restores must come after stack restores so that we have RAX and RCX free to
+            // use for stack restores and further RAX must be restored last since it holds the
+            // pointer to the InterpreterStackFrame from which we are restoring values.
+            // We must also track these restores using RecordDef in case the symbols are spilled.
+
+            IR::Instr* instr;
+
+            if (stackSym->IsConst())
+            {
+                instr = this->linearScan->InsertLoad(insertionPoint.instrInsertRegSym, stackSym, lifetime->reg);
+            }
+            else
+            {
+                IR::RegOpnd* dstRegOpnd = IR::RegOpnd::New(stackSym, stackSym->GetType(), this->func);
+                dstRegOpnd->SetReg(lifetime->reg);
+                instr = LinearScan::InsertMove(dstRegOpnd, srcOpnd, insertionPoint.instrInsertRegSym);
+            }
+
+            if (insertionPoint.instrInsertRegSym == insertionPoint.instrInsertStackSym)
+            {
+                // This is the first register sym, make sure we don't insert stack stores
+                // after this instruction so we can ensure rax and rcx remain free to use
+                // for restoring spilled stack syms.
+                insertionPoint.instrInsertStackSym = instr;
+            }
+
+            if (lifetime->reg == interpreterFrameRegOpnd->GetReg())
+            {
+                // Ensure rax is restored last
+                Assert(insertionPoint.instrInsertRegSym != insertionPoint.instrInsertStackSym);
+
+                insertionPoint.instrInsertRegSym = instr;
+
+                if (insertionPoint.raxRestoreInstr != nullptr)
+                {
+                    AssertMsg(false, "this is unexpected until copy prop is enabled");
+                    // rax was mapped to multiple bytecode registers.  Obviously only the first
+                    // restore we do will work so change all following stores to `mov rax, rax`.
+                    // We still need to keep them around for RecordDef in case the corresponding
+                    // dst sym is spilled later on.
+                    insertionPoint.raxRestoreInstr->FreeSrc1();
+                    insertionPoint.raxRestoreInstr->SetSrc1(this->interpreterFrameRegOpnd);
+                }
+
+                insertionPoint.raxRestoreInstr = instr;
+            }
+
+            this->linearScan->RecordDef(lifetime, instr, 0);
+        }
+    }
+    NEXT_BITSET_IN_SPARSEBV;
+}
+
+bool LinearScan::GeneratorBailIn::NeedsReloadingValueWhenBailIn(StackSym* sym, Lifetime* lifetime) const
+{
+    if (sym->IsConst())
+    {
+        if (this->func->GetJITFunctionBody()->RegIsConstant(sym->GetByteCodeRegSlot()))
+        {
+            return false;
+        }
+        else
+        {
+            return !lifetime->isSpilled;
+        }
+    }
+
+    // If we have for-in in the generator, don't need to reload the symbol again as it is done
+    // during the resume jump table
+    if (this->func->GetForInEnumeratorSymForGeneratorSym() && this->func->GetForInEnumeratorSymForGeneratorSym()->m_id == sym->m_id)
+    {
+        return false;
+    }
+
+    // Check for other special registers that are already initialized
+    return !this->initializedRegs.Test(sym->GetByteCodeRegSlot());
+}
+
+IR::SymOpnd* LinearScan::GeneratorBailIn::CreateGeneratorObjectOpnd() const
+{
+    StackSym* sym = StackSym::NewParamSlotSym(1, this->func);
+    this->func->SetArgOffset(sym, LowererMD::GetFormalParamOffset() * MachPtr);
+    return IR::SymOpnd::New(sym, TyMachPtr, this->func);
+}
+
+uint32 LinearScan::GeneratorBailIn::GetOffsetFromInterpreterStackFrame(Js::RegSlot regSlot) const
+{
+    // Some objects aren't stored in the local space in interpreter frame, but instead
+    // in their own fields. Use their offsets in such cases.
+    if (regSlot == this->jitFnBody->GetLocalFrameDisplayReg())
+    {
+        return Js::InterpreterStackFrame::GetOffsetOfLocalFrameDisplay();
+    }
+    else if (regSlot == this->jitFnBody->GetLocalClosureReg())
+    {
+        return Js::InterpreterStackFrame::GetOffsetOfLocalClosure();
+    }
+    else if (regSlot == this->jitFnBody->GetParamClosureReg())
+    {
+        return Js::InterpreterStackFrame::GetOffsetOfParamClosure();
+    }
+    else
+    {
+        return regSlot * sizeof(Js::Var) + Js::InterpreterStackFrame::GetOffsetOfLocals();
+    }
+}

lib/Backend/LinearScan.h

@@ -103,7 +103,7 @@ class LinearScan
         linearScanMD(func), opHelperSpilledLiveranges(NULL), currentOpHelperBlock(NULL),
         lastLabel(NULL), numInt32Regs(0), numFloatRegs(0), stackPackInUseLiveRanges(NULL), stackSlotsFreeList(NULL),
         totalOpHelperFullVisitedLength(0), curLoop(NULL), currentBlock(nullptr), currentRegion(nullptr), m_bailOutRecordCount(0),
-        globalBailOutRecordTables(nullptr), lastUpdatedRowIndices(nullptr)
+        globalBailOutRecordTables(nullptr), lastUpdatedRowIndices(nullptr), bailIn(GeneratorBailIn(func, this))
     {
     }
 
@@ -227,4 +227,56 @@ class LinearScan
     static IR::Instr *  InsertMove(IR::Opnd *dst, IR::Opnd *src, IR::Instr *const insertBeforeInstr);
     static IR::Instr *  InsertLea(IR::RegOpnd *dst, IR::Opnd *src, IR::Instr *const insertBeforeInstr);
 
+    class GeneratorBailIn {
+        // We need to rely on 2 registers `rax` and `rcx` to generate the bail-in code.
+        // At this point, since `rax` already has the address of the generator's interpreter frame,
+        // we can easily get the symbols' values through something like: mov dst [rax + appropriate offset]
+        //
+        // There are 4 types of symbols that we have to deal with:
+        //  - symbols that are currently on the stack at this point. We need 2 instructions:
+        //     - Load the value to `rcx`: mov rcx [rax + offset]
+        //     - Finally load the value to its stack slot: mov [rbp + stack offset] rcx
+        //  - symbol that is in rax
+        //  - symbol that is in rcx
+        //  - symbols that are in the other registers. We only need 1 instruction:
+        //     - mov reg [rax + offset]
+        //
+        // Since restoring symbols on the stack might mess up values that will be in rax/rcx,
+        // and we want to maintain the invariant that rax has to hold the value of the interpreter
+        // frame, we need to restore the symbols in the following order:
+        //  - symbols in stack
+        //  - symbols in registers
+        //  - symbol in rax
+        //
+        // The following 3 instructions indicate the insertion points for the above cases:
+        struct BailInInsertionPoint
+        {
+            IR::Instr* raxRestoreInstr;
+            IR::Instr* instrInsertStackSym;
+            IR::Instr* instrInsertRegSym;
+        };
+
+        Func* const func;
+        LinearScan* const linearScan;
+        const JITTimeFunctionBody* const jitFnBody;
+        BVSparse<JitArenaAllocator> initializedRegs;
+
+        static constexpr int regNum = 2;
+        const RegNum regs[regNum];
+        IR::RegOpnd* const interpreterFrameRegOpnd;
+        IR::RegOpnd* const tempRegOpnd;
+
+        bool NeedsReloadingValueWhenBailIn(StackSym* sym, Lifetime* lifetime) const;
+        uint32 GetOffsetFromInterpreterStackFrame(Js::RegSlot regSlot) const;
+        IR::SymOpnd* CreateGeneratorObjectOpnd() const;
+
+        void InsertRestoreSymbols(BVSparse<JitArenaAllocator>* symbols, BailInInsertionPoint& insertionPoint);
+
+    public:
+        GeneratorBailIn(Func* func, LinearScan* linearScan);
+        IR::Instr* GenerateBailIn(IR::Instr* resumeLabelInstr, BailOutInfo* bailOutInfo);
+        void SpillRegsForBailIn();
+    };
+
+    GeneratorBailIn bailIn;
 };