JIT: Propagate assertions into natural loops in global morph

src/coreclr/jit/assertionprop.cpp

@@ -672,8 +672,10 @@ void Compiler::optAssertionInit(bool isLocalProp)
             }
             else
             {
-                optMaxAssertionCount = (AssertionIndex)min(maxTrackedLocals, ((3 * lvaTrackedCount / 128) + 1) * 64);
+                optMaxAssertionCount = (AssertionIndex)(min(maxTrackedLocals, lvaTrackedCount) * 3 / 128 + 1) * 64;
             }
+
+            optMaxAssertionCount = optMaxAssertionCount * 2; // experiment
         }
         else
         {
@@ -1248,6 +1250,7 @@ AssertionIndex Compiler::optCreateAssertion(GenTree*         op1,
             assertion.op2.u1.iconVal = op2->AsIntCon()->gtIconVal;
             assertion.op2.vn         = optConservativeNormalVN(op2);
             assertion.op2.SetIconFlag(op2->GetIconHandleFlag());
+            assertion.op2.SetSimd(op2->IsCnsSimd());
 
             //
             // Ok everything has been set and the assertion looks good
@@ -1337,6 +1340,7 @@ AssertionIndex Compiler::optCreateAssertion(GenTree*         op1,
 
                         assertion.op2.u1.iconVal = iconVal;
                         assertion.op2.SetIconFlag(op2->GetIconHandleFlag(), op2->AsIntCon()->gtFieldSeq);
+                        assertion.op2.SetSimd(op2->IsCnsSimd());
                     }
                     else if (op2->gtOper == GT_CNS_LNG)
                     {
@@ -3340,6 +3344,10 @@ GenTree* Compiler::optConstantAssertionProp(AssertionDsc*        curAssertion,
             {
                 assert(varTypeIsIntegralOrI(tree));
                 newTree->BashToConst(curAssertion->op2.u1.iconVal, genActualType(tree));
+                if (curAssertion->op2.IsSimd())
+                {
+                    newTree->gtFlags |= GTF_ICON_SIMD_COUNT;
+                }
             }
             break;
 

src/coreclr/jit/compiler.h

@@ -2611,6 +2611,7 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
 */
 
 struct HWIntrinsicInfo;
+class LoopDefinitions;
 
 class Compiler
 {
@@ -5443,7 +5444,7 @@ class Compiler
     };
 
     PhaseStatus fgMorphBlocks();
-    void fgMorphBlock(BasicBlock* block, MorphUnreachableInfo* unreachableInfo = nullptr);
+    void fgMorphBlock(BasicBlock* block, MorphUnreachableInfo* unreachableInfo = nullptr, LoopDefinitions* loopDefs = nullptr);
     void fgMorphStmts(BasicBlock* block);
 
     void fgMergeBlockReturn(BasicBlock* block);
@@ -7829,6 +7830,7 @@ class Compiler
             optOp2Kind kind; // a const or copy assertion
         private:
             uint16_t m_encodedIconFlags; // encoded icon gtFlags, don't use directly
+            bool     m_isSimd;
         public:
             ValueNum vn;
             struct IntVal
@@ -7870,6 +7872,14 @@ class Compiler
                 m_encodedIconFlags = flags >> iconMaskTzc;
                 u1.fieldSeq        = fieldSeq;
             }
+            bool IsSimd() const
+            {
+                return m_isSimd;
+            }
+            void SetSimd(bool simd)
+            {
+                m_isSimd = simd;
+            }
         } op2;
 
         bool IsCheckedBoundArithBound()
@@ -12248,6 +12258,221 @@ class DomTreeVisitor
     }
 };
 
