[Wasm RyuJit] Zero init gc shadow stack slots in the prolog

src/coreclr/jit/codegencommon.cpp

@@ -3858,7 +3858,12 @@ void CodeGen::genCheckUseBlockInit()
     // If this logic changes, Compiler::fgVarNeedsExplicitZeroInit needs
     // to be modified.
 
-#ifdef TARGET_64BIT
+#ifdef TARGET_WASM
+
+    // On WASM we always use a single memory.fill for prolog zeroing.
+    genUseBlockInit = (genInitStkLclCnt > 0);
+
+#elif defined(TARGET_64BIT)
 #if defined(TARGET_AMD64)
 
     // We can clear using aligned SIMD so the threshold is lower,
@@ -5107,7 +5112,52 @@ void CodeGen::genFnProlog()
     regMaskTP initFltRegs = RBM_NONE; // FP registers which must be init'ed.
     regMaskTP initDblRegs = RBM_NONE;
 
-#ifndef TARGET_WASM
+#ifdef TARGET_WASM
+    // For Wasm we only need to zero locals on the shadow stack; wasm locals are
+    // automatically zeroed.
+    for (unsigned varNum = 0; varNum < m_compiler->lvaCount; ++varNum)
+    {
+        LclVarDsc* varDsc = m_compiler->lvaGetDesc(varNum);
+
+        if (varDsc->lvIsParam && !varDsc->lvIsRegArg)
+        {
+            continue;
+        }
+        if (!varDsc->lvOnFrame || !varDsc->lvMustInit)
+        {
+            continue;
+        }
+        if (m_compiler->lvaIsUnknownSizeLocal(varNum))
+        {
+            continue;
+        }
+        // Defensively skip the wasm frame-header locals so we never overwrite the
+        // funclet info written by genAllocLclFrame at fp[0].
+        if ((varNum == m_compiler->lvaWasmFunctionIndex) || (varNum == m_compiler->lvaWasmVirtualIP) ||
+            (varNum == m_compiler->lvaWasmResumeIP))
+        {
+            continue;
+        }
+
+        signed int loOffs = varDsc->GetStackOffset();
+        signed int hiOffs = loOffs + m_compiler->lvaLclStackHomeSize(varNum);
+
+        hasUntrLcl = true;
+
+        if (loOffs < untrLclLo)
+        {
+            untrLclLo = loOffs;
+        }
+        if (hiOffs > untrLclHi)
+        {
+            untrLclHi = hiOffs;
+        }
+    }
+
+    // GC spill temps are not produced on wasm today.
+    assert(regSet.tmpAllFree());
+    assert(regSet.tmpListBeg() == nullptr);
+#else // !TARGET_WASM
     unsigned   varNum;
     LclVarDsc* varDsc;
 
@@ -5266,6 +5316,7 @@ void CodeGen::genFnProlog()
             untrLclHi = hiOffs;
         }
     }
+#endif // !TARGET_WASM
 
     // TODO-Cleanup: Add suitable assert for the OSR case.
     assert(m_compiler->opts.IsOSR() || ((genInitStkLclCnt > 0) == hasUntrLcl));
@@ -5281,6 +5332,7 @@ void CodeGen::genFnProlog()
     }
 #endif
 
+#ifndef TARGET_WASM
 #ifdef TARGET_ARM
     // On the ARM we will spill any incoming struct args in the first instruction in the prolog
     // Ditto for all enregistered user arguments in a varargs method.
@@ -5569,13 +5621,13 @@ void CodeGen::genFnProlog()
     }
 #endif // TARGET_ARM
 
-#ifndef TARGET_WASM // TODO-WASM: temporary; un-undefine as needed.
     //
     // Zero out the frame as needed
     //
 
     genZeroInitFrame(untrLclHi, untrLclLo, initReg, &initRegZeroed);
 
+#ifndef TARGET_WASM // TODO-WASM: enable as needed.
     genReportGenericContextArg(initReg, &initRegZeroed);
 
 #ifdef JIT32_GCENCODER

