JIT: add more "simple" complexity helpers (#110486)

In early phases we haven't set tree costs or size estimates, so we fall back
to counting statements or nodes.

Capture some of this in some helper methods, both pure counting and "exceeds"
variants (which are potentially cheaper when comparing against a budget).

Update the loop cloning size check to use this.

Co-authored-by: Aman Khalid <amankhalid@microsoft.com>

src/coreclr/jit/block.cpp

@@ -1942,3 +1942,134 @@ StackEntry* BasicBlock::bbStackOnEntry() const
     assert(bbEntryState);
     return bbEntryState->esStack;
 }
+
+//------------------------------------------------------------------------
+// StatementCount: number of statements in the block.
+//
+// Returns:
+//   count of statements
+//
+// Notes:
+//   If you are calling this in order to compare the statement count
+//   against a limit, use StatementCountExceeds as it may do less work.
+//
+unsigned BasicBlock::StatementCount()
+{
+    unsigned count = 0;
+
+    for (Statement* const stmt : Statements())
+    {
+        count++;
+    }
+
+    return count;
+}
+
+//------------------------------------------------------------------------
+// StatementCountExceeds: check if the number of statements in the block
+//   exceeds some limit
+//
+// Arguments:
+//    limit  - limit on the number of statements
+//    count  - [out, optional] actual number of statements (if less than or equal to limit)
+//
+// Returns:
+//   true if the number of statements is greater than limit
+//
+bool BasicBlock::StatementCountExceeds(unsigned limit, unsigned* count /* = nullptr */)
+{
+    unsigned localCount = 0;
+    bool     overLimit  = false;
+
+    for (Statement* const stmt : Statements())
+    {
+        if (++localCount > limit)
+        {
+            overLimit = true;
+            break;
+        }
+    }
+
+    if (count != nullptr)
+    {
+        *count = localCount;
+    }
+
+    return overLimit;
+}
+
+//------------------------------------------------------------------------
+// ComplexityExceeds: check if the number of nodes in the trees in the block
+//   exceeds some limit
+//
+// Arguments:
+//    comp   - compiler instance
+//    limit  - limit on the number of nodes
+//    count  - [out, optional] actual number of nodes (if less than or equal to limit)
+//
+// Returns:
+//   true if the number of nodes is greater than limit
+//
+bool BasicBlock::ComplexityExceeds(Compiler* comp, unsigned limit, unsigned* count /* = nullptr */)
+{
+    unsigned localCount = 0;
+    bool     overLimit  = false;
+
+    for (Statement* const stmt : Statements())
+    {
+        unsigned slack  = limit - localCount;
+        unsigned actual = 0;
+        if (comp->gtComplexityExceeds(stmt->GetRootNode(), slack, &actual))
+        {
+            overLimit = true;
+            break;
+        }
+
+        localCount += actual;
+    }
+
+    if (count != nullptr)
+    {
+        *count = localCount;
+    }
+
+    return overLimit;
+}
+
+//------------------------------------------------------------------------
+// ComplexityExceeds: check if the number of nodes in the trees in the blocks
+//   in the range exceeds some limit
+//
+// Arguments:
+//    comp   - compiler instance
+//    limit  - limit on the number of nodes
+//    count  - [out, optional] actual number of nodes (if less than or equal to limit)
+//
+// Returns:
+//   true if the number of nodes is greater than limit
+//
+bool BasicBlockRangeList::ComplexityExceeds(Compiler* comp, unsigned limit, unsigned* count /* = nullptr */)
+{
+    unsigned localCount = 0;
+    bool     overLimit  = false;
+
+    for (BasicBlock* const block : *this)
+    {
+        unsigned slack  = limit - localCount;
+        unsigned actual = 0;
+        if (block->ComplexityExceeds(comp, slack, &actual))
+        {
+            overLimit = true;
+            break;
+        }
+
+        localCount += actual;
+    }
+
+    if (count != nullptr)
+    {
+        *count = localCount;
+    }
+
+    return overLimit;
+}

src/coreclr/jit/block.h

@@ -1667,19 +1667,6 @@ struct BasicBlock : private LIR::Range
 
     bool bbFallsThrough() const;
 
