[Wasm RyuJIT] Spill live ref/byref values to pinned stack slots at calls

src/coreclr/jit/compiler.cpp

@@ -5053,6 +5053,11 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
     // keep the Virtual IP updated.
     //
     DoPhase(this, PHASE_WASM_VIRTUAL_IP, &Compiler::fgWasmVirtualIP);
+
+    // Ensure that any refs or byrefs live at call sites are spilled
+    // to pinned stack slots so the objects aren't moved.
+    //
+    DoPhase(this, PHASE_WASM_SPILL_REFS, &Compiler::WasmSpillRefs);
 #endif
 
     FinalizeEH();

src/coreclr/jit/compiler.h

@@ -4244,6 +4244,7 @@ class Compiler
     unsigned lvaWasmVirtualIP = BAD_VAR_NUM; // Wasm virtual IP slot
     unsigned lvaWasmFunctionIndex = BAD_VAR_NUM; // Wasm function index slot
     unsigned lvaWasmResumeIP = BAD_VAR_NUM; // Wasm catch resumption IP slot
+    jitstd::vector<unsigned>* m_wasmSpillSlots = nullptr;
 #endif // defined(TARGET_WASM)
 
     unsigned lvaInlinedPInvokeFrameVar = BAD_VAR_NUM; // variable representing the InlinedCallFrame
