[MERGE #6170 @nhat-nguyen] Enable basic jitting without global optimizer for generators on x64

Merge pull request #6170 from nhat-nguyen:generator

- Cleaned up the bail-in code and the IR dumps with self-descriptive label names and added special opcodes just for generators; this made it easier to read the dumps in general and to perform some other logic using those opcodes
- Fixed epilogues in generators. We now have 2 epilogue labels in generators that can be jumped to. The first one nulls out the interpreter frame to signal generator completion and restores all callee-save registers - used for return statements. The other one only does the latter, used for jumps after bailing out. (the register restore part is shared between the 2 labels)
- We also already have the mechanism to deal with yield points that were optimized away (called `BailOutForElidedYield`), I also changed the logic a bit to work with these new epilogue labels
- The resume jump table must be placed close enough to the beginning of the jit'd function so that we don't accidentally re-run the code, but also further below enough so that we can reload "global" values such as constants / environment (right now, it is placed right after we load the environment). One thing to note is that objects like Environment are always loaded from the `ScriptFunction`, but some like the `FrameDisplay/LocalClosure` might be created on the heap and transferred over to the interpreter frame afterwards. So we want to place the jump table before anything that can created new objects on the heap
    - also made sure that the bail-in code doesn't blindly overwrite any of the values we reload before the jump table
- Right now the bail-in code does the reverse of what the bail-out code does: grabbing all those special field members and putting them back, the rest of the bytecode symbols are simply loaded from the usual `localSlots` space
- Made for-in enumerators work with generators by always using the space allocated on the interpreter frame
- Previously, in some cases, we would not have an interpreter frame the first time we are in the jit'd code and only create the frame once we bail out for the first time. This is fine for normal generators because we would bail-out to return the generator object, so the next time the interpreter frame would have already been created. However, `async` functions don't have this first yield point, so we cannot use the interpreter frame space to store the enumerators.
    - to fix this, we would always create an interpreter frame in the jit if the generator doesn't already have one
- Since the bail-in code relies on rax and rcx to do the reload, sometimes we would accidentally overwrite symbols reloaded before the jump table that might already be in those two registers. I made sure that we save/restore rax/rcx in the bail-in code if there are such symbols

Author: nhat-nguyen

lib/Backend/BackendApi.cpp

@@ -126,7 +126,7 @@ Js::JavascriptMethod GetCheckAsmJsCodeGenThunk()
 
 uint GetBailOutRegisterSaveSlotCount()
 {
-    // REVIEW: not all registers are used, we are allocating more space then necessary.
+    // REVIEW: not all registers are used, we are allocating more space than necessary.
     return LinearScanMD::GetRegisterSaveSlotCount();
 }
 

lib/Backend/BackwardPass.cpp

@@ -88,20 +88,26 @@ BackwardPass::DoMarkTempNumbers() const
 }
 
 bool
-BackwardPass::DoMarkTempObjects() const
-{
-    // only mark temp object on the backward store phase
-    return (tag == Js::BackwardPhase) && !PHASE_OFF(Js::MarkTempPhase, this->func) &&
-        !PHASE_OFF(Js::MarkTempObjectPhase, this->func) && func->DoGlobOpt() && func->GetHasTempObjectProducingInstr() &&
-        !func->IsJitInDebugMode() &&
-        func->DoGlobOptsForGeneratorFunc();
+BackwardPass::SatisfyMarkTempObjectsConditions() const {
+    return !PHASE_OFF(Js::MarkTempPhase, this->func) &&
+           !PHASE_OFF(Js::MarkTempObjectPhase, this->func) &&
+           func->DoGlobOpt() && func->GetHasTempObjectProducingInstr() &&
+           !func->IsJitInDebugMode() &&
+           func->DoGlobOptsForGeneratorFunc();
 
     // Why MarkTempObject is disabled under debugger:
     //   We add 'identified so far dead non-temp locals' to byteCodeUpwardExposedUsed in ProcessBailOutInfo,
     //   this may cause MarkTempObject to convert some temps back to non-temp when it sees a 'transferred exposed use'
     //   from a temp to non-temp. That's in general not a supported conversion (while non-temp -> temp is fine).
 }
 
+bool
+BackwardPass::DoMarkTempObjects() const
+{
+    // only mark temp object on the backward store phase
+    return (tag == Js::BackwardPhase) && SatisfyMarkTempObjectsConditions();
+}
+
 bool
 BackwardPass::DoMarkTempNumbersOnTempObjects() const
 {
@@ -113,8 +119,7 @@ bool
 BackwardPass::DoMarkTempObjectVerify() const
 {
     // only mark temp object on the backward store phase
-    return (tag == Js::DeadStorePhase) && !PHASE_OFF(Js::MarkTempPhase, this->func) &&
-        !PHASE_OFF(Js::MarkTempObjectPhase, this->func) && func->DoGlobOpt() && func->GetHasTempObjectProducingInstr();
+    return (tag == Js::DeadStorePhase) && SatisfyMarkTempObjectsConditions();
 }
 #endif
 
@@ -7709,7 +7714,7 @@ BackwardPass::ProcessDef(IR::Opnd * opnd)
         PropertySym *propertySym = sym->AsPropertySym();
         ProcessStackSymUse(propertySym->m_stackSym, isJITOptimizedReg);
 
-        if(IsCollectionPass())
+        if (IsCollectionPass())
         {
             return false;
         }
@@ -7796,7 +7801,7 @@ BackwardPass::ProcessDef(IR::Opnd * opnd)
             }
         }
 