-    // Our slop fraction is 1/50 of the block weight.
-    static weight_t GetSlopFraction(weight_t weightBlk)
-    {
-        return weightBlk / 50.0;
-    }
-
-    // Given an the edge b1 -> b2, calculate the slop fraction by
-    // using the higher of the two block weights
-    static weight_t GetSlopFraction(BasicBlock* b1, BasicBlock* b2)
-    {
-        return GetSlopFraction(max(b1->bbWeight, b2->bbWeight));
-    }
-
 #ifdef DEBUG
     unsigned        bbTgtStkDepth; // Native stack depth on entry (for throw-blocks)
     static unsigned s_nMaxTrees;   // The max # of tree nodes in any BB
@@ -1727,6 +1714,12 @@ struct BasicBlock : private LIR::Range
         return StatementList(FirstNonPhiDef());
     }
 
+    // Simple "size" estimates
+    //
+    unsigned StatementCount();
+    bool     StatementCountExceeds(unsigned limit, unsigned* count = nullptr);
+    bool     ComplexityExceeds(Compiler* comp, unsigned limit, unsigned* complexity = nullptr);
+
     GenTree* lastNode() const;
 
     bool endsWithJmpMethod(Compiler* comp) const;
@@ -2203,6 +2196,8 @@ class BasicBlockRangeList
     {
         return BasicBlockIterator(m_end->Next()); // walk until we see the block *following* the `m_end` block
     }
+
+    bool ComplexityExceeds(Compiler* comp, unsigned limit, unsigned* count = nullptr);
 };
 
 // BBswtDesc -- descriptor for a switch block

src/coreclr/jit/compiler.h

@@ -3661,7 +3661,7 @@ class Compiler
     // is #of nodes in subtree) of "tree" is greater than "limit".
     // (This is somewhat redundant with the "GetCostEx()/GetCostSz()" fields, but can be used
     // before they have been set.)
-    bool gtComplexityExceeds(GenTree* tree, unsigned limit);
+    bool gtComplexityExceeds(GenTree* tree, unsigned limit, unsigned* complexity = nullptr);
 
     GenTree* gtReverseCond(GenTree* tree);
 

src/coreclr/jit/gentree.cpp

@@ -17852,13 +17852,14 @@ ExceptionSetFlags Compiler::gtCollectExceptions(GenTree* tree)
 // of number of sub nodes.
 //
 // Arguments:
-//     tree  - The tree to check
-//     limit - The limit in terms of number of nodes
+//     tree       - The tree to check
+//     limit      - The limit in terms of number of nodes
+//     complexity - [out, optional] the actual node count (if not greater than limit)
 //
 // Return Value:
-//     True if there are mode sub nodes in tree; otherwise false.
+//     True if there are more than limit nodes in tree; otherwise false.
 //
-bool Compiler::gtComplexityExceeds(GenTree* tree, unsigned limit)
+bool Compiler::gtComplexityExceeds(GenTree* tree, unsigned limit, unsigned* complexity)
 {
     struct ComplexityVisitor : GenTreeVisitor<ComplexityVisitor>
     {
@@ -17883,13 +17884,26 @@ bool Compiler::gtComplexityExceeds(GenTree* tree, unsigned limit)
             return WALK_CONTINUE;
         }
 
+        unsigned NumNodes()
+        {
+            return m_numNodes;
+        }
+
     private:
         unsigned m_limit;
         unsigned m_numNodes = 0;
     };
 
     ComplexityVisitor visitor(this, limit);
-    return visitor.WalkTree(&tree, nullptr) == WALK_ABORT;
+
+    fgWalkResult result = visitor.WalkTree(&tree, nullptr);
+
+    if (complexity != nullptr)
+    {
+        *complexity = visitor.NumNodes();
+    }
+
+    return (result == WALK_ABORT);
 }
 
 bool GenTree::IsPhiNode()

src/coreclr/jit/jitconfigvalues.h

