JIT Generator (and Async func) loops (#6990)

* Disable JitES6Generators

* Jit Generator and async func LoopBodies

* Update Tests cases for Generator Jit

Author: rhuanjl

lib/Backend/IRBuilder.cpp

@@ -1,6 +1,6 @@
 //-------------------------------------------------------------------------------------------------------
 // Copyright (C) Microsoft. All rights reserved.
-// Copyright (c) 2021 ChakraCore Project Contributors. All rights reserved.
+// Copyright (c) ChakraCore Project Contributors. All rights reserved.
 // Licensed under the MIT license. See LICENSE.txt file in the project root for full license information.
 //-------------------------------------------------------------------------------------------------------
 #include "Backend.h"
@@ -118,7 +118,7 @@ IRBuilder::DoBailOnNoProfile()
         return false;
     }
 
-    if (m_func->GetTopFunc()->GetJITFunctionBody()->IsCoroutine())
+    if (m_func->GetTopFunc()->GetJITFunctionBody()->IsCoroutine() && !m_func->IsLoopBody())
     {
         return false;
     }
@@ -441,7 +441,7 @@ IRBuilder::Build()
     // Note that for generators, we insert the bailout after the jump table to allow
     // the generator's execution to proceed before bailing out. Otherwise, we would always
     // bail to the beginning of the function in the interpreter, creating an infinite loop.
-    if (m_func->IsJitInDebugMode() && !this->m_func->GetJITFunctionBody()->IsCoroutine())
+    if (m_func->IsJitInDebugMode() && (!this->m_func->GetJITFunctionBody()->IsCoroutine() || this->IsLoopBody()))
     {
         this->InsertBailOutForDebugger(m_functionStartOffset, IR::BailOutForceByFlag | IR::BailOutBreakPointInFunction | IR::BailOutStep, nullptr);
     }
