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

src/coreclr/jit/compiler.cpp

@@ -5049,6 +5049,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::fgWasmSpillRefs);
 #endif
 
  FinalizeEH();

src/coreclr/jit/compiler.h

@@ -4259,6 +4259,13 @@ class Compiler
  unsigned lvaWasmFunctionIndex = BAD_VAR_NUM; // Wasm function index slot
  unsigned lvaWasmResumeIP = BAD_VAR_NUM; // Wasm catch resumption IP slot
 
+ struct WasmSpillSlot
+ {
+ unsigned lclNum;
+ bool byRef;
+ bool inUse;
+ };
+ jitstd::vector<WasmSpillSlot>* m_wasmSpillSlots = nullptr;
 #endif // defined(TARGET_WASM)
 
  unsigned lvaInlinedPInvokeFrameVar = BAD_VAR_NUM; // variable representing the InlinedCallFrame
@@ -4661,6 +4668,7 @@ class Compiler
  unsigned lvaTrackedIndexToLclNum(unsigned trackedIndex)
  {
  assert(trackedIndex < lvaTrackedCount);
+ assert(trackedIndex < lvaTrackedToVarNumSize);
  unsigned lclNum = lvaTrackedToVarNum[trackedIndex];
  assert(lclNum < lvaCount);
  return lclNum;
@@ -6770,6 +6778,7 @@ class Compiler
  PhaseStatus fgWasmControlFlow();
  PhaseStatus fgWasmTransformSccs();
  PhaseStatus fgWasmVirtualIP();
+ PhaseStatus fgWasmSpillRefs();
 #ifdef DEBUG
  void fgDumpWasmControlFlow();
  void fgDumpWasmControlFlowDot();

src/coreclr/jit/compmemkind.h

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

src/coreclr/jit/compphases.h

@@ -129,6 +129,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

@@ -1759,6 +1759,189 @@ PhaseStatus Compiler::fgWasmControlFlow()
  return PhaseStatus::MODIFIED_EVERYTHING;
 }
 
