JIT: Improve Wasm PEP calls

src/coreclr/jit/compiler.cpp

@@ -5038,10 +5038,16 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
     m_pLowering->Run();
 
     // Set stack levels and analyze throw helper usage.
+    // This may create throw helper calls that get lowered.
+    //
     StackLevelSetter stackLevelSetter(this);
     stackLevelSetter.Run();
     m_pLowering->FinalizeOutgoingArgSpace();
 
+    // Run the final liveness pass after adding throw helper calls.
+    //
+    DoPhase(this, PHASE_LCLVARLIVENESS, &Compiler::fgLateLiveness);
+
 #ifdef TARGET_WASM
     // Determine if a Virtual IP is needed and add code as needed to
     // keep the Virtual IP updated.

src/coreclr/jit/compiler.h

@@ -7418,6 +7418,7 @@ class Compiler
     void fgAsyncLiveness();
     void fgPostLowerLiveness();
     PhaseStatus fgEarlyLiveness();
+    PhaseStatus fgLateLiveness();
 
     void fgAddHandlerLiveVars(BasicBlock* block, VARSET_TP& ehHandlerLiveVars, MemoryKindSet& memoryLiveness);
 

src/coreclr/jit/liveness.cpp

@@ -1302,26 +1302,45 @@ void Liveness<TLiveness>::DoLiveVarAnalysis()
         }
     } while (changed && dfsTree->HasCycle());
 
-    // Now that we create throw helper blocks after lower,
-    // we don't need to search for them and set up liveness
-    // during lower.
-    assert(!m_compiler->fgRngChkThrowAdded);
-
-#ifdef DEBUG
-    // Double-check that no unreachable throw helper blocks exist.
-    if (m_compiler->fgBBcount != dfsTree->GetPostOrderCount())
+    // If we had unremovable blocks that are not in the DFS tree then make
+    // the 'keepAlive' set live in them. This would normally not be
+    // necessary assuming those blocks are actually unreachable; however,
+    // in LIR, throw helpers fall into this category because we do not introduce flow
+    // to them until codegen. Fix that up here.
+    //
+    if (TLiveness::IsLIR && (m_compiler->fgBBcount != dfsTree->GetPostOrderCount()))
     {
+        JITDUMP("Checking for throw helpers...\n");
+
         for (BasicBlock* block : m_compiler->Blocks())
         {
             if (dfsTree->Contains(block))
             {
                 continue;
             }
 
-            assert(!block->HasFlag(BBF_THROW_HELPER));
+            if (!block->HasFlag(BBF_THROW_HELPER))
+            {
+                continue;
+            }
+
+            JITDUMP(FMT_BB " is a throw helper, computing liveness\n", block->bbNum);
+
+            // We know throw helpers do not impact global liveness, so we just
+            // recompute within the block itself.
+            //
+            m_compiler->fgSetThrowHelpBlockLiveness(block);
+
+            // Mark last uses in the throw helper block's IR.
+            //
+            VARSET_TP keepAliveVars(VarSetOps::MakeEmpty(m_compiler));
+            VARSET_TP life(VarSetOps::MakeCopy(m_compiler, block->bbLiveOut));
+            ComputeLifeLIR(life, block, keepAliveVars);
+            assert(VarSetOps::Equal(m_compiler, life, block->bbLiveIn));
         }
     }
 
+#ifdef DEBUG
     if (m_compiler->verbose)
     {
         printf("\nBB liveness after DoLiveVarAnalysis():\n\n");
@@ -1527,6 +1546,14 @@ void Compiler::fgSetThrowHelpBlockLiveness(BasicBlock* block)
         VarSetOps::AddElemD(this, block->bbLiveOut, thisVarIndex);
     }
 