-        if(IsCollectionPass())
+        if (IsCollectionPass())
         {
             return false;
         }
@@ -8247,9 +8252,20 @@ BackwardPass::ProcessBailOnNoProfile(IR::Instr *instr, BasicBlock *block)
         return false;
     }
 
+    // For generator functions, we don't want to move the BailOutOnNoProfile above
+    // certain instructions such as ResumeYield/ResumeYieldStar/CreateInterpreterStackFrameForGenerator
+    // This indicates the insertion point for the BailOutOnNoProfile in such cases.
+    IR::Instr *insertionPointForGenerator = nullptr;
+
     // Don't hoist if we see calls with profile data (recursive calls)
     while(!curInstr->StartsBasicBlock())
     {
+        if (curInstr->DontHoistBailOnNoProfileAboveInGeneratorFunction())
+        {
+            Assert(insertionPointForGenerator == nullptr);
+            insertionPointForGenerator = curInstr;
+        }
+
         // If a function was inlined, it must have had profile info.
         if (curInstr->m_opcode == Js::OpCode::InlineeEnd || curInstr->m_opcode == Js::OpCode::InlineBuiltInEnd || curInstr->m_opcode == Js::OpCode::InlineNonTrackingBuiltInEnd
             || curInstr->m_opcode == Js::OpCode::InlineeStart || curInstr->m_opcode == Js::OpCode::EndCallForPolymorphicInlinee)
@@ -8320,7 +8336,8 @@ BackwardPass::ProcessBailOnNoProfile(IR::Instr *instr, BasicBlock *block)
     // Now try to move this up the flowgraph to the predecessor blocks
     FOREACH_PREDECESSOR_BLOCK(pred, block)
     {
-        bool hoistBailToPred = true;
+        // Don't hoist BailOnNoProfile up past blocks containing ResumeYield/ResumeYieldStar
+        bool hoistBailToPred = (insertionPointForGenerator == nullptr);
 
         if (block->isLoopHeader && pred->loop == block->loop)
         {
@@ -8396,10 +8413,19 @@ BackwardPass::ProcessBailOnNoProfile(IR::Instr *instr, BasicBlock *block)
 #if DBG
         blockHeadInstr->m_noHelperAssert = true;
 #endif
-        block->beginsBailOnNoProfile = true;
 
         instr->m_func = curInstr->m_func;
-        curInstr->InsertAfter(instr);
+
+        if (insertionPointForGenerator != nullptr)
+        {
+            insertionPointForGenerator->InsertAfter(instr);
+            block->beginsBailOnNoProfile = false;
+        }
+        else
+        {
+            curInstr->InsertAfter(instr);
+            block->beginsBailOnNoProfile = true;
+        }
 
         bool setLastInstr = (curInstr == block->GetLastInstr());
         if (setLastInstr)

lib/Backend/BackwardPass.h

@@ -118,6 +118,8 @@ class BackwardPass
     bool DoByteCodeUpwardExposedUsed() const;
     bool DoCaptureByteCodeUpwardExposedUsed() const;
     void DoSetDead(IR::Opnd * opnd, bool isDead) const;
+
+    bool SatisfyMarkTempObjectsConditions() const;
     bool DoMarkTempObjects() const;
     bool DoMarkTempNumbers() const;
     bool DoMarkTempNumbersOnTempObjects() const;

lib/Backend/BailOut.cpp

@@ -989,6 +989,25 @@ BailOutRecord::RestoreValue(IR::BailOutKind bailOutKind, Js::JavascriptCallStack
         value = Js::JavascriptNumber::ToVar(int32Value, scriptContext);
         BAILOUT_VERBOSE_TRACE(newInstance->function->GetFunctionBody(), bailOutKind, _u(", value: %10d (ToVar: 0x%p)"), int32Value, value);
     }
+    else if (regSlot == newInstance->function->GetFunctionBody()->GetYieldRegister() && newInstance->function->GetFunctionBody()->IsCoroutine())
+    {
+        // This value can only either be:
+        // 1) the ResumeYieldData. Even though this value is on the stack, it is only used to extract the data as part of Op_ResumeYield.
+        //    So there is no need to box the value.
+        // 2) the object used as the return value for yield statement. This object is created on the heap, so no need to box either.
+        Assert(value);
+
+#if ENABLE_DEBUG_CONFIG_OPTIONS
+        if (ThreadContext::IsOnStack(value))
+        {
+            BAILOUT_VERBOSE_TRACE(newInstance->function->GetFunctionBody(), bailOutKind, _u(", value: 0x%p (ResumeYieldData)"), value);
+        }
+        else
+        {
+            BAILOUT_VERBOSE_TRACE(newInstance->function->GetFunctionBody(), bailOutKind, _u(", value: 0x%p (Yield Return Value)"), value);
+        }
+#endif
+    }
     else
     {
         BAILOUT_VERBOSE_TRACE(newInstance->function->GetFunctionBody(), bailOutKind, _u(", value: 0x%p"), value);
@@ -1507,63 +1526,32 @@ BailOutRecord::BailOutHelper(Js::JavascriptCallStackLayout * layout, Js::ScriptF
     if (executeFunction->IsCoroutine())
     {
         // If the FunctionBody is a generator then this call is being made by one of the three
-        // generator resuming methods: next(), throw(), or return().  They all pass the generator
-        // object as the first of two arguments.  The real user arguments are obtained from the
-        // generator object.  The second argument is the ResumeYieldData which is only needed
-        // when resuming a generator and not needed when yielding from a generator, as is occurring
-        // here.
+        // generator resuming methods: next(), throw(), or return(). They all pass the generator
+        // object as the first of two arguments. The real user arguments are obtained from the
+        // generator object. The second argument is the ResumeYieldData which is only needed when
+        // resuming a generator and not needed when yielding from a generator, as is occurring here.
         AssertMsg(args.Info.Count == 2, "Generator ScriptFunctions should only be invoked by generator APIs with the pair of arguments they pass in -- the generator object and a ResumeYieldData pointer");
         Js::JavascriptGenerator* generator = Js::VarTo<Js::JavascriptGenerator>(args[0]);
         newInstance = generator->GetFrame();
 
-        if (newInstance != nullptr)
-        {
-            // BailOut will recompute OutArg pointers based on BailOutRecord.  Reset them back
-            // to initial position before that happens so that OP_StartCall calls don't accumulate
-            // incorrectly over multiple yield bailouts.
-            newInstance->ResetOut();
-
-            // The debugger relies on comparing stack addresses of frames to decide when a step_out is complete so
-            // give the InterpreterStackFrame a legit enough stack address to make this comparison work.
-            newInstance->m_stackAddress = reinterpret_cast<DWORD_PTR>(&generator);
-        }
-        else
-        {
-            //
-            // Allocate a new InterpreterStackFrame instance on the recycler heap.
-            // It will live with the JavascriptGenerator object.
-            //
-            Js::Arguments generatorArgs = generator->GetArguments();
-            Js::InterpreterStackFrame::Setup setup(function, generatorArgs, true, isInlinee);
-            Assert(setup.GetStackAllocationVarCount() == 0);
-            size_t varAllocCount = setup.GetAllocationVarCount();
-            size_t varSizeInBytes = varAllocCount * sizeof(Js::Var);
-            DWORD_PTR stackAddr = reinterpret_cast<DWORD_PTR>(&generator); // as mentioned above, use any stack address from this frame to ensure correct debugging functionality
-            Js::LoopHeader* loopHeaderArray = executeFunction->GetHasAllocatedLoopHeaders() ? executeFunction->GetLoopHeaderArrayPtr() : nullptr;
-
-            allocation = RecyclerNewPlus(functionScriptContext->GetRecycler(), varSizeInBytes, Js::Var);
-
-            // Initialize the interpreter stack frame (constants) but not the param, the bailout record will restore the value
-#if DBG
-            // Allocate invalidVar on GC instead of stack since this InterpreterStackFrame will out live the current real frame
-            Js::Var invalidVar = (Js::RecyclableObject*)RecyclerNewPlusLeaf(functionScriptContext->GetRecycler(), sizeof(Js::RecyclableObject), Js::Var);
-            memset(invalidVar, 0xFE, sizeof(Js::RecyclableObject));
-#endif
+        // The jit relies on the interpreter stack frame to store various information such as
+        // for-in enumerators. Therefore, we always create an interpreter stack frame for generator
+        // as part of the resume jump table, at the beginning of the jit'd function, if it doesn't
+        // already exist.
+        Assert(newInstance != nullptr);
 
-            newInstance = setup.InitializeAllocation(allocation, nullptr, false, false, loopHeaderArray, stackAddr
-#if DBG
-                , invalidVar
-#endif
-                );
+        // BailOut will recompute OutArg pointers based on BailOutRecord. Reset them back
+        // to initial position before that happens so that OP_StartCall calls don't accumulate
+        // incorrectly over multiple yield bailouts.
+        newInstance->ResetOut();
 
-            newInstance->m_reader.Create(executeFunction);
-
-            generator->SetFrame(newInstance, varSizeInBytes);
-        }
+        // The debugger relies on comparing stack addresses of frames to decide when a step_out is complete so
+        // give the InterpreterStackFrame a legit enough stack address to make this comparison work.
+        newInstance->m_stackAddress = reinterpret_cast<DWORD_PTR>(&generator);
     }
     else
     {
-        Js::InterpreterStackFrame::Setup setup(function, args, true, isInlinee);
+        Js::InterpreterStackFrame::Setup setup(function, args, true /* bailedOut */, isInlinee);
         size_t varAllocCount = setup.GetAllocationVarCount();
         size_t stackVarAllocCount = setup.GetStackAllocationVarCount();
         size_t varSizeInBytes;
@@ -2826,7 +2814,7 @@ void BailOutRecord::CheckPreemptiveRejit(Js::FunctionBody* executeFunction, IR::
 
 Js::Var BailOutRecord::BailOutForElidedYield(void * framePointer)
 {
-    JIT_HELPER_REENTRANT_HEADER(NoSaveRegistersBailOutForElidedYield);
+    JIT_HELPER_NOT_REENTRANT_NOLOCK_HEADER(NoSaveRegistersBailOutForElidedYield);
     Js::JavascriptCallStackLayout * const layout = Js::JavascriptCallStackLayout::FromFramePointer(framePointer);
     Js::ScriptFunction ** functionRef = (Js::ScriptFunction **)&layout->functionObject;
     Js::ScriptFunction * function = *functionRef;

lib/Backend/FlowGraph.cpp

@@ -3318,7 +3318,7 @@ FlowGraph::RemoveInstr(IR::Instr *instr, GlobOpt * globOpt)
             if (opcode == Js::OpCode::Yield)
             {
                 IR::Instr *instrLabel = newByteCodeUseInstr->m_next;
-                while (instrLabel->m_opcode != Js::OpCode::Label)
+                while (instrLabel->m_opcode != Js::OpCode::GeneratorBailInLabel)
                 {
                     instrLabel = instrLabel->m_next;
                 }

lib/Backend/Func.cpp

@@ -148,6 +148,7 @@ Func::Func(JitArenaAllocator *alloc, JITTimeWorkItem * workItem,
     , m_forInEnumeratorArrayOffset(-1)
     , argInsCount(0)
     , m_globalObjTypeSpecFldInfoArray(nullptr)
+    , m_forInEnumeratorForGeneratorSym(nullptr)
 #if LOWER_SPLIT_INT64
     , m_int64SymPairMap(nullptr)
 #endif

lib/Backend/Func.h

@@ -205,7 +205,8 @@ class Func
         return
             !PHASE_OFF(Js::GlobOptPhase, this) && !IsSimpleJit() &&
             (!GetTopFunc()->HasTry() || GetTopFunc()->CanOptimizeTryCatch()) &&
-            (!GetTopFunc()->HasFinally() || GetTopFunc()->CanOptimizeTryFinally());
+            (!GetTopFunc()->HasFinally() || GetTopFunc()->CanOptimizeTryFinally()) &&
+            !GetTopFunc()->GetJITFunctionBody()->IsCoroutine();
     }
 
     bool DoInline() const
@@ -1061,11 +1062,23 @@ static const unsigned __int64 c_debugFillPattern8 = 0xcececececececece;
     StackSym* m_loopParamSym;
     StackSym* m_bailoutReturnValueSym;
     StackSym* m_hasBailedOutSym;
+    StackSym* m_forInEnumeratorForGeneratorSym;
 
 public:
     StackSym* tempSymDouble;
     StackSym* tempSymBool;
 
+    void SetForInEnumeratorSymForGeneratorSym(StackSym* sym)
+    {
+        Assert(this->m_forInEnumeratorForGeneratorSym == nullptr);
+        this->m_forInEnumeratorForGeneratorSym = sym;
+    }
+
+    StackSym* GetForInEnumeratorSymForGeneratorSym() const
+    {
+        return this->m_forInEnumeratorForGeneratorSym;
+    }
+
     // StackSyms' corresponding getters/setters
     void SetInlineeFrameStartSym(StackSym* sym)
     {

lib/Backend/GlobOpt.cpp

@@ -4704,6 +4704,14 @@ GlobOpt::ValueNumberDst(IR::Instr **pInstr, Value *src1Val, Value *src2Val)
     case Js::OpCode::Coerce_Str:
         AssertMsg(instr->GetDst()->GetValueType().IsString(),
             "Creator of this instruction should have set the type");
+
+        // Due to fall through and the fact that Ld_A only takes one source,
+        // free the other source here.
+        if (instr->GetSrc2() && !(this->IsLoopPrePass() || src1ValueInfo == nullptr || !src1ValueInfo->IsString()))
+        {
+            instr->FreeSrc2();
+        }
+
         // fall-through
     case Js::OpCode::Coerce_StrOrRegex:
         // We don't set the ValueType of src1 for Coerce_StrOrRegex, hence skip the ASSERT

lib/Backend/IR.cpp

@@ -1047,6 +1047,13 @@ bool IR::Instr::IsStElemVariant() const
         this->m_opcode == Js::OpCode::StElemC;
 }
 
+bool IR::Instr::DontHoistBailOnNoProfileAboveInGeneratorFunction() const
+{
+    return this->m_opcode == Js::OpCode::ResumeYield ||
+        this->m_opcode == Js::OpCode::ResumeYieldStar ||
+        this->m_opcode == Js::OpCode::GeneratorCreateInterpreterStackFrame;
+}
+
 bool IR::Instr::CanChangeFieldValueWithoutImplicitCall() const
 {
     // TODO: Why is InitFld necessary?

lib/Backend/IR.h

@@ -336,6 +336,8 @@ class Instr
     static Instr*   FindSingleDefInstr(Js::OpCode opCode, Opnd* src);
     bool            CanAggregateByteCodeUsesAcrossInstr(IR::Instr * instr);
 
+    bool            DontHoistBailOnNoProfileAboveInGeneratorFunction() const;
+
     // LazyBailOut
     bool            AreAllOpndsTypeSpecialized() const;
     bool            IsStFldVariant() const;
@@ -798,6 +800,7 @@ class LabelInstr : public Instr
     inline void             SetRegion(Region *);
     inline Region *         GetRegion(void) const;
     inline BOOL             IsUnreferenced(void) const;
+    inline BOOL             IsGeneratorEpilogueLabel(void) const;
 
     LabelInstr *            CloneLabel(BOOL fCreate);