@@ -1880,6 +1880,9 @@ IRBuilder::BuildReg2(Js::OpCode newOpcode, uint32 offset, Js::RegSlot R0, Js::Re
         break;
 
     case Js::OpCode::Yield:
+        // Jitting Loop Bodies containing Yield is not possible, blocked at callsites of GenerateLoopBody
+        AssertMsg(!this->IsLoopBody(), "Attempting to JIT loop body containing Yield");
+
         instr = IR::Instr::New(newOpcode, dstOpnd, src1Opnd, m_func);
         this->AddInstr(instr, offset);
         IR::Instr* yieldInstr = instr->ConvertToBailOutInstr(instr, IR::BailOutForGeneratorYield);
@@ -7849,6 +7852,7 @@ IRBuilder::GeneratorJumpTable::GeneratorJumpTable(Func* func, IRBuilder* irBuild
 IR::Instr*
 IRBuilder::GeneratorJumpTable::BuildJumpTable()
 {
+    AssertMsg(!this->m_func->IsLoopBody(), "Coroutine Loop Bodies can be jitted but should follow a different path");
     if (!this->m_func->GetJITFunctionBody()->IsCoroutine())
     {
         return this->m_irBuilder->m_lastInstr;

lib/Backend/Lower.cpp

@@ -5467,7 +5467,7 @@ Lowerer::LowerPrologEpilog()
     instr = m_func->m_exitInstr;
     AssertMsg(instr->IsExitInstr(), "Last instr isn't an ExitInstr...");
 
-    if (m_func->GetJITFunctionBody()->IsCoroutine())
+    if (m_func->GetJITFunctionBody()->IsCoroutine() && !m_func->IsLoopBody())
     {
         IR::LabelInstr* epilogueLabel = this->m_lowerGeneratorHelper.GetEpilogueForReturnStatements();
         this->m_lowerGeneratorHelper.InsertNullOutGeneratorFrameInEpilogue(epilogueLabel);
@@ -11527,6 +11527,7 @@ Lowerer::LowerArgIn(IR::Instr *instrArgIn)
 
             if (m_func->GetJITFunctionBody()->IsCoroutine())
             {
+                AssertMsg(!m_func->IsLoopBody(), "LoopBody Jit should not involve Rest params");
                 generatorArgsPtrOpnd = LoadGeneratorArgsPtr(instrArgIn);
             }
 
@@ -11544,7 +11545,7 @@ Lowerer::LowerArgIn(IR::Instr *instrArgIn)
     if (argIndex == 1)
     {
         // The "this" argument is not source-dependent and doesn't need to be checked.
-        if (m_func->GetJITFunctionBody()->IsCoroutine())
+        if (m_func->GetJITFunctionBody()->IsCoroutine() && !m_func->IsLoopBody())
         {
             generatorArgsPtrOpnd = LoadGeneratorArgsPtr(instrArgIn);
             ConvertArgOpndIfGeneratorFunction(instrArgIn, generatorArgsPtrOpnd);

lib/Backend/NativeCodeGenerator.cpp

@@ -1,5 +1,6 @@
 //-------------------------------------------------------------------------------------------------------
 // Copyright (C) Microsoft. All rights reserved.
+// Copyright (c) ChakraCore Project Contributors. All rights reserved.
 // Licensed under the MIT license. See LICENSE.txt file in the project root for full license information.
 //-------------------------------------------------------------------------------------------------------
 #include "Backend.h"
@@ -246,8 +247,11 @@ NativeCodeGenerator::GenerateAllFunctions(Js::FunctionBody * fn)
             for (uint i = 0; i < fn->GetLoopCount(); i++)
             {
                 Js::LoopHeader * loopHeader = fn->GetLoopHeader(i);
-                Js::EntryPointInfo * entryPointInfo = loopHeader->GetCurrentEntryPointInfo();
-                this->GenerateLoopBody(fn, loopHeader, entryPointInfo);
+                if (loopHeader->hasYield == false)
+                {
+                    Js::EntryPointInfo * entryPointInfo = loopHeader->GetCurrentEntryPointInfo();
+                    this->GenerateLoopBody(fn, loopHeader, entryPointInfo);
+                }
             }
         }
         else
@@ -631,6 +635,7 @@ void NativeCodeGenerator::GenerateLoopBody(Js::FunctionBody * fn, Js::LoopHeader
     ASSERT_THREAD();
     Assert(fn->GetScriptContext()->GetNativeCodeGenerator() == this);
     Assert(entryPoint->jsMethod == nullptr);
+    Assert(!loopHeader->hasYield);
 
 #if DBG_DUMP
     if (PHASE_TRACE1(Js::JITLoopBodyPhase))

lib/Common/ConfigFlagsList.h

@@ -675,12 +675,12 @@ PHASE(All)
 #define DEFAULT_CONFIG_ESNullishCoalescingOperator (true)
 #define DEFAULT_CONFIG_ESGlobalThis            (true)
 
-// Jitting generators has not been tested on ARM
-// enabled only for x86 and x64 for now
+// Jitting generator functions is not functional on ARM
+// Also still contains significant bugs on x86/x64 hence disabled
 #ifdef _M_ARM32_OR_ARM64
     #define DEFAULT_CONFIG_JitES6Generators            (false)
 #else
-    #define DEFAULT_CONFIG_JitES6Generators            (true)
+    #define DEFAULT_CONFIG_JitES6Generators            (false)
 #endif
 
 #ifdef COMPILE_DISABLE_ES6RegExPrototypeProperties

lib/Runtime/Base/FunctionBody.h

@@ -1,5 +1,6 @@
 //-------------------------------------------------------------------------------------------------------
 // Copyright (C) Microsoft. All rights reserved.
+// Copyright (c) ChakraCore Project Contributors. All rights reserved.
 // Licensed under the MIT license. See LICENSE.txt file in the project root for full license information.
 //-------------------------------------------------------------------------------------------------------
 #pragma once
@@ -725,6 +726,7 @@ namespace Js
 #if ENABLE_NATIVE_CODEGEN
         Field(uint) rejitCount;
 #endif
+        Field(bool) hasYield;
         Field(bool) isNested;
         Field(bool) isInTry;
         Field(bool) isInTryFinally;
@@ -1153,8 +1155,9 @@ namespace Js
 
         bool IsJitLoopBodyPhaseEnabled() const
         {
-            // Consider: Allow JitLoopBody in generator functions for loops that do not yield.
-            return !PHASE_OFF(JITLoopBodyPhase, this) && !PHASE_OFF(FullJitPhase, this) && !this->IsCoroutine();
+            return !PHASE_OFF(JITLoopBodyPhase, this) && !PHASE_OFF(FullJitPhase, this) &&
+                (!this->IsCoroutine() || !CONFIG_FLAG(JitES6Generators) || this->IsModule());
+                // Jitting loop bodies is currently disabled when testing the Jitting of whole generator functions
         }
 
         bool IsJitLoopBodyPhaseForced() const

lib/Runtime/ByteCode/AsmJsByteCodeWriter.cpp

@@ -1,5 +1,6 @@
 //-------------------------------------------------------------------------------------------------------
 // Copyright (C) Microsoft Corporation and contributors. All rights reserved.
+// Copyright (c) ChakraCore Project Contributors. All rights reserved.
 // Licensed under the MIT license. See LICENSE.txt file in the project root for full license information.
 //-------------------------------------------------------------------------------------------------------
 
@@ -632,7 +633,7 @@ namespace Js
         uint loopId = m_functionWrite->IncrLoopCount();
         Assert((uint)m_loopHeaders->Count() == loopId);
 
-        m_loopHeaders->Add(LoopHeaderData(m_byteCodeData.GetCurrentOffset(), 0, m_loopNest > 0));
+        m_loopHeaders->Add(LoopHeaderData(m_byteCodeData.GetCurrentOffset(), 0, m_loopNest > 0, false));
         m_loopNest++;
         Js::OpCodeAsmJs loopBodyOpcode = Js::OpCodeAsmJs::AsmJsLoopBodyStart;
         this->MarkAsmJsLabel(loopEntrance);

lib/Runtime/ByteCode/ByteCodeEmitter.cpp

@@ -1,6 +1,6 @@
 //-------------------------------------------------------------------------------------------------------
 // Copyright (C) Microsoft. All rights reserved.
-// Copyright (c) 2021 ChakraCore Project Contributors. All rights reserved.
+// Copyright (c) ChakraCore Project Contributors. All rights reserved.
 // Licensed under the MIT license. See LICENSE.txt file in the project root for full license information.
 //-------------------------------------------------------------------------------------------------------
 #include "RuntimeByteCodePch.h"
@@ -10477,6 +10477,9 @@ void EmitYieldAndResume(
 
     auto* writer = byteCodeGenerator->Writer();
 
+    // If in a loop mark it as containing Yield and hence not eligible for Jit loop body
+    writer->SetCurrentLoopHasYield();
+
     if (inputReg != funcInfo->yieldRegister)
         writer->Reg2(Js::OpCode::Ld_A, funcInfo->yieldRegister, inputReg);
 

lib/Runtime/ByteCode/ByteCodeWriter.cpp

@@ -1,6 +1,6 @@
 //-------------------------------------------------------------------------------------------------------
 // Copyright (C) Microsoft. All rights reserved.
-// Copyright (c) 2021 ChakraCore Project Contributors. All rights reserved.
+// Copyright (c) ChakraCore Project Contributors. All rights reserved.
 // Licensed under the MIT license. See LICENSE.txt file in the project root for full license information.
 //-------------------------------------------------------------------------------------------------------
 #include "RuntimeByteCodePch.h"
@@ -210,7 +210,7 @@ namespace Js
             }
         }
 
-        if (this->DoJitLoopBodies() &&
+        if (this->DoJitLoopBodies() && this->HasLoopWithoutYield() &&
             !(this->m_functionWrite->GetFunctionBody()->GetHasTry() && PHASE_OFF(Js::JITLoopBodyInTryCatchPhase, this->m_functionWrite)) &&
             !(this->m_functionWrite->GetFunctionBody()->GetHasFinally() && PHASE_OFF(Js::JITLoopBodyInTryFinallyPhase, this->m_functionWrite)))
         {
@@ -252,6 +252,7 @@ namespace Js
             loopHeader->startOffset = data.startOffset;
             loopHeader->endOffset = data.endOffset;
             loopHeader->isNested = data.isNested;
+            loopHeader->hasYield = data.hasYield;
         });
     }
 
@@ -3224,7 +3225,7 @@ namespace Js
         uint loopId = m_functionWrite->IncrLoopCount();
         Assert((uint)m_loopHeaders->Count() == loopId);
 
-        m_loopHeaders->Add(LoopHeaderData(m_byteCodeData.GetCurrentOffset(), 0, m_loopNest > 0));
+        m_loopHeaders->Add(LoopHeaderData(m_byteCodeData.GetCurrentOffset(), 0, m_loopNest > 0, false));
         m_loopNest++;
         m_functionWrite->SetHasNestedLoop(m_loopNest > 1);
 
@@ -3259,6 +3260,20 @@ namespace Js
         m_loopHeaders->Item(loopId).endOffset = m_byteCodeData.GetCurrentOffset();
     }
 
+    void ByteCodeWriter::SetCurrentLoopHasYield()
+    {
+        if (m_loopNest > 0)
+        {
+            for (int i = 0; i < m_loopHeaders->Count(); ++i)
+            {
+                if (m_loopHeaders->Item(i).endOffset == 0) // check for loops we're currently inside
+                {
+                    m_loopHeaders->Item(i).hasYield = true;
+                }
+            }
+        }
+    }
+
     void ByteCodeWriter::IncreaseByteCodeCount()
     {
         m_byteCodeCount++;

lib/Runtime/ByteCode/ByteCodeWriter.h

@@ -1,6 +1,6 @@
 //-------------------------------------------------------------------------------------------------------
 // Copyright (C) Microsoft. All rights reserved.
-// Copyright (c) 2021 ChakraCore Project Contributors. All rights reserved.
+// Copyright (c) ChakraCore Project Contributors. All rights reserved.
 // Licensed under the MIT license. See LICENSE.txt file in the project root for full license information.
 //-------------------------------------------------------------------------------------------------------
 #pragma once
@@ -114,8 +114,9 @@ namespace Js
             uint startOffset;
             uint endOffset;
             bool isNested;
+            bool hasYield;
             LoopHeaderData() {}
-            LoopHeaderData(uint startOffset, uint endOffset, bool isNested) : startOffset(startOffset), endOffset(endOffset), isNested(isNested){}
+            LoopHeaderData(uint startOffset, uint endOffset, bool isNested, bool hasYield) : startOffset(startOffset), endOffset(endOffset), isNested(isNested), hasYield(hasYield){}
         };
 
         JsUtil::List<uint, ArenaAllocator> * m_labelOffsets;          // Label offsets, once defined
@@ -384,6 +385,18 @@ namespace Js
 
         uint EnterLoop(Js::ByteCodeLabel loopEntrance);
         void ExitLoop(uint loopId);
+        void SetCurrentLoopHasYield();
+        bool HasLoopWithoutYield()
+        {
+            for (int i = 0; i < m_loopHeaders->Count(); ++i)
+            {
+                if(!m_loopHeaders->Item(i).hasYield)
+                {
+                    return true;
+                }
+            }
+            return false;
+        }
 
         bool DoJitLoopBodies() const { return m_doJitLoopBodies; }
         bool DoInterruptProbes() const { return m_doInterruptProbe; }

lib/Runtime/ByteCode/WasmByteCodeWriter.cpp

@@ -1,5 +1,6 @@
 //-------------------------------------------------------------------------------------------------------
 // Copyright (C) Microsoft Corporation and contributors. All rights reserved.
+// Copyright (c) ChakraCore Project Contributors. All rights reserved.
 // Licensed under the MIT license. See LICENSE.txt file in the project root for full license information.
 //-------------------------------------------------------------------------------------------------------
 #include "RuntimeByteCodePch.h"
@@ -45,7 +46,7 @@ uint32 WasmByteCodeWriter::WasmLoopStart(ByteCodeLabel loopEntrance, __in_ecount
     uint loopId = m_functionWrite->IncrLoopCount();
     Assert((uint)m_loopHeaders->Count() == loopId);
 
-    m_loopHeaders->Add(LoopHeaderData(m_byteCodeData.GetCurrentOffset(), 0, m_loopNest > 0));
+    m_loopHeaders->Add(LoopHeaderData(m_byteCodeData.GetCurrentOffset(), 0, m_loopNest > 0, false));
     m_loopNest++;
     this->MarkAsmJsLabel(loopEntrance);
     MULTISIZE_LAYOUT_WRITE(WasmLoopStart, Js::OpCodeAsmJs::WasmLoopBodyStart, loopId, curRegs);