+#ifdef TARGET_WASM
+    LclVarDsc* wasmSpVarDsc = lvaGetDesc(lvaWasmSpArg);
+    if (wasmSpVarDsc->lvTracked)
+    {
+        VarSetOps::AddElemD(this, block->bbLiveOut, wasmSpVarDsc->lvVarIndex);
+    }
+#endif // TARGET_WASM
+
     if (block->HasPotentialEHSuccs(this))
     {
         block->VisitEHSuccs(this, [=](BasicBlock* succ) {

src/coreclr/jit/lower.cpp

@@ -8946,6 +8946,12 @@ void Lowering::WidenSIMD12IfNecessary(GenTreeLclVarCommon* node)
 #endif // FEATURE_SIMD
 }
 
+//------------------------------------------------------------------------
+// Lowering::DoPhase -- lower the IR
+//
+// Returns:
+//    suitable phase status
+//
 PhaseStatus Lowering::DoPhase()
 {
     // If we have any PInvoke calls, insert the one-time prolog code. We'll insert the epilog code in the
@@ -9002,54 +9008,62 @@ PhaseStatus Lowering::DoPhase()
 
     AfterLowerBlocks();
 
-#ifdef DEBUG
-    JITDUMP("Lower has completed modifying nodes.\n");
-    if (VERBOSE)
-    {
-        m_compiler->fgDispBasicBlocks(true);
-    }
-#endif
-
-    // Recompute local var ref counts before potentially sorting for liveness.
-    // Note this does minimal work in cases where we are not going to sort.
-    const bool isRecompute    = true;
-    const bool setSlotNumbers = false;
-    m_compiler->lvaComputeRefCounts(isRecompute, setSlotNumbers);
-
     if (m_compiler->m_dfsTree == nullptr)
     {
         // Compute DFS tree. We want to remove dead blocks even in MinOpts, so we
-        // do this everywhere. The dead blocks are removed below, however, some of
-        // lowering may use the DFS tree, so we compute that here.
+        // do this everywhere.
         m_compiler->m_dfsTree = m_compiler->fgComputeDfs();
     }
 
-    // Remove dead blocks. We want to remove unreachable blocks even in
-    // MinOpts.
+    // Remove dead blocks before stack level setting analyzes throw helper usage.
+    //
     m_compiler->fgRemoveBlocksOutsideDfsTree();
 
-    if (m_compiler->backendRequiresLocalVarLifetimes())
+    return PhaseStatus::MODIFIED_EVERYTHING;
+}
+
+//------------------------------------------------------------------------
+// fgLateLiveness -- rerun liveness after lower / stacklevelsetter
+//
+// Returns:
+//    suitable phase status
+//
+PhaseStatus Compiler::fgLateLiveness()
+{
+    if (!backendRequiresLocalVarLifetimes())
     {
-        assert(m_compiler->opts.OptimizationEnabled());
+        fgInvalidateDfsTree();
+        return PhaseStatus::MODIFIED_NOTHING;
+    }
 
-        m_compiler->fgPostLowerLiveness();
-        // local var liveness can delete code, which may create empty blocks
-        bool modified = m_compiler->fgUpdateFlowGraph(/* doTailDuplication */ false, /* isPhase */ false);
+    assert(backendRequiresLocalVarLifetimes());
+    assert(m_dfsTree != nullptr);
 
-        if (modified)
-        {
-            m_compiler->fgDfsBlocksAndRemove();
-            JITDUMP("had to run another liveness pass:\n");
-            m_compiler->fgPostLowerLiveness();
-        }
+    // Recompute local var ref counts before potentially sorting for liveness.
+    // Note this does minimal work in cases where we are not going to sort.
+    const bool isRecompute    = true;
+    const bool setSlotNumbers = false;
+    lvaComputeRefCounts(isRecompute, setSlotNumbers);
+
+    assert(opts.OptimizationEnabled());
 
-        // Recompute local var ref counts again after liveness to reflect
-        // impact of any dead code removal. Note this may leave us with
-        // tracked vars that have zero refs.
-        m_compiler->lvaComputeRefCounts(isRecompute, setSlotNumbers);
+    fgPostLowerLiveness();
+    // local var liveness can delete code, which may create empty blocks
+    bool modified = fgUpdateFlowGraph(/* doTailDuplication */ false, /* isPhase */ false);
+
+    if (modified)
+    {
+        fgDfsBlocksAndRemove();
+        JITDUMP("had to run another liveness pass:\n");
+        fgPostLowerLiveness();
     }
 
-    m_compiler->fgInvalidateDfsTree();
+    // Recompute local var ref counts again after liveness to reflect
+    // impact of any dead code removal. Note this may leave us with
+    // tracked vars that have zero refs.
+    lvaComputeRefCounts(isRecompute, setSlotNumbers);
+
+    fgInvalidateDfsTree();
 
     return PhaseStatus::MODIFIED_EVERYTHING;
 }