src/coreclr/jit/codegenwasm.cpp

@@ -199,6 +199,53 @@ void CodeGen::genEnregisterOSRArgsAndLocals(regNumber initReg, bool* pInitRegZer
     unreached(); // OSR not supported on WASM.
 }
 
+//------------------------------------------------------------------------
+// genZeroInitFrame: zero the lvMustInit untracked range of the frame using
+// a single memory.fill.
+//
+// Arguments:
+//    untrLclHi      - High offset (exclusive) of the untracked range relative to FP.
+//    untrLclLo      - Low offset of the untracked range relative to FP.
+//    initReg        - Unused on WASM (no scratch register concept).
+//    pInitRegZeroed - Unused on WASM.
+//
+void CodeGen::genZeroInitFrame(int untrLclHi, int untrLclLo, regNumber initReg, bool* pInitRegZeroed)
+{
+    assert(m_compiler->compGeneratingProlog);
+    if (!genUseBlockInit)
+    {
+        // Nothing to zero (genCheckUseBlockInit forces block-init for any non-empty range on wasm).
+        assert(genInitStkLclCnt == 0);
+        return;
+    }
+    assert(untrLclHi > untrLclLo);
+    assert(untrLclLo >= 0); // Wasm locals are at non-negative offsets from FP.
+
+    emitter* emit = GetEmitter();
+
+    // Push destination address: FP (+ untrLclLo if non-zero).
+    emit->emitIns_I(INS_local_get, EA_PTRSIZE, GetFramePointerRegIndex());
+    if (untrLclLo != 0)
+    {
+        emit->emitIns_I(INS_I_const, EA_PTRSIZE, untrLclLo);
+        emit->emitIns(INS_I_add);
+    }
+    // Push fill value (zero, always i32 per memory.fill spec).
+    emit->emitIns_I(INS_i32_const, EA_4BYTE, 0);
+    // Push length (pointer-sized: i32 on wasm32, i64 on wasm64).
+    emit->emitIns_I(INS_I_const, EA_PTRSIZE, untrLclHi - untrLclLo);
+    // memory.fill 0
+    emit->emitIns_I(INS_memory_fill, EA_4BYTE, LINEAR_MEMORY_INDEX);
+}
+
+//------------------------------------------------------------------------
+// genZeroInitFrameUsingBlockInit: Unused on WASM (genZeroInitFrame handles block init inline).
+//
+void CodeGen::genZeroInitFrameUsingBlockInit(int untrLclHi, int untrLclLo, regNumber initReg, bool* pInitRegZeroed)
+{
+    unreached();
+}
+
 //------------------------------------------------------------------------
 // genOSRHandleTier0CalleeSavedRegistersAndFrame:
 //   Not called for WASM without OSR support.

src/coreclr/jit/compiler.hpp

@@ -4354,6 +4354,11 @@ bool Compiler::fgVarNeedsExplicitZeroInit(unsigned varNum, bool bbInALoop, bool
         // Below conditions guarantee block initialization, which will initialize
         // all struct fields. If the logic for block initialization in CodeGen::genCheckUseBlockInit()
         // changes, these conditions need to be updated.
+#ifdef TARGET_WASM
+        // On WASM the prolog always uses a single memory.fill to zero any
+        // locals that need initialization, regardless of size.
+        return false;
+#else // !TARGET_WASM
         unsigned stackHomeSize = lvaLclStackHomeSize(varNum);
 #ifdef TARGET_64BIT
 #if defined(TARGET_AMD64)
@@ -4369,6 +4374,7 @@ bool Compiler::fgVarNeedsExplicitZeroInit(unsigned varNum, bool bbInALoop, bool
         {
             return false;
         }
+#endif // !TARGET_WASM
     }
 
     return !info.compInitMem || (varDsc->lvIsTemp && !varDsc->HasGCPtr());