Reduce memory imprint by saving only the current register for each types when entering a loop.

Use a different opcode to discern yields from normal load then check with the saved register to determine if we are yielding outside of the loop.

lib/Backend/IRBuilderAsmJs.cpp

@@ -100,34 +100,13 @@ IRBuilderAsmJs::Build()
         // those as if they are local for yielding
         this->m_ldSlots = BVFixed::New<JitArenaAllocator>(GetLastTmp(WAsmJs::LastType), m_tempAlloc);
         this->m_stSlots = BVFixed::New<JitArenaAllocator>(GetLastTmp(WAsmJs::LastType), m_tempAlloc);
-        uint count = m_func->GetWorkItem()->GetWorkItemData()->wasmLoopYieldRegCount;
-        if (count > 0)
-        {
+        this->m_yieldRegs = nullptr;
 #ifdef ENABLE_WASM
-            AssertOrFailFast(m_func->GetJITFunctionBody()->IsWasmFunction());
+        if (m_func->GetJITFunctionBody()->IsWasmFunction())
+        {
             this->m_yieldRegs = BVFixed::New<JitArenaAllocator>(GetLastTmp(WAsmJs::LastType), m_tempAlloc);
-            WasmRegisterInfoIDL* infos = m_func->GetWorkItem()->GetWorkItemData()->wasmLoopYieldRegs;
-            AssertOrFailFast(infos);
-            for (uint i = 0; i < count; ++i)
-            {
-                Js::RegSlot loc = infos[i].location;
-                CompileAssert(sizeof(Wasm::WasmTypes::WasmType) == sizeof(infos[i].type));
-                Wasm::WasmTypes::WasmType type = (Wasm::WasmTypes::WasmType)infos[i].type;
-                Js::RegSlot regSlot = 0;
-                switch (type)
-                {
-                case Wasm::WasmTypes::I32: regSlot = GetRegSlotFromTypedReg(loc, WAsmJs::INT32); break;
-                case Wasm::WasmTypes::I64: regSlot = GetRegSlotFromTypedReg(loc, WAsmJs::INT64); break;
-                case Wasm::WasmTypes::F32: regSlot = GetRegSlotFromTypedReg(loc, WAsmJs::FLOAT32); break;
-                case Wasm::WasmTypes::F64: regSlot = GetRegSlotFromTypedReg(loc, WAsmJs::FLOAT64); break;
-                default:
-                    Assert(UNREACHED);
-                    break;
-                }
-                this->m_yieldRegs->Set(regSlot);
-            }
-#endif
         }
+#endif
         this->m_loopBodyRetIPSym = StackSym::New(TyInt32, this->m_func);
 #if DBG
         if (m_tempCount > 0)
@@ -286,7 +265,11 @@ IR::RegOpnd *
 IRBuilderAsmJs::BuildDstOpnd(Js::RegSlot dstRegSlot, IRType type)
 {
     SymID symID = static_cast<SymID>(dstRegSlot);
-    if (IsLoopBody() && (RegIsVar(dstRegSlot) || RegIsYield(dstRegSlot)))
+    if (RegIsYield(dstRegSlot))
+    {
+        // Use default symID for yield registers
+    }
+    else if (IsLoopBody() && RegIsVar(dstRegSlot))
     {
         // Loop body and not constants
         this->m_stSlots->Set(dstRegSlot);
@@ -383,7 +366,7 @@ IRBuilderAsmJs::BuildIntConstOpnd(Js::RegSlot regSlot)
     Assert(info.constSrcByteOffset != Js::Constants::InvalidOffset);
     AssertOrFailFast(info.constSrcByteOffset < UInt32Math::Mul<sizeof(Js::Var)>(m_func->GetJITFunctionBody()->GetConstCount()));
     int* intConstTable = reinterpret_cast<int*>(((byte*)constTable) + info.constSrcByteOffset);
-    Js::RegSlot srcReg = GetTypedRegFromRegSlot(regSlot, WAsmJs::INT32);
+    uint32 srcReg = GetTypedRegFromRegSlot(regSlot, WAsmJs::INT32);
     AssertOrFailFast(srcReg >= Js::FunctionBody::FirstRegSlot && srcReg < info.constCount);
     const int32 value = intConstTable[srcReg];
     IR::IntConstOpnd *opnd = IR::IntConstOpnd::New(value, TyInt32, m_func);
@@ -464,7 +447,7 @@ IRBuilderAsmJs::AddStatementBoundary(uint statementIndex, uint offset)
     return m_statementReader.MoveNextStatementBoundary();
 }
 
-Js::RegSlot IRBuilderAsmJs::GetTypedRegFromRegSlot(Js::RegSlot reg, WAsmJs::Types type)
+uint32 IRBuilderAsmJs::GetTypedRegFromRegSlot(Js::RegSlot reg, WAsmJs::Types type)
 {
     const auto typedInfo = m_asmFuncInfo->GetTypedSlotInfo(type);
     Js::RegSlot srcReg = reg;
@@ -652,6 +635,32 @@ IRBuilderAsmJs::RegIsYield(Js::RegSlot reg)
     return m_yieldRegs && m_yieldRegs->Test(reg);
 }
 
+void
+IRBuilderAsmJs::CheckIsYieldOutsideOfLoop(Js::RegSlot reg, IRType type)
+{
+#ifdef ENABLE_WASM
+    if (m_yieldRegs && !m_yieldRegs->Test(reg))
+    {
+        Assert(m_func->IsLoopBody());
+        WAsmJs::Types wasmType = WAsmJs::FromIRType(type);
+        Assert(wasmType < WAsmJs::LIMIT);
+        uint32 typedReg = GetTypedRegFromRegSlot(reg, wasmType);
+        if (typedReg < m_func->GetWorkItem()->GetWorkItemData()->wasmLoopYieldMins[wasmType])
+        {
+            Assert(!RegIsTemp(reg) || GetMappedTemp(reg) == (SymID)reg || GetMappedTemp(reg) == 0);
+            m_yieldRegs->Set(reg);
+            // Make sure we set that slot when done with the loop
+            m_stSlots->Set(reg);
+            // We need to make sure that the symbols is loaded as well
+            // so that the sym will be defined on all path.
+            this->EnsureLoopBodyAsmJsLoadSlot(reg, type);
+        }
+    }
+#else
+    Assert(!m_yieldRegs);
+#endif
+}
+
 bool
 IRBuilderAsmJs::RegIsSimd128ReturnVar(Js::RegSlot reg)
 {
@@ -2434,12 +2443,7 @@ IRBuilderAsmJs::BuildInt2(Js::OpCodeAsmJs newOpcode, uint32 offset, Js::RegSlot
 
     case Js::OpCodeAsmJs::Return_Int:
         instr = IR::Instr::New(Js::OpCode::Ld_I4, dstOpnd, srcOpnd, m_func);
-        if (m_func->IsLoopBody())
-        {
-            // Make sure we set that slot when done with the loop
-            this->m_stSlots->Set(dstRegSlot);
-        }
-
+        CheckIsYieldOutsideOfLoop(dstRegSlot, TyInt32);
         break;
 
     case Js::OpCodeAsmJs::Eqz_Int:
@@ -2700,11 +2704,7 @@ IRBuilderAsmJs::BuildDouble2(Js::OpCodeAsmJs newOpcode, uint32 offset, Js::RegSl
         break;
     case Js::OpCodeAsmJs::Return_Db:
         instr = IR::Instr::New(Js::OpCode::Ld_A, dstOpnd, srcOpnd, m_func);
-        if (m_func->IsLoopBody())
-        {
-            // Make sure we set that slot when done with the loop
-            this->m_stSlots->Set(dstRegSlot);
-        }
+        CheckIsYieldOutsideOfLoop(dstRegSlot, TyFloat64);
         break;
     case Js::OpCodeAsmJs::Trunc_Db:
         instr = IR::Instr::New(Js::OpCode::Trunc_A, dstOpnd, srcOpnd, m_func);
@@ -2753,11 +2753,7 @@ IRBuilderAsmJs::BuildFloat2(Js::OpCodeAsmJs newOpcode, uint32 offset, Js::RegSlo
         break;
     case Js::OpCodeAsmJs::Return_Flt:
         instr = IR::Instr::New(Js::OpCode::Ld_A, dstOpnd, srcOpnd, m_func);
-        if (m_func->IsLoopBody())
-        {
-            // Make sure we set that slot when done with the loop
-            this->m_stSlots->Set(dstRegSlot);
-        }
+        CheckIsYieldOutsideOfLoop(dstRegSlot, TyFloat32);
         break;
     case Js::OpCodeAsmJs::Trunc_Flt:
         instr = IR::Instr::New(Js::OpCode::Trunc_A, dstOpnd, srcOpnd, m_func);
@@ -3105,11 +3101,7 @@ IRBuilderAsmJs::BuildLong2(Js::OpCodeAsmJs newOpcode, uint32 offset, Js::RegSlot
 
     case Js::OpCodeAsmJs::Return_Long:
         instr = IR::Instr::New(Js::OpCode::Ld_I4, dstOpnd, srcOpnd, m_func);
-        if (m_func->IsLoopBody())
-        {
-            // Make sure we set that slot when done with the loop
-            this->m_stSlots->Set(dstRegSlot);
-        }
+        CheckIsYieldOutsideOfLoop(dstRegSlot, TyInt64);
         break;
     case Js::OpCodeAsmJs::I64Extend8_s:
         instr = CreateSignExtendInstr(dstOpnd, srcOpnd, TyInt8);
@@ -6298,13 +6290,7 @@ void IRBuilderAsmJs::BuildSimd_2(Js::OpCodeAsmJs newOpcode, uint32 offset, Js::R
     case Js::OpCodeAsmJs::Simd128_Return_B4:
     case Js::OpCodeAsmJs::Simd128_Return_B8:
     case Js::OpCodeAsmJs::Simd128_Return_B16:
-        if (m_func->IsLoopBody())
-        {
-            // Make sure we set that slot when done with the loop
-            this->m_stSlots->Set(dstRegSlot);
-        }
-        opcode = Js::OpCode::Ld_A;
-        break;
+        CheckIsYieldOutsideOfLoop(dstRegSlot, simdType);
     case Js::OpCodeAsmJs::Simd128_Ld_F4:
     case Js::OpCodeAsmJs::Simd128_Ld_I4:
     case Js::OpCodeAsmJs::Simd128_Ld_I8:

lib/Backend/IRBuilderAsmJs.h

@@ -87,7 +87,7 @@ private:
 #if DBG
     BVFixed *               m_usedAsTemp;
 #endif
-    Js::RegSlot             GetTypedRegFromRegSlot(Js::RegSlot reg, WAsmJs::Types type);
+    uint32                  GetTypedRegFromRegSlot(Js::RegSlot reg, WAsmJs::Types type);
     Js::RegSlot             GetRegSlotFromTypedReg(Js::RegSlot srcReg, WAsmJs::Types type);
     Js::RegSlot             GetRegSlotFromPtrReg(Js::RegSlot srcReg)
     {
@@ -121,6 +121,7 @@ private:
     bool                    RegIsTypedTmp(Js::RegSlot reg, WAsmJs::Types type);
     bool                    RegIs(Js::RegSlot reg, WAsmJs::Types type);
     bool                    RegIsYield(Js::RegSlot reg);
+    void                    CheckIsYieldOutsideOfLoop(Js::RegSlot reg, IRType type);
 
     void                    BuildArgOut(IR::Opnd* srcOpnd, uint32 dstRegSlot, uint32 offset, IRType type, ValueType valueType = ValueType::Uninitialized);
     void                    BuildFromVar(uint32 offset, Js::RegSlot dstRegSlot, Js::RegSlot srcRegSlot, IRType irType, ValueType valueType);

lib/Backend/NativeCodeGenerator.cpp

@@ -683,18 +683,9 @@ void NativeCodeGenerator::GenerateLoopBody(Js::FunctionBody * fn, Js::LoopHeader
             CodeGenWorkItemIDL* data = workitem->GetJITData();
             Js::AsmJsFunctionInfo* asmInfo = fn->GetAsmJsFunctionInfoWithLock();
             AssertOrFailFast(data->loopNumber < (uint)asmInfo->GetWasmLoopsYieldInfo()->Count());
-            Js::WasmLoopYieldInfo* loopYieldInfo = asmInfo->GetWasmLoopsYieldInfo()->Item(data->loopNumber);
-            if (loopYieldInfo)
-            {
-                data->wasmLoopYieldRegCount = loopYieldInfo->Count();
-                data->wasmLoopYieldRegs = (WasmRegisterInfoIDL*)loopYieldInfo->GetBuffer();
-            }
-            else
-            {
-                Assert(UNREACHED);
-                data->wasmLoopYieldRegCount = 0;
-                data->wasmLoopYieldRegs = nullptr;
-            }
+            Js::WasmLoopYieldInfo& loopYieldInfo = asmInfo->GetWasmLoopsYieldInfo()->Item(data->loopNumber);
+            CompileAssert(sizeof(data->wasmLoopYieldMins) == sizeof(loopYieldInfo.minYield));
+            memcpy_s(data->wasmLoopYieldMins, sizeof(data->wasmLoopYieldMins), loopYieldInfo.minYield, sizeof(loopYieldInfo.minYield));
         }
 #endif
 

lib/JITIDL/JITTypes.h

@@ -418,12 +418,6 @@ typedef struct WasmSignatureIDL
     IDL_DEF([size_is(paramsCount)]) int * params;
 } WasmSignatureIDL;
 
-typedef struct WasmRegisterInfoIDL
-{
-    unsigned int location;
-    unsigned int type;
-} WasmRegisterInfoIDL;
-
 typedef struct TypedSlotInfo
 {
     boolean isValidType;
@@ -703,7 +697,8 @@ typedef struct CodeGenWorkItemIDL
     unsigned int loopNumber;
     unsigned int inlineeInfoCount;
     unsigned int symIdToValueTypeMapCount;
-    unsigned int wasmLoopYieldRegCount;
+    // This is used only by WebAssembly Jit loop body
+    unsigned int wasmLoopYieldMins[5];
     XProcNumberPageSegment * xProcNumberPageSegment;
 
     PolymorphicInlineCacheInfoIDL * selfInfo;
@@ -712,8 +707,6 @@ typedef struct CodeGenWorkItemIDL
 
     // TODO: OOP JIT, move loop body data to separate struct
     IDL_DEF([size_is(symIdToValueTypeMapCount)]) unsigned short * symIdToValueTypeMap;
-    // This is used only by WebAssembly Jit loop body
-    IDL_DEF([size_is(wasmLoopYieldRegCount)]) WasmRegisterInfoIDL * wasmLoopYieldRegs;
 
     FunctionJITTimeDataIDL * jitData;
     CHAKRA_PTR jittedLoopIterationsSinceLastBailoutAddr;

lib/Runtime/Language/AsmJsTypes.h

@@ -896,8 +896,11 @@ namespace Js
         }
     };
 
-    typedef JsUtil::List<WasmRegisterInfoIDL> WasmLoopYieldInfo;
-    typedef JsUtil::List<WasmLoopYieldInfo*> WasmLoopsYieldInfo;
+    struct WasmLoopYieldInfo
+    {
+        uint32 minYield[WAsmJs::LIMIT] = {0,0,0,0,0};
+    };
+    typedef JsUtil::List<WasmLoopYieldInfo> WasmLoopsYieldInfo;
     class AsmJsFunctionInfo
     {
         Field(WAsmJs::TypedSlotInfo) mTypedSlotInfos[WAsmJs::LIMIT];

lib/Runtime/Language/WAsmjsUtils.h

@@ -140,6 +140,13 @@ namespace WAsmJs
             return mNextLocation++;
         }
 
+        RegSlot PeekNextTmpRegister()
+        {
+            // Make sure this function is called correctly
+            Assert(mNextLocation <= mRegisterCount && mNextLocation >= mFirstTmpReg);
+            return mNextLocation;
+        }
+
         // Release a location for a temporary register, must be the last location acquired
         void ReleaseTmpRegister( RegSlot tmpReg )
         {

lib/WasmReader/WasmByteCodeGenerator.cpp

@@ -944,17 +944,16 @@ EmitInfo WasmBytecodeGenerator::EmitLoop()
 
     // Save all the possible blocks the loop could yield to
     Js::WasmLoopYieldInfo* loopYieldInfo = EnsureWasmLoopYieldInfo(loopId);
-    m_blockInfos.Map([loopYieldInfo](int i, BlockInfo block)
+    for (WAsmJs::Types type = WAsmJs::Types(0); type != WAsmJs::LIMIT; type = WAsmJs::Types(type + 1))
     {
-        if (block.HasYield())
+        uint32 minYield = 0;
+        if (!mTypedRegisterAllocator.IsTypeExcluded(type))
         {
-            EmitInfo info = block.yieldInfo->info;
-            CompileAssert(sizeof(WasmRegisterInfoIDL) == sizeof(EmitInfo));
-            CompileAssert(offsetof(WasmRegisterInfoIDL, location) == offsetof(EmitInfo, location));
-            CompileAssert(offsetof(WasmRegisterInfoIDL, type) == offsetof(EmitInfo, type));
-            loopYieldInfo->Add(*(WasmRegisterInfoIDL*)&info);
+            CompileAssert(sizeof(minYield) == sizeof(Js::RegSlot));
+            minYield = static_cast<uint32>(mTypedRegisterAllocator.GetRegisterSpace(type)->PeekNextTmpRegister());
         }
-    });
+        loopYieldInfo->minYield[type] = minYield;
+    }
 
     // Internally we create a block for loop to exit, but semantically, they don't exist so pop it
     m_blockInfos.Pop();
@@ -1614,7 +1613,7 @@ void WasmBytecodeGenerator::YieldToBlock(BlockInfo blockInfo, EmitInfo expr)
         if (!IsUnreachable())
         {
             blockInfo.yieldInfo->didYield = true;
-            m_writer->AsmReg2(GetLoadOp(expr.type), yieldInfo.location, expr.location);
+            m_writer->AsmReg2(GetReturnOp(expr.type), yieldInfo.location, expr.location);
         }
     }
 }
@@ -1651,16 +1650,10 @@ Js::WasmLoopYieldInfo* WasmBytecodeGenerator::EnsureWasmLoopYieldInfo(uint32 _lo
     int32 oldCount = loopsInfo->Count();
     if (loopId >= oldCount)
     {
-        loopsInfo->SetItem(loopId, RecyclerNew(m_module->GetRecycler(), Js::WasmLoopYieldInfo, m_module->GetRecycler()));
-
-        // Initialize any new item in the list to nullptr
-        for (int32 count = oldCount; count < loopId; ++count)
-        {
-            loopsInfo->SetItem(count, nullptr);
-        }
+        loopsInfo->SetItem(loopId, Js::WasmLoopYieldInfo());
     }
 
-    return loopsInfo->Item(loopId);
+    return &loopsInfo->Item(loopId);
 }
 
 WasmRegisterSpace* WasmBytecodeGenerator::GetRegisterSpace(WasmTypes::WasmType type)

test/wasm/loopyieldregress.js

@@ -0,0 +1,89 @@
+//-------------------------------------------------------------------------------------------------------
+// Copyright (C) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE.txt file in the project root for full license information.
+//-------------------------------------------------------------------------------------------------------
+function getBuf() {
+  if (WebAssembly.wabt) {
+    const buf = WebAssembly.wabt.convertWast2Wasm(`(module
+      (memory 1 1)
+      (func $_main (export "foo") (param i32 i32) (result i32)
+      (local i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i32 i64)
+      block
+      block
+      loop  ;; label = @5
+        block  ;; label = @6
+          get_local 3
+          i32.const -3
+          i32.and
+          set_local 3
+          get_local 3
+          i32.const 16
+          i32.shl
+          i32.const 16
+          i32.shr_s
+          set_local 1
+          get_local 0
+          get_local 1
+          i32.lt_s
+          if (result i32)  ;; label = @7
+            get_local 4
+            i32.load8_s
+            set_local 1
+            get_local 1
+            i32.const 1
+            i32.ne
+            br_if 1 (;@6;)
+            get_local 9
+            i32.const -2115251625
+            i32.store
+            i32.const -2115251625
+            set_local 1
+            i32.const 1
+          else
+            get_local 7
+            i32.const 0
+            i32.store
+            get_local 0
+            set_local 1
+            get_local 3
+          end
+          set_local 0
+          i32.const 22112
+          get_local 6
+          i32.const 1
+          i32.add
+          tee_local 6
+          i32.store
+          get_local 6
+          i32.const 20
+          i32.ge_s
+          br_if 2 (;@4;)
+          get_local 0
+          set_local 3
+          get_local 1
+          set_local 0
+          br 1 (;@5;)
+        end
+      end
+      end
+      end
+      (i32.const 0)
+    )
+    )`);
+    //const view = new Uint8Array(buf);
+    //console.log(view.join(","));
+    return buf;
+  } else {
+    const arr = [0,97,115,109,1,0,0,0,1,7,1,96,2,127,127,1,127,3,2,1,0,5,4,1,1,1,1,7,7,1,3,102,111,111,0,0,10,131,1,1,128,1,2,158,1,127,1,126,2,64,2,64,3,64,2,64,32,3,65,125,113,33,3,32,3,65,16,116,65,16,117,33,1,32,0,32,1,72,4,127,32,4,44,0,0,33,1,32,1,65,1,71,13,1,32,9,65,215,164,175,143,120,54,2,0,65,215,164,175,143,120,33,1,65,1,5,32,7,65,0,54,2,0,32,0,33,1,32,3,11,33,0,65,224,172,1,32,6,65,1,106,34,6,54,2,0,32,6,65,20,78,13,2,32,0,33,3,32,1,33,0,12,1,11,11,11,11,65,0,11]
+    const buf = new ArrayBuffer(arr.length);
+    const view = new Uint8Array(buf);
+    for (let i = 0; i < arr.length; ++i) {
+      view[i] = arr[i];
+    }
+    return buf;
+  }
+}
+
+const mod = new WebAssembly.Module(getBuf());
+const {exports} = new WebAssembly.Instance(mod);
+console.log(exports.foo() == 0 ? "pass" : "failed");

test/wasm/rlexe.xml

@@ -391,4 +391,11 @@
     <tags>exclude_jshost,exclude_win7,exclude_xplat</tags>
   </default>
 </test>
+<test>
+  <default>
+    <files>loopyieldregress.js</files>
+    <compile-flags>-lic:0 -bgjit-</compile-flags>
+    <tags>exclude_win7,exclude_xplat</tags>
+  </default>
+</test>
 </regress-exe>