[JumpThreading] Allow duplicating a basic block into preds when its branch condition is freeze(phi)

[aqjune]

Summary

This is the last JumpThreading patch for getting the performance numbers shown at
https://reviews.llvm.org/D84940#2184653 .

This patch makes ProcessBlock call ProcessBranchOnPHI when the branch condition
is freeze(phi) as well (originally it calls the function when the condition is
phi only).

Since what ProcessBranchOnPHI does is to duplicate the basic block into
predecessors if profitable, it is still valid when the condition is freeze(phi)
too.

    p = phi [a, pred1] [b, pred2]
    p.fr = freeze p
    br p.fr, ...
=>
  pred1:
    p.fr = freeze a
    br p.fr, ...
  pred2:
    p.fr2 = freeze b
    br p.fr2, ...

[efriedma]

llvm/lib/Transforms/Scalar/JumpThreading.cpp

1207

Would it make sense to use SimplifyValue here? (That sort of interacts with your other patch, I think?)

[aqjune]

Rebase

llvm/lib/Transforms/Scalar/JumpThreading.cpp

1207

I think it is tricky, since SimplifyValue may be extracted from icmp as well (line 1179)

[efriedma]

llvm/lib/Transforms/Scalar/JumpThreading.cpp

1207

By the same logic that makes it okay to look through freeze, it would be legal to look through icmps, I think? But maybe not profitable, and out of scope for this patch, I guess.

1778

This comment is a bit misleading. Really, this transform is valid on any PHI node, whether or not it feeds into the branch. It's just a heuristic that we only do this for i1 PHIs that feed directly into a branch.

[aqjune]

Update the comment

[aqjune]

llvm/lib/Transforms/Scalar/JumpThreading.cpp

1778

Rephrased the comment a bit

[efriedma]

LGTM

llvm/lib/Transforms/Scalar/JumpThreading.cpp

 //===- JumpThreading.cpp - Thread control through conditional blocks ------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
if (PN->getParent() == BB && isa<BranchInst>(BB->getTerminator()))
       updatePredecessorProfileMetadata(PN, BB);
 
   // Handle a variety of cases where we are branching on something derived from
   // a PHI node in the current block.  If we can prove that any predecessors
   // compute a predictable value based on a PHI node, thread those predecessors.
   if (ProcessThreadableEdges(CondInst, BB, Preference, Terminator))
     return true;
 
-  // If this is an otherwise-unfoldable branch on a phi node in the current
-  // block, see if we can simplify.
-  if (PHINode *PN = dyn_cast<PHINode>(CondInst))
-    if (PN->getParent() == BB && isa<BranchInst>(BB->getTerminator()))
-      return ProcessBranchOnPHI(PN);
+  // If this is an otherwise-unfoldable branch on a phi node or freeze(phi) in
+  // the current block, see if we can simplify.
+  PHINode *PN = dyn_cast<PHINode>(
+      isa<FreezeInst>(CondInst) ? cast<FreezeInst>(CondInst)->getOperand(0)
+                                : CondInst);
+
+  if (PN && PN->getParent() == BB && isa<BranchInst>(BB->getTerminator()))
 
   // If this is an otherwise-unfoldable branch on a XOR, see if we can simplify.
   if (CondInst->getOpcode() == Instruction::Xor &&
       CondInst->getParent() == BB && isa<BranchInst>(BB->getTerminator()))
     return ProcessBranchOnXOR(cast<BinaryOperator>(CondInst));
 
   // Search for a stronger dominating condition that can be used to simplify a
   // conditional branch leaving BB.
if (!MostPopularDest)
     MostPopularDest = BB->getTerminator()->
                             getSuccessor(GetBestDestForJumpOnUndef(BB));
 
   // Ok, try to thread it!
   return TryThreadEdge(BB, PredsToFactor, MostPopularDest);
 }
 
 /// ProcessBranchOnPHI - We have an otherwise unthreadable conditional branch on
-/// a PHI node in the current block.  See if there are any simplifications we
-/// can do based on inputs to the phi node.
+/// a PHI node (or freeze PHI) in the current block.  See if there are any
+/// simplifications we can do based on inputs to the phi node.
 bool JumpThreadingPass::ProcessBranchOnPHI(PHINode *PN) {
   BasicBlock *BB = PN->getParent();
 
   // TODO: We could make use of this to do it once for blocks with common PHI
   // values.
   SmallVector<BasicBlock*, 1> PredBBs;
   PredBBs.resize(1);
 
   // If any of the predecessor blocks end in an unconditional branch, we can
   // *duplicate* the conditional branch into that block in order to further
   // encourage jump threading and to eliminate cases where we have branch on a
   // phi of an icmp (branch on icmp is much better).
+  // This transformation is correct when a frozen PN is used as a branch
+  // condition as well, because it does not remove the freeze instruction.
+  //
+  //   p = phi [a, pred1] [b, pred2]
+  //   p.fr = freeze p
+  //   br p.fr, ...
+  // =>
+  // pred1:
+  //   p.fr = freeze a
+  //   br p.fr, ...
+  // pred2:
+  //   p.fr2 = freeze b
+  //   br p.fr2, ...
+  //
+  // This is beneficial because when a and b are icmps CodeGenPrepare can
+  // canonicalize them to br(freeze(icmp)).
   for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i) {
     BasicBlock *PredBB = PN->getIncomingBlock(i);
     if (BranchInst *PredBr = dyn_cast<BranchInst>(PredBB->getTerminator()))
       if (PredBr->isUnconditional()) {
         PredBBs[0] = PredBB;
         // Try to duplicate BB into PredBB.
         if (DuplicateCondBranchOnPHIIntoPred(BB, PredBBs))
           return true;

llvm/test/Transforms/JumpThreading/phi-copy-to-pred.ll

 EXIT1:
   ret i32 0
 EXIT2:
   ret i32 1
 }
 
 define i32 @test2(i1 %cond, i1 %a, i1 %b) {
 ; CHECK-LABEL: @test2(
-; CHECK-NEXT:    br i1 [[COND:%.*]], label [[A:%.*]], label [[B:%.*]]
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[A:%.*]], label [[C:%.*]]
 ; CHECK:       A:
 ; CHECK-NEXT:    call void @f()
-; CHECK-NEXT:    br label [[C:%.*]]
-; CHECK:       B:
-; CHECK-NEXT:    call void @g()
-; CHECK-NEXT:    br label [[C]]
+; CHECK-NEXT:    [[P_FR1:%.*]] = freeze i1 [[A:%.*]]
+; CHECK-NEXT:    br i1 [[P_FR1]], label [[EXIT1:%.*]], label [[EXIT2:%.*]]
 ; CHECK:       C:
-; CHECK-NEXT:    [[P:%.*]] = phi i1 [ [[A:%.*]], [[A]] ], [ [[B:%.*]], [[B]] ]
-; CHECK-NEXT:    [[P_FR:%.*]] = freeze i1 [[P]]
-; CHECK-NEXT:    br i1 [[P_FR]], label [[EXIT1:%.*]], label [[EXIT2:%.*]]
+; CHECK-NEXT:    call void @g()
+; CHECK-NEXT:    [[P_FR:%.*]] = freeze i1 [[B:%.*]]
+; CHECK-NEXT:    br i1 [[P_FR]], label [[EXIT1]], label [[EXIT2]]
 ; CHECK:       EXIT1:
 ; CHECK-NEXT:    ret i32 0
 ; CHECK:       EXIT2:
 ; CHECK-NEXT:    ret i32 1
 ;
   br i1 %cond, label %A, label %B
 A:
   call void @f()