JIT: Enable cross-block local assertion prop for natural loops

src/coreclr/jit/assertionprop.cpp

@@ -603,7 +603,6 @@ ASSERT_TP& Compiler::GetAssertionDep(unsigned lclNum)
 void Compiler::optAssertionTraitsInit(AssertionIndex assertionCount)
 {
     apTraits = new (this, CMK_AssertionProp) BitVecTraits(assertionCount, this);
-    apFull   = BitVecOps::MakeFull(apTraits);
 }
 
 /*****************************************************************************
@@ -1610,6 +1609,28 @@ AssertionIndex Compiler::optFinalizeCreatingAssertion(AssertionDsc* assertion)
     return optAddAssertion(assertion);
 }
 
+//------------------------------------------------------------------------
+// optRemoveCopyAssertions: Clear all indices that represent copy assertions in
+// the specified bit set of assertions.
+//
+// Arguments:
+//    assertions - Bitset of assertions
+//
+void Compiler::optRemoveCopyAssertions(ASSERT_TP& assertions)
+{
+    BitVecOps::Iter iter(apTraits, assertions);
+    unsigned        bvIndex = 0;
+    while (iter.NextElem(&bvIndex))
+    {
+        AssertionIndex const index     = GetAssertionIndex(bvIndex);
+        AssertionDsc* const  assertion = optGetAssertion(index);
+        if (assertion->IsCopyAssertion())
+        {
+            BitVecOps::RemoveElemD(apTraits, assertions, bvIndex);
+        }
+    }
+}
+
 /*****************************************************************************
  *
  * If tree is a constant node holding an integral value, retrieve the value in
@@ -6253,7 +6274,7 @@ ASSERT_TP* Compiler::optInitAssertionDataflowFlags()
     // The local assertion gen phase may have created unreachable blocks.
     // They will never be visited in the dataflow propagation phase, so they need to
     // be initialized correctly. This means that instead of setting their sets to
-    // apFull (i.e. all possible bits set), we need to set the bits only for valid
+    // all possible bits set, we need to set the bits only for valid
     // assertions (note that at this point we are not creating any new assertions).
     // Also note that assertion indices start from 1.
     ASSERT_TP apValidFull = BitVecOps::MakeEmpty(apTraits);

src/coreclr/jit/compiler.cpp

@@ -4856,9 +4856,6 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
     // Expose candidates for implicit byref last-use copy elision.
     DoPhase(this, PHASE_IMPBYREF_COPY_OMISSION, &Compiler::fgMarkImplicitByRefCopyOmissionCandidates);
 
-    // Locals tree list is no longer kept valid.
-    fgNodeThreading = NodeThreading::None;
-
     // Apply the type update to implicit byref parameters; also choose (based on address-exposed
     // analysis) which implicit byref promotions to keep (requires copy to initialize) or discard.
     //
@@ -4875,8 +4872,15 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
         // Build post-order that morph will use, and remove dead blocks
         //
         DoPhase(this, PHASE_DFS_BLOCKS, &Compiler::fgDfsBlocksAndRemove);
+
+        // Find loops and annotations that morph will use for assertion prop
+        //
+        DoPhase(this, PHASE_FIND_LOOPS_BEFORE_MORPH, &Compiler::optFindLoopsBeforeMorph);
     }
 
+    // Locals tree list is no longer kept valid.
+    fgNodeThreading = NodeThreading::None;
+
     // Morph the trees in all the blocks of the method
     //
     unsigned const preMorphBBCount = fgBBcount;

src/coreclr/jit/compiler.h

@@ -2460,6 +2460,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
@@ -5042,6 +5076,7 @@ class Compiler
     // flow graph.
     FlowGraphNaturalLoops* m_loops;
     LoopSideEffects* m_loopSideEffects;
+    LoopDefinitions* m_loopDefinitions;
     BlockToNaturalLoopMap* m_blockToLoop;
     // Dominator tree used by SSA construction and copy propagation (the two are expected to use the same tree
     // in order to avoid the need for SSA reconstruction and an "out of SSA" phase).
@@ -6854,6 +6889,8 @@ class Compiler
     bool optSwitchConvert(BasicBlock* firstBlock, int testsCount, ssize_t* testValues, weight_t falseLikelihood, GenTree* nodeToTest);
     bool optSwitchDetectAndConvert(BasicBlock* firstBlock);
 
+    PhaseStatus optFindLoopsBeforeMorph();
+
     PhaseStatus optInvertLoops();    // Invert loops so they're entered at top and tested at bottom.
     PhaseStatus optOptimizeFlow();   // Simplify flow graph and do tail duplication
     PhaseStatus optOptimizeLayout(); // Optimize the BasicBlock layout of the method
@@ -7499,7 +7536,6 @@ class Compiler
 public:
     // Data structures for assertion prop
     BitVecTraits* apTraits;
-    ASSERT_TP     apFull;
     ASSERT_TP     apLocal;
     ASSERT_TP     apLocalIfTrue;
 
@@ -7853,6 +7889,8 @@ class Compiler
 
     AssertionIndex optFinalizeCreatingAssertion(AssertionDsc* assertion);
 
+    void optRemoveCopyAssertions(ASSERT_TP& assertions);
+
     bool optTryExtractSubrangeAssertion(GenTree* source, IntegralRange* pRange);
 
     void optCreateComplementaryAssertion(AssertionIndex assertionIndex,

src/coreclr/jit/compphases.h

@@ -31,6 +31,7 @@ CompPhaseNameMacro(PHASE_POST_IMPORT,                "Post-import",
 CompPhaseNameMacro(PHASE_IBCPREP,                    "Profile instrumentation prep",   false, -1, false)
 CompPhaseNameMacro(PHASE_IBCINSTR,                   "Profile instrumentation",        false, -1, false)
 CompPhaseNameMacro(PHASE_INCPROFILE,                 "Profile incorporation",          false, -1, false)
+CompPhaseNameMacro(PHASE_FIND_LOOPS_BEFORE_MORPH,    "Find loops before morph",        false, -1, false)
 CompPhaseNameMacro(PHASE_MORPH_INIT,                 "Morph - Init",                   false, -1, false)
 CompPhaseNameMacro(PHASE_MORPH_INLINE,               "Morph - Inlining",               false, -1, true)
 CompPhaseNameMacro(PHASE_MORPH_ADD_INTERNAL,         "Morph - Add internal blocks",    false, -1, true)

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

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

src/coreclr/jit/morph.cpp

@@ -13752,7 +13752,9 @@ void Compiler::fgMorphBlock(BasicBlock* block)
             }
             else
             {
-                bool hasPredAssertions = false;
+                bool                  hasPredAssertions = false;
+                FlowGraphNaturalLoop* loop              = nullptr;
+                bool                  searchedLoop      = false;
 
                 for (BasicBlock* const pred : block->PredBlocks())
                 {
@@ -13764,10 +13766,33 @@ void Compiler::fgMorphBlock(BasicBlock* block)
                     //
                     if (pred->bbPostorderNum <= block->bbPostorderNum)
                     {
-                        JITDUMP(FMT_BB " pred " FMT_BB " not processed; clearing assertions in\n", block->bbNum,
+                        if (!searchedLoop)
+                        {
+                            loop         = m_loops->GetLoopByHeader(block);
+                            searchedLoop = true;
+                        }
+
+                        if (loop == nullptr)
+                        {
+                            JITDUMP(FMT_BB " pred " FMT_BB " not processed; clearing assertions in\n", block->bbNum,
+                                    pred->bbNum);
+                            hasPredAssertions = false;
+                            break;
+                        }
+
+#ifdef DEBUG
+                        bool foundBackedge = false;
+                        for (FlowEdge* backedge : loop->BackEdges())
+                        {
+                            foundBackedge |= backedge->getSourceBlock() == pred;
+                        }
+
+                        assert(foundBackedge);
+#endif
+
+                        JITDUMP(FMT_BB " pred " FMT_BB " not processed, but is a backedge\n", block->bbNum,
                                 pred->bbNum);
-                        hasPredAssertions = false;
-                        break;
+                        continue;
                     }
 
                     // Yes, pred assertions are available.
@@ -13822,6 +13847,18 @@ void Compiler::fgMorphBlock(BasicBlock* block)
                     //
                     canUsePredAssertions = false;
                 }
+                else if (loop != nullptr)
+                {
+                    // Kill all assertions about locals that will be stored inside the loop.
+                    m_loopDefinitions->VisitDefinedLocalNums(loop, [=](unsigned lclNum) {
+                        BitVecOps::DiffD(apTraits, apLocal, GetAssertionDep(lclNum));
+                    });
+
+                    // Kill all copy prop assertions. Copy propagation into
+                    // loops creates long-lived lifetimes that has very mixed
+                    // benefits.
+                    optRemoveCopyAssertions(apLocal);
+                }
             }
 
             if (!canUsePredAssertions)
@@ -15196,7 +15233,7 @@ PhaseStatus Compiler::fgRetypeImplicitByRefArgs()
                     GenTree* data  = (varDsc->TypeGet() == TYP_STRUCT) ? gtNewBlkIndir(varDsc->GetLayout(), addr)
                                                                        : gtNewIndir(varDsc->TypeGet(), addr);
                     GenTree* store = gtNewStoreLclVarNode(newLclNum, data);
-                    fgNewStmtAtBeg(fgFirstBB, store);
+                    fgInsertStmtAtBeg(fgFirstBB, fgNewStmtFromTree(store));
                 }
 
                 // Update the locals corresponding to the promoted fields.