+PhaseStatus Compiler::fgWasmSpillRefs()
+{
+ bool anyChanges = false;
+
+ jitstd::vector<GenTree*> defs(getAllocator(CMK_WasmSpillRefs));
+ jitstd::vector<unsigned> spillSlotsToZeroAtEndOfBlock(getAllocator(CMK_WasmSpillRefs));
+
+ for (BasicBlock* const block : Blocks())
+ {
+ // The defs list should already be empty because we walked all the nodes in the previous block,
+ // which should have led to visiting all uses
+ assert(defs.empty());
+
+ // LIR edges cannot span blocks, so we can safely clear the list of live values per-block
+ defs.clear();
+
+ if (m_wasmSpillSlots != nullptr)
+ {
+ // Flag all our spill slot vars as no longer in use so they can be reused in the new block
+ for (WasmSpillSlot& slot : *m_wasmSpillSlots)
+ {
+ slot.inUse = false;
+ }
+ }
+
+ for (GenTree* tree : LIR::AsRange(block))
+ {
+ if (tree->IsCall())
+ {
+ // 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.empty())
+ {
+ if (m_wasmSpillSlots == nullptr)
+ {
+ m_wasmSpillSlots = new (this, CMK_WasmSpillRefs)
+ jitstd::vector<WasmSpillSlot>(getAllocator(CMK_WasmSpillRefs));
+ }
+
+ JITDUMP("Spilling %zu live ref(s) for call\n", defs.size());
+ DISPNODE(tree);
+ for (GenTree* def : defs)
+ {
+ JITDUMP(" ");
+ DISPNODE(def);
+
+ int spillSlot = -1;
+ // Find an existing spill slot of the right type (byref or ref) that isn't in use
+ for (WasmSpillSlot& slot : *m_wasmSpillSlots)
+ {
+ if (slot.inUse)
+ continue;
+ if (slot.byRef != def->TypeIs(TYP_BYREF))
+ continue;
+
+ spillSlot = slot.lclNum;
+ slot.inUse = true;
+ break;
+ }
+
+ // We didn't find an available spill slot so make a new one
+ if (spillSlot == -1)
+ {
+ spillSlot = lvaGrabTemp(false DEBUGARG("WasmSpillRefs spill slot"));
+ LclVarDsc* const varDsc = lvaGetDesc(spillSlot);
+ varDsc->lvType = def->TypeGet();
+ varDsc->lvPinned = true;
+ varDsc->lvMustInit = true;
+ lvaSetVarDoNotEnregister(spillSlot DEBUGARG(DoNotEnregisterReason::WasmGCVisibility));
+ WasmSpillSlot slotDesc;
+ slotDesc.lclNum = spillSlot;
+ slotDesc.byRef = def->TypeIs(TYP_BYREF);
+ slotDesc.inUse = true;
+ m_wasmSpillSlots->push_back(slotDesc);
+ }
+
+ GenTreeLclVar* spill = gtNewStoreLclVarNode(spillSlot, def);
+ GenTreeLclVar* reload = gtNewLclVarNode(spillSlot, def->TypeGet());
+ LIR::Use use;
+ noway_assert(LIR::AsRange(block).TryGetUse(def, &use));
+ LIR::AsRange(block).InsertAfter(def, spill);
+ LIR::AsRange(block).InsertAfter(spill, reload);
+ use.ReplaceWith(reload);
+
+ // The user will expect to have child nodes that have the multiply-used flag set, so when
+ // we replace the expression with the reload node, we need to transfer the flag
+ 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;
+ }
+
+ spillSlotsToZeroAtEndOfBlock.push_back(spillSlot);
+
+ 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() || !tree->TypeIs(TYP_REF, TYP_BYREF))
+ {
+ 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))
+ {
+ GenTreeLclVarCommon* lclVar = tree->AsLclVarCommon();
+ LclVarDsc* dsc = lvaGetDesc(lclVar);
+ 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.
+ defs.push_back(tree);
+ }
+
+ // For the spill slots we used, zero them at the end of the block to avoid keeping objects pinned way longer
+ // than absolutely necessary.
+ if (spillSlotsToZeroAtEndOfBlock.size())
+ {
+ // There's no point in wasting time zeroing slots immediately before a return at the end of the block
+ // since the GC will have no opportunity to inspect the slots. We could potentially do this for throws
+ // too, but it's possible the exception would be caught and the method would continue running
+ if (!block->KindIs(BBJ_RETURN))
+ {
+ for (unsigned lclNum : spillSlotsToZeroAtEndOfBlock)
+ {
+ GenTree* zero = gtNewZeroConNode(TYP_I_IMPL);
+ GenTree* store = gtNewStoreLclVarNode(lclNum, zero);
+ LIR::Range range = LIR::SeqTree(this, store);
+ LIR::InsertBeforeTerminator(block, std::move(range));
+ }
+ }
+
+ spillSlotsToZeroAtEndOfBlock.clear();
+ }
+ }
+
+ if (m_wasmSpillSlots != nullptr)
+ {
+ JITDUMP("Total allocated spill slot count was %zu\n", m_wasmSpillSlots->size());
+ }
+ return anyChanges ? PhaseStatus::MODIFIED_EVERYTHING : PhaseStatus::MODIFIED_NOTHING;
+}
+
 #ifdef DEBUG
 
 //------------------------------------------------------------------------

src/coreclr/jit/gentree.cpp

@@ -13145,6 +13145,20 @@ void Compiler::gtGetLclVarNameInfo(unsigned lclNum, const char** ilKindOut, cons
  const char* ilName = nullptr;
 
  unsigned ilNum = compMap2ILvarNum(lclNum);
+#if defined(TARGET_WASM)
+ int wasmSpillSlotIndex = -1;
+ if (m_wasmSpillSlots != nullptr)
+ {
+ for (unsigned i = 0; i < m_wasmSpillSlots->size(); i++)
+ {
+ if (m_wasmSpillSlots->at(i).lclNum == lclNum)
+ {
+ wasmSpillSlotIndex = i;
+ break;
+ }
+ }
+ }
+#endif
 
  if (ilNum == (unsigned)ICorDebugInfo::RETBUF_ILNUM)
  {
@@ -13214,12 +13228,6 @@ void Compiler::gtGetLclVarNameInfo(unsigned lclNum, const char** ilKindOut, cons
  {
  ilName = "AsyncCont";
  }
-#if defined(TARGET_WASM)
- else if (lclNum == lvaWasmSpArg)
- {
- ilName = "SP";
- }
-#endif // defined(TARGET_WASM)
  else
  {
  ilKind = "tmp";
@@ -13261,6 +13269,26 @@ void Compiler::gtGetLclVarNameInfo(unsigned lclNum, const char** ilKindOut, cons
  }
  }
 
+#if defined(TARGET_WASM)
+ if (lclNum == lvaWasmSpArg)
+ {
+ 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)
+
  *ilKindOut = ilKind;
  *ilNameOut = ilName;
  *ilNumOut = ilNum;