From 6b8d056a665fa20d321ea1d49a93a37a83960663 Mon Sep 17 00:00:00 2001
From: Egor Bogatov <egorbo@gmail.com>
Date: Sat, 2 Jul 2022 12:13:03 +0200
Subject: [PATCH] JIT: Run optRedundantBranches twice and try to compute more
 accurate doms (#70907)

Co-authored-by: Andy Ayers <andya@microsoft.com>
---
 src/coreclr/jit/compiler.h              |  2 +
 src/coreclr/jit/flowgraph.cpp           | 39 ++++++++++++++++++++
 src/coreclr/jit/redundantbranchopts.cpp | 49 ++++++++++++++++++++++++-
 3 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h
index f56a4a3730d6b..ca9bcdfcb2698 100644
--- a/src/coreclr/jit/compiler.h
+++ b/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);
diff --git a/src/coreclr/jit/flowgraph.cpp b/src/coreclr/jit/flowgraph.cpp
index f1e74ff10094c..1e3806b88fe8b 100644
--- a/src/coreclr/jit/flowgraph.cpp
+++ b/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
diff --git a/src/coreclr/jit/redundantbranchopts.cpp b/src/coreclr/jit/redundantbranchopts.cpp
index 75670400c9414..38911bc879540 100644
--- a/src/coreclr/jit/redundantbranchopts.cpp
+++ b/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.
                         //