JIT: Propagate LCL_ADDRs into natural loops

src/coreclr/jit/compiler.h

@@ -2462,6 +2462,40 @@ struct LoopSideEffects
     void AddModifiedElemType(Compiler* comp, CORINFO_CLASS_HANDLE structHnd);
 };
 
+typedef JitHashTable<unsigned, JitSmallPrimitiveKeyFuncs<unsigned>, bool> LocalNumSet;
+
+class LoopDefinitions
+{
+    FlowGraphNaturalLoops* m_loops;
+    LocalNumSet** m_sets;
+
+    LoopDefinitions(FlowGraphNaturalLoops* loops);
+
+    LocalNumSet* GetOrCreate(FlowGraphNaturalLoop* loop);
+    void Add(FlowGraphNaturalLoop* loop, unsigned lclNum);
+    void IncludeIntoParent(FlowGraphNaturalLoop* loop);
+public:
+    FlowGraphNaturalLoops* GetLoops()
+    {
+        return m_loops;
+    }
+
+    template<typename Visitor>
+    void VisitDefinedLocalNums(FlowGraphNaturalLoop* loop, Visitor visitor)
+    {
+        LocalNumSet* set = m_sets[loop->GetIndex()];
+        if (set != nullptr)
+        {
+            for (unsigned lclNum : LocalNumSet::KeyIteration(set))
+            {
+                visitor(lclNum);
+            }
+        }
+    }
+
+    static LoopDefinitions* Find(FlowGraphNaturalLoops* loops);
+};
+
 //  The following holds information about instr offsets in terms of generated code.
 
 enum class IPmappingDscKind
@@ -5118,6 +5152,7 @@ class Compiler
     // optimization phases. They are invalidated once RBO runs and modifies the
     // flow graph.
     FlowGraphNaturalLoops* m_loops;
+    LoopDefinitions* m_loopDefinitions;
     LoopSideEffects* m_loopSideEffects;
     BlockToNaturalLoopMap* m_blockToLoop;
     // Dominator tree used by SSA construction and copy propagation (the two are expected to use the same tree

src/coreclr/jit/fgdiagnostic.cpp

@@ -4760,6 +4760,7 @@ void Compiler::fgDebugCheckFlowGraphAnnotations()
     assert((m_loops == nullptr) || (m_loops->GetDfsTree() == m_dfsTree));
     assert((m_domTree == nullptr) || (m_domTree->GetDfsTree() == m_dfsTree));
     assert((m_reachabilitySets == nullptr) || (m_reachabilitySets->GetDfsTree() == m_dfsTree));
+    assert((m_loopDefinitions == nullptr) || (m_loopDefinitions->GetLoops() == m_loops));
 }
 
 /*****************************************************************************/

src/coreclr/jit/flowgraph.cpp

@@ -4119,6 +4119,7 @@ void Compiler::fgInvalidateDfsTree()
 {
     m_dfsTree          = nullptr;
     m_loops            = nullptr;
+    m_loopDefinitions  = nullptr;
     m_domTree          = nullptr;
     m_reachabilitySets = nullptr;
     fgSsaValid         = false;

src/coreclr/jit/lclmorph.cpp

@@ -289,17 +289,45 @@ class LocalEqualsLocalAddrAssertions
             return;
         }
 
-        m_currentAssertions = UINT64_MAX;
+        m_currentAssertions = 0;
+        bool first = true;
+        FlowGraphNaturalLoop* loop = nullptr;
         for (BasicBlock* pred : block->PredBlocks())
         {
             assert(m_comp->m_dfsTree->Contains(pred));
             if (pred->bbPostorderNum <= block->bbPostorderNum)
             {
+                loop = m_comp->m_loops->GetLoopByHeader(block);
+                if (loop != nullptr)
+                {
+                    JITDUMP("Ignoring loop backedge " FMT_BB "->" FMT_BB "\n", pred->bbNum, block->bbNum);
+                    continue;
+                }
+
+                JITDUMP("Found non-loop backedge " FMT_BB "->" FMT_BB ", clearing assertions\n", pred->bbNum, block->bbNum);
                 m_currentAssertions = 0;
                 break;
             }
 
-            m_currentAssertions &= m_outgoingAssertions[pred->bbPostorderNum];
+            if (first)
+            {
+                m_currentAssertions = m_outgoingAssertions[pred->bbPostorderNum];
+                first = false;
+            }
+            else
+            {
+                m_currentAssertions &= m_outgoingAssertions[pred->bbPostorderNum];
+            }
+        }
+
+        if ((loop != nullptr) && (m_currentAssertions != 0))
+        {
+            JITDUMP("Block " FMT_BB " is a loop header; clearing assertions about defined locals\n");
+            auto clearAssertion = [=](unsigned lclNum) {
+                Clear(lclNum);
+                };
+
+            m_comp->m_loopDefinitions->VisitDefinedLocalNums(loop, clearAssertion);
         }
 
 #ifdef DEBUG
@@ -1965,6 +1993,9 @@ PhaseStatus Compiler::fgMarkAddressExposedLocals()
         }
 #endif
 
+        m_loops = FlowGraphNaturalLoops::Find(m_dfsTree);
+        m_loopDefinitions = LoopDefinitions::Find(m_loops);
+
         LocalSequencer      sequencer(this);
         LocalAddressVisitor visitor(this, &sequencer, pAssertions);
         for (unsigned i = m_dfsTree->GetPostOrderCount(); i != 0; i--)

src/coreclr/jit/optimizer.cpp

@@ -4430,6 +4430,161 @@ void LoopSideEffects::AddModifiedElemType(Compiler* comp, CORINFO_CLASS_HANDLE s
     ArrayElemTypesModified->Set(structHnd, true, ClassHandleSet::Overwrite);
 }
 