@@ -4646,6 +4647,7 @@ class Compiler
     unsigned lvaTrackedIndexToLclNum(unsigned trackedIndex)
     {
         assert(trackedIndex < lvaTrackedCount);
+        assert(trackedIndex < lvaTrackedToVarNumSize);
         unsigned lclNum = lvaTrackedToVarNum[trackedIndex];
         assert(lclNum < lvaCount);
         return lclNum;
@@ -6757,6 +6759,7 @@ class Compiler
     PhaseStatus fgWasmControlFlow();
     PhaseStatus fgWasmTransformSccs();
     PhaseStatus fgWasmVirtualIP();
+    PhaseStatus WasmSpillRefs();
 #ifdef DEBUG
     void fgDumpWasmControlFlow();
     void fgDumpWasmControlFlowDot();

src/coreclr/jit/compmemkind.h

@@ -73,6 +73,7 @@ CompMemKindMacro(RangeCheckCloning)
 CompMemKindMacro(WasmSccTransform)
 CompMemKindMacro(WasmCfgLowering)
 CompMemKindMacro(WasmEH)
+CompMemKindMacro(WasmSpillRefs)
 //clang-format on
 
 #undef CompMemKindMacro

src/coreclr/jit/compphases.h

@@ -130,6 +130,7 @@ CompPhaseNameMacro(PHASE_DFS_BLOCKS_WASM,            "Wasm remove unreachable bl
 CompPhaseNameMacro(PHASE_WASM_EH_FLOW,               "Wasm eh control flow",           false, -1, false)
 CompPhaseNameMacro(PHASE_WASM_TRANSFORM_SCCS,        "Wasm transform sccs",            false, -1, false)
 CompPhaseNameMacro(PHASE_WASM_CONTROL_FLOW,          "Wasm control flow",              false, -1, false)
+CompPhaseNameMacro(PHASE_WASM_SPILL_REFS,            "Wasm spill refs",                false, -1, false)
 CompPhaseNameMacro(PHASE_WASM_VIRTUAL_IP,            "Wasm virtual IP",                false, -1, false)
 
 CompPhaseNameMacro(PHASE_ASYNC,                      "Transform async",                false, -1, true)

src/coreclr/jit/fgwasm.cpp

@@ -1671,6 +1671,137 @@ PhaseStatus Compiler::fgWasmControlFlow()
     return PhaseStatus::MODIFIED_EVERYTHING;
 }
 
+PhaseStatus Compiler::WasmSpillRefs()
+{
+    bool anyChanges = false;
+
+    size_t highWaterMark = 0;
+    jitstd::vector<GenTree*> defs(getAllocator(CMK_WasmSpillRefs));
+
+    for (BasicBlock* const block : Blocks())
+    {
+        // LIR edges cannot span blocks, so we can safely clear the list of live values per-block
+        defs.clear();
+
+        for (GenTree* tree : LIR::AsRange(block))
+        {
+            if (tree->IsCall())
+            {
+                highWaterMark = std::max(highWaterMark, defs.size());
+
+                // For any ref/byref values live at the point of a call, spill them into pinned slots
+                //  on the stack where the GC can see them so it won't move them.
+                if (defs.size())
+                {
+                    anyChanges = true;
+
+                    if (!m_wasmSpillSlots)
+                    {
+                        m_wasmSpillSlots = new (this, CMK_WasmSpillRefs) jitstd::vector<unsigned>(getAllocator(CMK_WasmSpillRefs));
+                    }
+
+                    unsigned spillSlotIndex = 0;
+                    JITDUMP("Spilling %zu live ref(s) for call\n", defs.size());
+                    DISPNODE(tree);
+                    for (GenTree* def : defs)
+                    {
+                        JITDUMP("    ");
+                        DISPNODE(def);
+
+                        unsigned spillSlot;
+                        if (spillSlotIndex < m_wasmSpillSlots->size())
+                        {
+                            spillSlot = m_wasmSpillSlots->at(spillSlotIndex);
+                        }
+                        else
+                        {
+                            spillSlot = lvaGrabTemp(false DEBUGARG("WasmSpillRefs spill slot"));
+                            LclVarDsc* const varDsc = lvaGetDesc(spillSlot);
+                            varDsc->lvType = TYP_BYREF;
+                            varDsc->lvPinned = true;
+                            varDsc->lvImplicitlyReferenced = true;
+                            varDsc->lvMustInit = true;
+                            lvaSetVarDoNotEnregister(spillSlot, DoNotEnregisterReason::WasmGCVisibility);
+                            m_wasmSpillSlots->push_back(spillSlot);
+                        }
+                        spillSlotIndex++;
+
+                        GenTreeLclVar *spill = gtNewStoreLclVarNode(spillSlot, def);
+                        GenTreeLclVar *reload = gtNewLclVarNode(spillSlot);
+                        LIR::Use use;
+                        noway_assert(LIR::AsRange(block).TryGetUse(def, &use));
+                        use.ReplaceWith(reload);
+                        LIR::AsRange(block).InsertAfter(def, spill);
+                        LIR::AsRange(block).InsertAfter(spill, reload);
+                        if (def->gtLIRFlags & LIR::Flags::MultiplyUsed)
+                        {
+                            JITDUMP("Transferring multiply-used flag from [%06u] to [%06u] for spill\n", Compiler::dspTreeID(def), Compiler::dspTreeID(reload));
+                            def->gtLIRFlags &= ~LIR::Flags::MultiplyUsed;
+                            reload->gtLIRFlags |= LIR::Flags::MultiplyUsed;
+                        }
+                        anyChanges = true;
+                    }
+
+                    defs.clear();
+                }
+            }
+
+            // FIXME: Should this happen before the spilling of the live defs list?
+            // I think the answer is no, because live defs being passed as arguments to the current call
+            //  are not guaranteed to ever end up in memory where the GC can see them unless we spill
+            //  them. If we can somehow guarantee that all callees will spill their ref parameters
+            //  immediately, we could do this before the block above.
+
+            // Remove used nodes from defs list, they're no longer meaningfully 'live'.
+            tree->VisitOperands([&defs](GenTree* op) {
+                if (!op->IsValue())
+                    return GenTree::VisitResult::Continue;
+                if (!op->TypeIs(TYP_REF, TYP_BYREF))
+                    return GenTree::VisitResult::Continue;
+
+                for (size_t i = defs.size(); i > 0; i--)
+                {
+                    if (op == defs[i - 1])
+                    {
+                        defs[i - 1] = defs[defs.size() - 1];
+                        defs.pop_back();
+                        break;
+                    }
+                }
+
+                return GenTree::VisitResult::Continue;
+            });
+
+            // We only care about used values, and invariant nodes can't produce movable GC refs, so skip
+            //  nodes appropriately
+            if (!tree->IsValue() || tree->IsUnusedValue() || tree->IsInvariant())
+            {
+                continue;
+            }
+
+            // If a value is just a GT_LCL_VAR that isn't address-exposed, by construction we ensure that
+            //  it won't be mutated between its def (here) and its use (the call that would produce a spill)
+            //  and we won't need to spill it.
+            if (tree->OperIs(GT_LCL_VAR))
+            {
+                LclVarDsc* dsc = lvaGetDesc(tree->AsLclVarCommon());
+                if (!dsc->IsAddressExposed())
+                    continue;
+            }
+
+            // We have a ref sourced from something like a call result or an indirection that hasn't been
+            //  spilled yet, so record it for potential spilling at the next call.
+            if (tree->TypeIs(TYP_REF, TYP_BYREF))
+            {
+                defs.push_back(tree);
+            }
+        }
+    }
+
+    JITDUMP("High water mark for refs was %zu\n", highWaterMark);
+    return anyChanges ? PhaseStatus::MODIFIED_EVERYTHING : PhaseStatus::MODIFIED_NOTHING;
+}
+
 #ifdef DEBUG
 
 //------------------------------------------------------------------------

src/coreclr/jit/gentree.cpp

@@ -13126,6 +13126,20 @@ void Compiler::gtGetLclVarNameInfo(unsigned lclNum, const char** ilKindOut, cons
     const char* ilName = nullptr;
 
     unsigned ilNum = compMap2ILvarNum(lclNum);
+#if TARGET_WASM
+    int wasmSpillSlotIndex = -1;
+    if (m_wasmSpillSlots)
+    {
+        for (unsigned i = 0; i < m_wasmSpillSlots->size(); i++)
+        {
+            if (m_wasmSpillSlots->at(i) == lclNum)
+            {
+                wasmSpillSlotIndex = i;
+                break;
+            }
+        }
+    }
+#endif
 
     if (ilNum == (unsigned)ICorDebugInfo::RETBUF_ILNUM)
     {
@@ -13200,6 +13214,19 @@ void Compiler::gtGetLclVarNameInfo(unsigned lclNum, const char** ilKindOut, cons
             {
                 ilName = "SP";
             }
+            else if (lclNum == lvaWasmVirtualIP)
+            {
+                ilName = "VirtualIP";
+            }
+            else if (lclNum == lvaWasmFunctionIndex)
+            {
+                ilName = "FuncIndex";
+            }
+            else if (wasmSpillSlotIndex > -1)
+            {
+                ilKind = "spill";
+                ilNum = wasmSpillSlotIndex;
+            }
 #endif // defined(TARGET_WASM)
             else
             {