+// Data structure that keeps track of local definitions inside loops.
+class LoopDefinitions
+{
+    typedef JitHashTable<unsigned, JitSmallPrimitiveKeyFuncs<unsigned>, bool> LocalDefinitionsMap;
+
+    FlowGraphNaturalLoops* m_loops;
+    // For every loop, we track all definitions exclusive to that loop.
+    // Definitions in descendant loops are not kept in their ancestor's maps.
+    LocalDefinitionsMap** m_maps;
+    // Blocks whose IR we have visited to find local definitions in.
+    BitVec m_visitedBlocks;
+    // Whether we are optimizing for cross-block definitions.
+    // If true, we will also consider locals that are used as address in a block.
+    bool m_optForCrossBlock = false;
+
+    LocalDefinitionsMap* GetOrCreateMap(FlowGraphNaturalLoop* loop);
+
+    template <typename TFunc>
+    bool VisitLoopNestMaps(FlowGraphNaturalLoop* loop, TFunc& func);
+public:
+    LoopDefinitions(FlowGraphNaturalLoops* loops, bool optForCrossBlock);
+
+    template <typename TFunc>
+    void VisitDefinedLocalNums(FlowGraphNaturalLoop* loop, TFunc func);
+};
+
+inline LoopDefinitions::LoopDefinitions(FlowGraphNaturalLoops* loops, bool optForCrossBlock)
+    : m_loops(loops)
+    , m_optForCrossBlock(optForCrossBlock)
+{
+    Compiler* comp = loops->GetDfsTree()->GetCompiler();
+    m_maps = loops->NumLoops() == 0 ? nullptr : new (comp, CMK_LoopOpt) LocalDefinitionsMap* [loops->NumLoops()] {};
+    BitVecTraits poTraits = loops->GetDfsTree()->PostOrderTraits();
+    m_visitedBlocks       = BitVecOps::MakeEmpty(&poTraits);
+}
+
+//------------------------------------------------------------------------------
+// LoopDefinitions:GetOrCreateMap:
+//   Get or create the map of occurrences exclusive to a single loop.
+//
+// Parameters:
+//   loop - The loop
+//
+// Returns:
+//   Map of occurrences.
+//
+// Remarks:
+//   As a precondition occurrences of all descendant loops must already have
+//   been found.
+//
+inline LoopDefinitions::LocalDefinitionsMap* LoopDefinitions::GetOrCreateMap(FlowGraphNaturalLoop* loop)
+{
+    LocalDefinitionsMap* map = m_maps[loop->GetIndex()];
+    if (map != nullptr)
+    {
+        return map;
+    }
+
+    BitVecTraits poTraits = m_loops->GetDfsTree()->PostOrderTraits();
+
+#ifdef DEBUG
+    // As an invariant the map contains only the locals exclusive to each loop
+    // (i.e. occurrences inside descendant loops are not contained in ancestor
+    // loop maps). Double check that we've already computed the child maps to
+    // make sure we do not visit descendant blocks below.
+    for (FlowGraphNaturalLoop* child = loop->GetChild(); child != nullptr; child = child->GetSibling())
+    {
+        assert(BitVecOps::IsMember(&poTraits, m_visitedBlocks, child->GetHeader()->bbPostorderNum));
+    }
+#endif
+
+    Compiler* comp           = m_loops->GetDfsTree()->GetCompiler();
+    map                      = new (comp, CMK_LoopOpt) LocalDefinitionsMap(comp->getAllocator(CMK_LoopOpt));
+    m_maps[loop->GetIndex()] = map;
+
+    struct LocalsVisitor : GenTreeVisitor<LocalsVisitor>
+    {
+        enum
+        {
+            DoPreOrder    = true,
+            DoLclVarsOnly = true,
+        };
+
+        LocalsVisitor(Compiler* comp, LoopDefinitions::LocalDefinitionsMap* map, bool optForCrossBlock)
+            : GenTreeVisitor(comp)
+            , m_map(map)
+            , m_optForCrossBlock(optForCrossBlock)
+        {
+        }
+
+        fgWalkResult PreOrderVisit(GenTree** use, GenTree* user)
+        {
+            GenTreeLclVarCommon* lcl = (*use)->AsLclVarCommon();
+
+            if (m_optForCrossBlock)
+            {
+                if (!lcl->OperIsLocalStore() &&
+                    !(lcl->OperIs(GT_LCL_ADDR) && m_compiler->gtNodeHasSideEffects(user, GTF_PERSISTENT_SIDE_EFFECTS)))
+                {
+                    return Compiler::WALK_CONTINUE;
+                }
+            }
+            else
+            {
+                if (!lcl->OperIsLocalStore())
+                {
+                    return Compiler::WALK_CONTINUE;
+                }
+            }
+
+            m_map->Set(lcl->GetLclNum(), true, LocalDefinitionsMap::Overwrite);
+
+            LclVarDsc* lclDsc = m_compiler->lvaGetDesc(lcl);
+            if (m_compiler->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);
+                m_map->Set(lclDsc->lvFieldLclStart, true, LocalDefinitionsMap::Overwrite);
+                lclDsc = m_compiler->lvaGetDesc(lclDsc->lvFieldLclStart);
+            }
+
+            if (lclDsc->lvPromoted)
+            {
+                for (unsigned i = 0; i < lclDsc->lvFieldCnt; i++)
+                {
+                    unsigned fieldLclNum = lclDsc->lvFieldLclStart + i;
+                    m_map->Set(fieldLclNum, true, LocalDefinitionsMap::Overwrite);
+                }
+            }
+            else if (lclDsc->lvIsStructField)
+            {
+                m_map->Set(lclDsc->lvParentLcl, true, LocalDefinitionsMap::Overwrite);
+            }
+
+            return Compiler::WALK_CONTINUE;
+        }
+
+    private:
+        LoopDefinitions::LocalDefinitionsMap* m_map;
+        bool                                  m_optForCrossBlock;
+    };
+
+    LocalsVisitor visitor(comp, map, m_optForCrossBlock);
+
+    loop->VisitLoopBlocksReversePostOrder([=, &poTraits, &visitor](BasicBlock* block) {
+        if (!BitVecOps::TryAddElemD(&poTraits, m_visitedBlocks, block->bbPostorderNum))
+        {
+            return BasicBlockVisit::Continue;
+        }
+
+        for (Statement* stmt : block->NonPhiStatements())
+        {
+            visitor.WalkTree(stmt->GetRootNodePointer(), nullptr);
+        }
+
+        return BasicBlockVisit::Continue;
+    });
+
+    return map;
+}
+
+//------------------------------------------------------------------------------
+// LoopDefinitions:VisitLoopNestMaps:
+//   Visit all occurrence maps of the specified loop nest.
+//
+// Type parameters:
+//   TFunc - bool(LocalToOccurrenceMap*) functor that returns true to continue
+//           the visit and false to abort.
+//
+// Parameters:
+//   loop - Root loop of the nest.
+//   func - Functor instance
+//
+// Returns:
+//   True if the visit completed; false if "func" returned false for any map.
+//
+template <typename TFunc>
+bool LoopDefinitions::VisitLoopNestMaps(FlowGraphNaturalLoop* loop, TFunc& func)
+{
+    for (FlowGraphNaturalLoop* child = loop->GetChild(); child != nullptr; child = child->GetSibling())
+    {
+        if (!VisitLoopNestMaps(child, func))
+        {
+            return false;
+        }
+    }
+
+    return func(GetOrCreateMap(loop));
+}
+
+//------------------------------------------------------------------------------
+// LoopDefinitions:VisitDefinedLocalNums:
+//   Call a callback for all locals that are defined in the specified loop.
+//
+// Parameters:
+//   loop - The loop
+//   func - The callback
+//
+template <typename TFunc>
+void LoopDefinitions::VisitDefinedLocalNums(FlowGraphNaturalLoop* loop, TFunc func)
+{
+    auto visit = [=, &func](LocalDefinitionsMap* map) {
+        for (unsigned lclNum : LocalDefinitionsMap::KeyIteration(map))
+        {
+            func(lclNum);
+        }
+
+        return true;
+    };
+
+    VisitLoopNestMaps(loop, visit);
+}
+
 // EHClauses: adapter class for forward iteration of the exception handling table using range-based `for`, e.g.:
 //    for (EHblkDsc* const ehDsc : EHClauses(compiler))
 //

src/coreclr/jit/gentree.h

@@ -2236,6 +2236,11 @@ struct GenTree
         return (gtOper == GT_CNS_INT) ? (gtFlags & GTF_ICON_HDL_MASK) : GTF_EMPTY;
     }
 
+    bool IsCnsSimd() const
+    {
+        return (gtOper == GT_CNS_INT) ? ((gtFlags & GTF_ICON_SIMD_COUNT) != 0) : false;
+    }
+
     bool IsTlsIconHandle()
     {
         if (IsIconHandle())