+//------------------------------------------------------------------------
+// LoopDefinitions::LoopDefinitions: Construct a new instance for the specified
+// collection of loops.
+//
+// Arguments:
+//   loops - The loops
+//
+LoopDefinitions::LoopDefinitions(FlowGraphNaturalLoops* loops)
+    : m_loops(loops)
+{
+    if (loops->NumLoops() == 0)
+    {
+        m_sets = nullptr;
+    }
+    else
+    {
+        m_sets = new (loops->GetDfsTree()->GetCompiler(), CMK_Loops) LocalNumSet* [loops->NumLoops()] {};
+    }
+}
+
+//------------------------------------------------------------------------
+// LoopDefinitions::GetOrCreate: Get or create the set of local definitions for
+// the specified loop.
+//
+// Arguments:
+//   loop - The loop
+//
+// Returns:
+//   Set of defined locals.
+//
+LocalNumSet* LoopDefinitions::GetOrCreate(FlowGraphNaturalLoop* loop)
+{
+    if (m_sets[loop->GetIndex()] == nullptr)
+    {
+        Compiler* comp           = m_loops->GetDfsTree()->GetCompiler();
+        m_sets[loop->GetIndex()] = new (comp, CMK_Loops) LocalNumSet(comp->getAllocator(CMK_Loops));
+    }
+
+    return m_sets[loop->GetIndex()];
+}
+
+//------------------------------------------------------------------------
+// LoopDefinitions::Add: Record that a local is defined in a loop.
+//
+// Arguments:
+//   loop   - The loop
+//   lclNum - The local being defined
+//
+void LoopDefinitions::Add(FlowGraphNaturalLoop* loop, unsigned lclNum)
+{
+    GetOrCreate(loop)->Set(lclNum, true, LocalNumSet::Overwrite);
+}
+
+//------------------------------------------------------------------------
+// LoopDefinitions::IncludeIntoParent: Include all definitions recorded to
+// happen inside "loop" into its parent loop.
+//
+// Arguments:
+//   loop - The loop
+//
+void LoopDefinitions::IncludeIntoParent(FlowGraphNaturalLoop* loop)
+{
+    if (loop->GetParent() == nullptr)
+    {
+        return;
+    }
+
+    if (m_sets[loop->GetIndex()] == nullptr)
+    {
+        return;
+    }
+
+    LocalNumSet* parentSet = GetOrCreate(loop->GetParent());
+
+    for (unsigned lclNum : LocalNumSet::KeyIteration(m_sets[loop->GetIndex()]))
+    {
+        parentSet->Set(lclNum, true, LocalNumSet::Overwrite);
+    }
+}
+
+//------------------------------------------------------------------------
+// LoopDefinitions::Find: Construct the data structure that tracks all local
+// definitions inside all loops.
+//
+// Arguments:
+//   loops - The loops
+//
+// Returns:
+//   Instance of the data structure.
+//
+LoopDefinitions* LoopDefinitions::Find(FlowGraphNaturalLoops* loops)
+{
+    const FlowGraphDfsTree* dfsTree = loops->GetDfsTree();
+    Compiler*               comp    = dfsTree->GetCompiler();
+
+    LoopDefinitions* result        = new (comp, CMK_Loops) LoopDefinitions(loops);
+    BitVecTraits     poTraits      = dfsTree->PostOrderTraits();
+    BitVec           visitedBlocks = BitVecOps::MakeEmpty(&poTraits);
+
+    assert(comp->fgNodeThreading == NodeThreading::AllLocals);
+
+    for (FlowGraphNaturalLoop* loop : loops->InPostOrder())
+    {
+        loop->VisitLoopBlocks([=, &poTraits, &visitedBlocks](BasicBlock* block) {
+            if (!BitVecOps::TryAddElemD(&poTraits, visitedBlocks, block->bbPostorderNum))
+            {
+                return BasicBlockVisit::Continue;
+            }
+
+            for (Statement* stmt : block->Statements())
+            {
+                for (GenTreeLclVarCommon* lcl : stmt->LocalsTreeList())
+                {
+                    if (!lcl->OperIsLocalStore())
+                    {
+                        continue;
+                    }
+
+                    result->Add(loop, lcl->GetLclNum());
+
+                    LclVarDsc* lclDsc = comp->lvaGetDesc(lcl);
+                    if (comp->lvaIsImplicitByRefLocal(lcl->GetLclNum()) && lclDsc->lvPromoted)
+                    {
+                        // fgRetypeImplicitByRefArgs created a new promoted
+                        // struct local to represent this arg. The stores will
+                        // be rewritten by morph.
+                        assert(lclDsc->lvFieldLclStart != 0);
+                        result->Add(loop, lclDsc->lvFieldLclStart);
+                        lclDsc = comp->lvaGetDesc(lclDsc->lvFieldLclStart);
+                    }
+
+                    if (lclDsc->lvPromoted)
+                    {
+                        for (unsigned i = 0; i < lclDsc->lvFieldCnt; i++)
+                        {
+                            unsigned fieldLclNum = lclDsc->lvFieldLclStart + i;
+                            result->Add(loop, fieldLclNum);
+                        }
+                    }
+                    else if (lclDsc->lvIsStructField)
+                    {
+                        result->Add(loop, lclDsc->lvParentLcl);
+                    }
+                }
+            }
+
+            return BasicBlockVisit::Continue;
+        });
+
+        result->IncludeIntoParent(loop);
+    }
+
+    return result;
+}
+
 //------------------------------------------------------------------------
 // optHoistLoopBlocks: Hoist invariant expression out of the loop.
 //