JIT: Rework *ComplexityExceeds helpers to take caller-defined complexity computation (#118046)

Allow Compiler::gtComplexityExceeds et al to compute the complexity of a tree as something other than the count of its nodes by relying on a callback function. We don't diverge from this definition of complexity anywhere at the moment, though this means callers can compute additional information, like the total node count for a block or loop, without needing out parameters by simply including such computations in the callback. I was motivated to do this by a want for additional insight into the optimization potential of loop inversion candidates, like whether a loop might also benefit from cloning because it contains bounds checks. Such information can be included in the callback to the complexity check.

In the near future, we might want to diverge from the number of nodes in a code segment as our definition of complexity and employ something more sophisticated. These changes should make that easy to try out.

Author: amanasifkhalid

src/coreclr/jit/block.cpp

@@ -1735,79 +1735,3 @@ bool BasicBlock::StatementCountExceeds(unsigned limit, unsigned* count /* = null
 
     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

@@ -1779,7 +1779,9 @@ struct BasicBlock : private LIR::Range
     //
     unsigned StatementCount();
     bool     StatementCountExceeds(unsigned limit, unsigned* count = nullptr);
-    bool     ComplexityExceeds(Compiler* comp, unsigned limit, unsigned* complexity = nullptr);
+
+    template <typename TFunc>
+    bool ComplexityExceeds(Compiler* comp, unsigned limit, TFunc getTreeComplexity);
 
     GenTree* lastNode() const;
 
@@ -2086,7 +2088,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);
+    template <typename TFunc>
+    bool ComplexityExceeds(Compiler* comp, unsigned limit, TFunc getTreeComplexity);
 };
 
 // BBJumpTable -- descriptor blocks with N successors

src/coreclr/jit/compiler.h

@@ -3602,11 +3602,11 @@ class Compiler
 
     void gtUpdateNodeOperSideEffects(GenTree* tree);
 
-    // Returns "true" iff the complexity (not formally defined, but first interpretation
-    // is #of nodes in subtree) of "tree" is greater than "limit".
+    // Returns "true" iff the complexity (defined by 'getComplexity') 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, unsigned* complexity = nullptr);
+    // before they have been set, if 'getComplexity' is independent of them.)
+    template <typename TFunc>
+    bool gtComplexityExceeds(GenTree* tree, unsigned limit, TFunc getComplexity);
 
     GenTree* gtReverseCond(GenTree* tree);
 
@@ -7007,7 +7007,8 @@ class Compiler
     bool optCanonicalizeExits(FlowGraphNaturalLoop* loop);
     bool optCanonicalizeExit(FlowGraphNaturalLoop* loop, BasicBlock* exit);
 
-    bool optLoopComplexityExceeds(FlowGraphNaturalLoop* loop, unsigned limit);
+    template <typename TFunc>
+    bool optLoopComplexityExceeds(FlowGraphNaturalLoop* loop, unsigned limit, TFunc getTreeComplexity);
 
     PhaseStatus optCloneLoops();
     PhaseStatus optRangeCheckCloning();

src/coreclr/jit/compiler.hpp

@@ -5411,6 +5411,178 @@ BasicBlockVisit FlowGraphNaturalLoop::VisitRegularExitBlocks(TFunc func)
     return BasicBlockVisit::Continue;
 }
 
+//-----------------------------------------------------------
+// gtComplexityExceeds: Check if a tree exceeds a specified complexity limit.
+//
+// Type parameters:
+//   TFunc - Callback functor type
+
+// Arguments:
+//    tree              - The tree to check
+//    limit             - complexity limit
+//    getTreeComplexity - Callback functor that takes a GenTree* and returns its complexity
+//
+// Return Value:
+//     True if 'tree' exceeds the complexity limit, otherwise false.
+//
+template <typename TFunc>
+bool Compiler::gtComplexityExceeds(GenTree* tree, unsigned limit, TFunc getComplexity)
+{
+    struct ComplexityVisitor : GenTreeVisitor<ComplexityVisitor>
+    {
+        enum
+        {
+            DoPreOrder = true,
+        };
+
+        ComplexityVisitor(Compiler* comp, unsigned limit, TFunc getComplexity)
+            : GenTreeVisitor<ComplexityVisitor>(comp)
+            , m_complexity(0)
+            , m_limit(limit)
+            , m_getComplexity(getComplexity)
+        {
+        }
+
+        fgWalkResult PreOrderVisit(GenTree** use, GenTree* user)
+        {
+            m_complexity += m_getComplexity(*use);
+            return (m_complexity > m_limit) ? WALK_ABORT : WALK_CONTINUE;
+        }
+
+    private:
+        unsigned       m_complexity;
+        const unsigned m_limit;
+        TFunc          m_getComplexity;
+    };
+
+    assert(tree != nullptr);
+
+    ComplexityVisitor visitor(this, limit, getComplexity);
+
+    fgWalkResult result = visitor.WalkTree(&tree, nullptr);
+
+    return (result == WALK_ABORT);
+}
+
+//------------------------------------------------------------------------
+// ComplexityExceeds: Check if the trees in a block exceed a specified complexity limit.
+//
+// Type parameters:
+//   TFunc - Callback functor type
+//
+// Arguments:
+//    comp              - compiler instance
+//    limit             - complexity limit
+//    getTreeComplexity - Callback functor that takes a GenTree* and returns its complexity
+//
+// Returns:
+//   True if the trees in the block exceed the complexity limit, otherwise false.
+//
+template <typename TFunc>
+bool BasicBlock::ComplexityExceeds(Compiler* comp, unsigned limit, TFunc getTreeComplexity)
+{
+    assert(comp != nullptr);
+
+    unsigned localCount    = 0;
+    auto     getComplexity = [&](GenTree* tree) -> unsigned {
+        const unsigned treeComplexity = getTreeComplexity(tree);
+        localCount += treeComplexity;
+        return treeComplexity;
+    };
+
+    for (Statement* const stmt : Statements())
+    {
+        const unsigned slack = limit - localCount;
+        if (comp->gtComplexityExceeds(stmt->GetRootNode(), slack, getComplexity))
+        {
+            return true;
+        }
+    }
+
+    return false;
+}
+
+//------------------------------------------------------------------------
+// ComplexityExceeds: Check if the trees in a range of blocks exceed a specified complexity limit.
+//
+// Type parameters:
+//   TFunc - Callback functor type
+//
+// Arguments:
+//    comp              - compiler instance
+//    limit             - complexity limit
+//    getTreeComplexity - Callback functor that takes a GenTree* and returns its complexity
+//
+// Returns:
+//    True if the trees in the block range exceed the complexity limit, otherwise false.
+//
+template <typename TFunc>
+bool BasicBlockRangeList::ComplexityExceeds(Compiler* comp, unsigned limit, TFunc getTreeComplexity)
+{
+    assert(comp != nullptr);
+
+    unsigned localCount    = 0;
+    auto     getComplexity = [&](GenTree* tree) -> unsigned {
+        const unsigned treeComplexity = getTreeComplexity(tree);
+        localCount += treeComplexity;
+        return treeComplexity;
+    };
+
+    for (BasicBlock* const block : *this)
+    {
+        const unsigned slack = limit - localCount;
+        if (block->ComplexityExceeds(comp, slack, getComplexity))
+        {
+            return true;
+        }
+    }
+
+    return false;
+}
+
+//------------------------------------------------------------------------
+// optLoopComplexityExceeds: Check if the trees in a loop exceed a specified complexity limit.
+//
+// Type parameters:
+//   TFunc - Callback functor type
+//
+// Arguments:
+//    comp              - compiler instance
+//    limit             - complexity limit
+//    getTreeComplexity - Callback functor that takes a GenTree* and returns its complexity
+//
+// Returns:
+//     True if the trees in 'loop' exceed the complexity limit, otherwise false.
+//
+template <typename TFunc>
+bool Compiler::optLoopComplexityExceeds(FlowGraphNaturalLoop* loop, unsigned limit, TFunc getTreeComplexity)
+{
+    assert(loop != nullptr);
+
+    unsigned loopComplexity = 0;
+    auto     getComplexity  = [&](GenTree* tree) -> unsigned {
+        const unsigned treeComplexity = getTreeComplexity(tree);
+        loopComplexity += treeComplexity;
+        return treeComplexity;
+    };
+
+    BasicBlockVisit const result = loop->VisitLoopBlocks([&](BasicBlock* block) {
+        assert(limit >= loopComplexity);
+        const unsigned slack = limit - loopComplexity;
+        return block->ComplexityExceeds(this, slack, getComplexity) ? BasicBlockVisit::Abort
+                                                                    : BasicBlockVisit::Continue;
+    });
+
+    if (result == BasicBlockVisit::Abort)
+    {
+        JITDUMP("Loop " FMT_LP ": exceeds complexity limit %u\n", loop->GetIndex(), limit);
+        return true;
+    }
+
+    JITDUMP("Loop " FMT_LP ": complexity %u does not exceed limit %u\n", loop->GetIndex(), loopComplexity, limit);
+    return false;
+}
+
 /*****************************************************************************/
 #endif //_COMPILER_HPP_
 /*****************************************************************************/

src/coreclr/jit/forwardsub.cpp

@@ -569,8 +569,11 @@ bool Compiler::fgForwardSubStatement(Statement* stmt)
     // Consider instead using the height of the fwdSubNode.
     //
     unsigned const nodeLimit = 16;
+    auto           countNode = [](GenTree* tree) -> unsigned {
+        return 1;
+    };
 
-    if (gtComplexityExceeds(fwdSubNode, nodeLimit))
+    if (gtComplexityExceeds(fwdSubNode, nodeLimit, countNode))
     {
         JITDUMP(" tree to sub has more than %u nodes\n", nodeLimit);
         return false;
@@ -633,7 +636,7 @@ bool Compiler::fgForwardSubStatement(Statement* stmt)
     // height of the fwdSubNode.
     //
     unsigned const nextTreeLimit = 200;
-    if ((fsv.GetComplexity() > nextTreeLimit) && gtComplexityExceeds(fwdSubNode, 1))
+    if ((fsv.GetComplexity() > nextTreeLimit) && gtComplexityExceeds(fwdSubNode, 1, countNode))
     {
         JITDUMP(" next stmt tree is too large (%u)\n", fsv.GetComplexity());
         return false;

src/coreclr/jit/gentree.cpp

@@ -17867,65 +17867,6 @@ ExceptionSetFlags Compiler::gtCollectExceptions(GenTree* tree)
     return walker.GetFlags();
 }
 
-//-----------------------------------------------------------
-// gtComplexityExceeds: Check if a tree exceeds a specified complexity in terms
-// of number of sub nodes.
-//
-// Arguments:
-//     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 more than limit nodes in tree; otherwise false.
-//
-bool Compiler::gtComplexityExceeds(GenTree* tree, unsigned limit, unsigned* complexity)
-{
-    struct ComplexityVisitor : GenTreeVisitor<ComplexityVisitor>
-    {
-        enum
-        {
-            DoPreOrder = true,
-        };
-
-        ComplexityVisitor(Compiler* comp, unsigned limit)
-            : GenTreeVisitor(comp)
-            , m_limit(limit)
-        {
-        }
-
-        fgWalkResult PreOrderVisit(GenTree** use, GenTree* user)
-        {
-            if (++m_numNodes > m_limit)
-            {
-                return WALK_ABORT;
-            }
-
-            return WALK_CONTINUE;
-        }
-
-        unsigned NumNodes()
-        {
-            return m_numNodes;
-        }
-
-    private:
-        unsigned m_limit;
-        unsigned m_numNodes = 0;
-    };
-
-    ComplexityVisitor visitor(this, limit);
-
-    fgWalkResult result = visitor.WalkTree(&tree, nullptr);
-
-    if (complexity != nullptr)
-    {
-        *complexity = visitor.NumNodes();
-    }
-
-    return (result == WALK_ABORT);
-}
-
 bool GenTree::IsPhiNode()
 {
     return (OperGet() == GT_PHI_ARG) || (OperGet() == GT_PHI) || IsPhiDefn();