From bf21b1667886cd491575596399f15a512702e473 Mon Sep 17 00:00:00 2001
From: Jakob Botsch Nielsen <jakob.botsch.nielsen@gmail.com>
Date: Thu, 24 Oct 2024 13:49:05 +0200
Subject: [PATCH] JIT: Propagate LCL_ADDRs into handlers

Expand the support for LCL_ADDR propagation to keep assertions around in
handler blocks when they are true at all points in all predecessor
blocks of those handlers. This requires keeping track of the "always
true" assertions for each block.
---
 src/coreclr/jit/lclmorph.cpp | 50 ++++++++++++++++++++++++++++++------
 1 file changed, 42 insertions(+), 8 deletions(-)

diff --git a/src/coreclr/jit/lclmorph.cpp b/src/coreclr/jit/lclmorph.cpp
index 90750945082aea..02a0ded1035cb2 100644
--- a/src/coreclr/jit/lclmorph.cpp
+++ b/src/coreclr/jit/lclmorph.cpp
@@ -224,11 +224,15 @@ class LocalEqualsLocalAddrAssertions
     ArrayStack<LocalEqualsLocalAddrAssertion> m_assertions;
     AssertionToIndexMap                       m_map;
     uint64_t*                                 m_lclAssertions;
-    uint64_t*                                 m_outgoingAssertions;
-    BitVec                                    m_localsToExpose;
+    // Assertions true going out of each block
+    uint64_t* m_outgoingAssertions;
+    // Assertions true at all points within each block
+    uint64_t* m_alwaysTrueAssertions;
+    BitVec    m_localsToExpose;
 
 public:
     uint64_t CurrentAssertions = 0;
+    uint64_t AlwaysAssertions  = 0;
 
     LocalEqualsLocalAddrAssertions(Compiler* comp)
         : m_comp(comp)
@@ -237,7 +241,8 @@ class LocalEqualsLocalAddrAssertions
     {
         m_lclAssertions =
             comp->lvaCount == 0 ? nullptr : new (comp, CMK_LocalAddressVisitor) uint64_t[comp->lvaCount]{};
-        m_outgoingAssertions = new (comp, CMK_LocalAddressVisitor) uint64_t[comp->m_dfsTree->GetPostOrderCount()]{};
+        m_outgoingAssertions   = new (comp, CMK_LocalAddressVisitor) uint64_t[comp->m_dfsTree->GetPostOrderCount()]{};
+        m_alwaysTrueAssertions = new (comp, CMK_LocalAddressVisitor) uint64_t[comp->m_dfsTree->GetPostOrderCount()]{};
 
         BitVecTraits localsTraits(comp->lvaCount, comp);
         m_localsToExpose = BitVecOps::MakeEmpty(&localsTraits);
@@ -284,25 +289,46 @@ class LocalEqualsLocalAddrAssertions
     //
     void StartBlock(BasicBlock* block)
     {
-        if ((m_assertions.Height() == 0) || (block->bbPreds == nullptr) || m_comp->bbIsHandlerBeg(block))
+        FlowEdge* preds;
+        if ((m_assertions.Height() == 0) || ((preds = m_comp->BlockPredsWithEH(block)) == nullptr))
         {
             CurrentAssertions = 0;
+            AlwaysAssertions  = 0;
             return;
         }
 
         CurrentAssertions = UINT64_MAX;
-        for (BasicBlock* pred : block->PredBlocks())
+
+        uint64_t* assertionMap = m_comp->bbIsHandlerBeg(block) ? m_alwaysTrueAssertions : m_outgoingAssertions;
+
+        INDEBUG(bool anyReachablePred = false);
+
+        for (FlowEdge* predEdge = preds; predEdge != nullptr; predEdge = predEdge->getNextPredEdge())
         {
-            assert(m_comp->m_dfsTree->Contains(pred));
+            BasicBlock* pred = predEdge->getSourceBlock();
+            if (!m_comp->m_dfsTree->Contains(pred))
+            {
+                // Edges induced due to implicit EH flow can come from
+                // unreachable blocks; skip those.
+                continue;
+            }
+
+            INDEBUG(anyReachablePred = true);
+
             if (pred->bbPostorderNum <= block->bbPostorderNum)
             {
                 CurrentAssertions = 0;
                 break;
             }
 
-            CurrentAssertions &= m_outgoingAssertions[pred->bbPostorderNum];
+            CurrentAssertions &= assertionMap[pred->bbPostorderNum];
         }
 
+        // There should always be at least one reachable pred for all blocks.
+        assert(anyReachablePred);
+
+        AlwaysAssertions = CurrentAssertions;
+
 #ifdef DEBUG
         if (CurrentAssertions != 0)
         {
@@ -329,7 +355,8 @@ class LocalEqualsLocalAddrAssertions
     //
     void EndBlock(BasicBlock* block)
     {
-        m_outgoingAssertions[block->bbPostorderNum] = CurrentAssertions;
+        m_outgoingAssertions[block->bbPostorderNum]   = CurrentAssertions;
+        m_alwaysTrueAssertions[block->bbPostorderNum] = AlwaysAssertions;
     }
 
     //-------------------------------------------------------------------
@@ -403,6 +430,7 @@ class LocalEqualsLocalAddrAssertions
     void Clear(unsigned dstLclNum)
     {
         CurrentAssertions &= ~m_lclAssertions[dstLclNum];
+        AlwaysAssertions &= CurrentAssertions;
     }
 
     //-----------------------------------------------------------------------------------
@@ -1148,6 +1176,9 @@ class LocalAddressVisitor final : public GenTreeVisitor<LocalAddressVisitor>
 
         if (m_lclAddrAssertions != nullptr)
         {
+            // Save and restore CurrentAssertions for each branch.
+            // AlwaysAssertions only shrinks, so it does not require special
+            // handling around QMARKs.
             uint64_t origAssertions = m_lclAddrAssertions->CurrentAssertions;
 
             if (WalkTree(&qmark->gtOp2->AsOp()->gtOp1, qmark->gtOp2) == Compiler::WALK_ABORT)
@@ -2128,6 +2159,9 @@ PhaseStatus Compiler::fgMarkAddressExposedLocals()
     }
     else
     {
+        // We'll be using BlockPredsWithEH, so clear its cache.
+        m_blockToEHPreds = nullptr;
+
         LocalEqualsLocalAddrAssertions  assertions(this);
         LocalEqualsLocalAddrAssertions* pAssertions = &assertions;