Force spill rax/rcx before Yield so that we don't have to save them

Author: nhat-nguyen

lib/Backend/LinearScan.cpp

@@ -3358,6 +3358,23 @@ LinearScan::KillImplicitRegs(IR::Instr *instr)
         this->RecordLoopUse(nullptr, LowererMDArch::GetRegIMulHighDestLower());
         return;
     }
+
+    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]);
+        }
+
+        return;
+    }
 #endif
 
     this->TrackInlineeArgLifetimes(instr);

lib/Backend/amd64/LinearScanMD.cpp

@@ -517,35 +517,21 @@ IR::Instr* LinearScanMD::GeneratorBailIn::GenerateBailIn(IR::Instr* resumeLabelI
         instrAfter  /* instrInsertRegSym */
     };
 
-    SaveInitializedRegister saveInitializedReg { false /* rax */, false /* rcx */ };
-
     // 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, saveInitializedReg);
+    this->InsertRestoreSymbols(bailOutInfo->byteCodeUpwardExposedUsed, insertionPoint);
     Assert(!this->func->IsStackArgsEnabled());
 
-    // 5) Save/restore rax/rcx if needed
-    if (saveInitializedReg.rax)
-    {
-        this->InsertSaveAndRestore(resumeLabelInstr, instrAfter, raxRegOpnd);
-    }
-
-    if (saveInitializedReg.rcx)
-    {
-        this->InsertSaveAndRestore(resumeLabelInstr, instrAfter, rcxRegOpnd);
-    }
-
     return instrAfter;
 }
 
 void LinearScanMD::GeneratorBailIn::InsertRestoreSymbols(
     BVSparse<JitArenaAllocator>* symbols,
-    BailInInsertionPoint& insertionPoint,
-    SaveInitializedRegister& saveInitializedReg
+    BailInInsertionPoint& insertionPoint
 )
 {
     if (symbols == nullptr)
@@ -558,89 +544,77 @@ void LinearScanMD::GeneratorBailIn::InsertRestoreSymbols(
         StackSym* stackSym = this->func->m_symTable->FindStackSym(symId);
         Lifetime* lifetime = stackSym->scratch.linearScan.lifetime;
 
-        if (this->NeedsReloadingValueWhenBailIn(stackSym, lifetime))
+        if (!this->NeedsReloadingValueWhenBailIn(stackSym, lifetime))
         {
-            Js::RegSlot regSlot = stackSym->GetByteCodeRegSlot();
-            IR::Opnd* srcOpnd = IR::IndirOpnd::New(
-                this->raxRegOpnd,
-                this->GetOffsetFromInterpreterStackFrame(regSlot),
-                stackSym->GetType(),
-                this->func
-            );
-
-            if (lifetime->isSpilled)
+            continue;
+        }
+
+        Js::RegSlot regSlot = stackSym->GetByteCodeRegSlot();
+        IR::Opnd* srcOpnd = IR::IndirOpnd::New(
+            this->raxRegOpnd,
+            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->rcxRegOpnd, srcOpnd, insertionPoint.instrInsertStackSym);
+            LinearScan::InsertMove(dstOpnd, this->rcxRegOpnd, 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())
             {
-                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->rcxRegOpnd, srcOpnd, insertionPoint.instrInsertStackSym);
-                LinearScan::InsertMove(dstOpnd, this->rcxRegOpnd, insertionPoint.instrInsertStackSym);
+                instr = this->linearScanMD->linearScan->InsertLoad(insertionPoint.instrInsertRegSym, stackSym, lifetime->reg);
             }
             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::RegOpnd* dstRegOpnd = IR::RegOpnd::New(stackSym, stackSym->GetType(), this->func);
+                dstRegOpnd->SetReg(lifetime->reg);
+                instr = LinearScan::InsertMove(dstRegOpnd, srcOpnd, insertionPoint.instrInsertRegSym);
+            }
 
-                IR::Instr* instr;
-                
-                if (stackSym->IsConst())
-                {
-                    instr = this->linearScanMD->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 (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 == RegRAX)
+            {
+                // Ensure rax is restored last
+                Assert(insertionPoint.instrInsertRegSym != insertionPoint.instrInsertStackSym);
 
-                if (lifetime->reg == RegRAX)
+                insertionPoint.instrInsertRegSym = instr;
+
+                if (insertionPoint.raxRestoreInstr != nullptr)
                 {
-                    // 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->raxRegOpnd);
-                    }
-
-                    insertionPoint.raxRestoreInstr = instr;
+                    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->raxRegOpnd);
                 }
 
-                this->linearScanMD->linearScan->RecordDef(lifetime, instr, 0);
-            }
-        }
-        else
-        {
-            if (lifetime->reg == RegRAX)
-            {
-                Assert(!saveInitializedReg.rax);
-                saveInitializedReg.rax = true;
+                insertionPoint.raxRestoreInstr = instr;
             }
 
-            if (lifetime->reg == RegRCX)
-            {
-                Assert(!saveInitializedReg.rcx);
-                saveInitializedReg.rcx = true;
-            }
+            this->linearScanMD->linearScan->RecordDef(lifetime, instr, 0);
         }
     }
     NEXT_BITSET_IN_SPARSEBV;
@@ -699,11 +673,3 @@ uint32 LinearScanMD::GeneratorBailIn::GetOffsetFromInterpreterStackFrame(Js::Reg
         return regSlot * sizeof(Js::Var) + Js::InterpreterStackFrame::GetOffsetOfLocals();
     }
 }
-
-void LinearScanMD::GeneratorBailIn::InsertSaveAndRestore(IR::Instr* start, IR::Instr* end, IR::RegOpnd* reg)
-{
-    IR::Instr* push = IR::Instr::New(Js::OpCode::PUSH, nullptr /* dst */, reg /* src1 */, this->func);
-    IR::Instr* pop = IR::Instr::New(Js::OpCode::POP, reg /* dst */, this->func);
-    start->InsertAfter(push);
-    end->InsertBefore(pop);
-}

lib/Backend/amd64/LinearScanMD.h

@@ -88,17 +88,6 @@ class LinearScanMD : public LinearScanMDShared
             IR::Instr* instrInsertRegSym;
         };
 
-        // There are symbols that we don't need to restore such as constant values,
-        // ScriptFunction's Environment (through LdEnv) and the address pointing to for-in enumerator
-        // on the interpreter frame because their values are already loaded before (or as part) of the
-        // generator resume jump table. In such cases, they could already be in either rax or rcx.
-        // So we would need to save their values (and restore afterwards) before generating the bail-in code.
-        struct SaveInitializedRegister
-        {
-            bool rax;
-            bool rcx;
-        };
-
         Func* const func;
         LinearScanMD* const linearScanMD;
         const JITTimeFunctionBody* const jitFnBody;
@@ -110,13 +99,7 @@ class LinearScanMD : public LinearScanMDShared
         uint32 GetOffsetFromInterpreterStackFrame(Js::RegSlot regSlot) const;
         IR::SymOpnd* CreateGeneratorObjectOpnd() const;
 
-        void InsertSaveAndRestore(IR::Instr* start, IR::Instr* end, IR::RegOpnd* reg);
-
-        void InsertRestoreSymbols(
-            BVSparse<JitArenaAllocator>* symbols,
-            BailInInsertionPoint& insertionPoint,
-            SaveInitializedRegister& saveInitializedReg
-        );
+        void InsertRestoreSymbols(BVSparse<JitArenaAllocator>* symbols, BailInInsertionPoint& insertionPoint);
 
     public:
         GeneratorBailIn(Func* func, LinearScanMD* linearScanMD);