JIT: Run optRedundantBranches twice and try to compute more accurate doms (#70907)

Co-authored-by: Andy Ayers <[email protected]>

Author: EgorBo

src/coreclr/jit/compiler.h

@@ -5178,6 +5178,8 @@ class Compiler
 
     BasicBlock* fgEndBBAfterMainFunction();
 
+    BasicBlock* fgGetDomSpeculatively(const BasicBlock* block);
+
     void fgUnlinkRange(BasicBlock* bBeg, BasicBlock* bEnd);
 
     void fgRemoveBlock(BasicBlock* block, bool unreachable);

src/coreclr/jit/flowgraph.cpp

@@ -3116,6 +3116,45 @@ BasicBlock* Compiler::fgLastBBInMainFunction()
     return fgLastBB;
 }
 
+//------------------------------------------------------------------------------
+// fgGetDomSpeculatively: Try determine a more accurate dominator than cached bbIDom
+//
+// Arguments:
+//    block - Basic block to get a dominator for
+//
+// Return Value:
+//    Basic block that dominates this block
+//
+BasicBlock* Compiler::fgGetDomSpeculatively(const BasicBlock* block)
+{
+    assert(fgDomsComputed);
+    BasicBlock* lastReachablePred = nullptr;
+
+    // Check if we have unreachable preds
+    for (const flowList* predEdge : block->PredEdges())
+    {
+        BasicBlock* predBlock = predEdge->getBlock();
+        if (predBlock == block)
+        {
+            continue;
+        }
+
+        // We check pred's count of InEdges - it's quite conservative.
+        // We, probably, could use fgReachable(fgFirstBb, pred) here to detect unreachable preds
+        if (predBlock->countOfInEdges() > 0)
+        {
+            if (lastReachablePred != nullptr)
+            {
+                // More than one of "reachable" preds - return cached result
+                return block->bbIDom;
+            }
+            lastReachablePred = predBlock;
+        }
+    }
+
+    return lastReachablePred == nullptr ? block->bbIDom : lastReachablePred;
+}
+
 /*****************************************************************************************************
  *
  *  Function to return the first basic block after the main part of the function. With funclets, it is

src/coreclr/jit/redundantbranchopts.cpp

@@ -47,7 +47,29 @@ PhaseStatus Compiler::optRedundantBranches()
             if (block->bbJumpKind == BBJ_COND)
             {
                 madeChanges |= m_compiler->optRedundantRelop(block);
+
+                BasicBlock* bbNext = block->bbNext;
+                BasicBlock* bbJump = block->bbJumpDest;
+
                 madeChanges |= m_compiler->optRedundantBranch(block);
+
+                // It's possible that either bbNext or bbJump were unlinked and it's proven
+                // to be profitable to pay special attention to their successors.
+                if (madeChanges && (bbNext->countOfInEdges() == 0))
+                {
+                    for (BasicBlock* succ : bbNext->Succs())
+                    {
+                        m_compiler->optRedundantBranch(succ);
+                    }
+                }
+
+                if (madeChanges && (bbJump->countOfInEdges() == 0))
+                {
+                    for (BasicBlock* succ : bbJump->Succs())
+                    {
+                        m_compiler->optRedundantBranch(succ);
+                    }
+                }
             }
         }
     };
@@ -293,8 +315,27 @@ bool Compiler::optRedundantBranch(BasicBlock* const block)
                 relopValue == 0 ? "false" : "true");
     }
 
-    while ((relopValue == -1) && (domBlock != nullptr))
+    bool trySpeculativeDom = false;
+    while ((relopValue == -1) && !trySpeculativeDom)
     {
+        if (domBlock == nullptr)
+        {
+            // It's possible that bbIDom is not up to date at this point due to recent BB modifications
+            // so let's try to quickly calculate new one
+            domBlock = fgGetDomSpeculatively(block);
+            if (domBlock == block->bbIDom)
+            {
+                // We already checked this one
+                break;
+            }
+            trySpeculativeDom = true;
+        }
+
+        if (domBlock == nullptr)
+        {
+            break;
+        }
+
         // Check the current dominator
         //
         if (domBlock->bbJumpKind == BBJ_COND)
@@ -351,6 +392,12 @@ bool Compiler::optRedundantBranch(BasicBlock* const block)
 
                     if (trueReaches && falseReaches && rii.canInferFromTrue && rii.canInferFromFalse)
                     {
+                        // JIT-TP: it didn't produce diffs so let's skip it
+                        if (trySpeculativeDom)
+                        {
+                            break;
+                        }
+
                         // Both dominating compare outcomes reach the current block so we can't infer the
                         // value of the relop.
                         //