@@ -62,7 +62,7 @@ CONFIG_INTEGER(JitCloneLoops, "JitCloneLoops", 1)         // If 0, don't clone.
 CONFIG_INTEGER(JitCloneLoopsWithEH, "JitCloneLoopsWithEH", 0) // If 0, don't clone loops containing EH regions
 CONFIG_INTEGER(JitCloneLoopsWithGdvTests, "JitCloneLoopsWithGdvTests", 1)     // If 0, don't clone loops based on
                                                                               // invariant type/method address tests
-RELEASE_CONFIG_INTEGER(JitCloneLoopsSizeLimit, "JitCloneLoopsSizeLimit", 400) // limit cloning to loops with less
+RELEASE_CONFIG_INTEGER(JitCloneLoopsSizeLimit, "JitCloneLoopsSizeLimit", 400) // limit cloning to loops with no more
                                                                               // than this many tree nodes
 CONFIG_INTEGER(JitDebugLogLoopCloning, "JitDebugLogLoopCloning", 0) // In debug builds log places where loop cloning
                                                                     // optimizations are performed on the fast path.

src/coreclr/jit/loopcloning.cpp

@@ -1798,68 +1798,34 @@ void Compiler::optPerformStaticOptimizations(FlowGraphNaturalLoop*     loop,
 //
 bool Compiler::optShouldCloneLoop(FlowGraphNaturalLoop* loop, LoopCloneContext* context)
 {
-    // Compute loop size
+    // See if loop size exceeds the limit.
     //
-    unsigned loopSize = 0;
+    const int      sizeConfig = JitConfig.JitCloneLoopsSizeLimit();
+    unsigned const sizeLimit  = (sizeConfig >= 0) ? (unsigned)sizeConfig : UINT_MAX;
+    unsigned       size       = 0;
 
-    // For now we use a very simplistic model where each tree node
-    // has the same code size.
-    //
-    // CostSz is not available until later.
-    //
-    struct TreeCostWalker : GenTreeVisitor<TreeCostWalker>
-    {
-        enum
-        {
-            DoPreOrder = true,
-        };
-
-        unsigned m_nodeCount;
-
-        TreeCostWalker(Compiler* comp)
-            : GenTreeVisitor(comp)
-            , m_nodeCount(0)
+    BasicBlockVisit result = loop->VisitLoopBlocks([&](BasicBlock* block) {
+        assert(sizeLimit >= size);
+        unsigned const slack     = sizeLimit - size;
+        unsigned       blockSize = 0;
+        if (block->ComplexityExceeds(this, slack, &blockSize))
         {
+            return BasicBlockVisit::Abort;
         }
 
-        fgWalkResult PreOrderVisit(GenTree** use, GenTree* user)
-        {
-            m_nodeCount++;
-            return WALK_CONTINUE;
-        }
-
-        void Reset()
-        {
-            m_nodeCount = 0;
-        }
-        unsigned Cost()
-        {
-            return m_nodeCount;
-        }
-    };
-
-    TreeCostWalker costWalker(this);
-
-    loop->VisitLoopBlocks([&](BasicBlock* block) {
-        weight_t normalizedWeight = block->getBBWeight(this);
-        for (Statement* const stmt : block->Statements())
-        {
-            costWalker.Reset();
-            costWalker.WalkTree(stmt->GetRootNodePointer(), nullptr);
-            loopSize += costWalker.Cost();
-        }
+        size += blockSize;
         return BasicBlockVisit::Continue;
     });
 
-    int const sizeLimit = JitConfig.JitCloneLoopsSizeLimit();
-
-    if ((sizeLimit >= 0) && (loopSize >= (unsigned)sizeLimit))
+    if (result == BasicBlockVisit::Abort)
     {
-        JITDUMP("Loop cloning: rejecting loop " FMT_LP " of size %u, size limit %d\n", loop->GetIndex(), loopSize,
-                sizeLimit);
+        JITDUMP("Loop cloning: rejecting loop " FMT_LP ": exceeds size limit %u\n", loop->GetIndex(), sizeLimit);
         return false;
     }
 
+    JITDUMP("Loop cloning: loop " FMT_LP ": size %u does not exceed size limit %u\n", loop->GetIndex(), size,
+            sizeLimit);
+
     return true;
 }