[JumpThreading] Fold br(freeze(undef))

[aqjune]

Summary

This patch makes JumpThreading fold br(freeze(undef)) if the freeze instruction
is only used by the branch.

[efriedma]

I'm a little surprised this is relevant, but sure, 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 (Instruction *I = dyn_cast<Instruction>(Condition)) {
     if (SimpleVal) {
       I->replaceAllUsesWith(SimpleVal);
       if (isInstructionTriviallyDead(I, TLI))
         I->eraseFromParent();
       Condition = SimpleVal;
     }
   }
 
-  // If the terminator is branching on an undef, we can pick any of the
-  // successors to branch to.  Let GetBestDestForJumpOnUndef decide.
-  if (isa<UndefValue>(Condition)) {
+  // If the terminator is branching on an undef or freeze undef, we can pick any
+  // of the successors to branch to.  Let GetBestDestForJumpOnUndef decide.
+  auto *FI = dyn_cast<FreezeInst>(Condition);
+  if (isa<UndefValue>(Condition) ||
+      (FI && isa<UndefValue>(FI->getOperand(0)) && FI->hasOneUse())) {
     unsigned BestSucc = GetBestDestForJumpOnUndef(BB);
     std::vector<DominatorTree::UpdateType> Updates;
 
     // Fold the branch/switch.
     Instruction *BBTerm = BB->getTerminator();
     Updates.reserve(BBTerm->getNumSuccessors());
     for (unsigned i = 0, e = BBTerm->getNumSuccessors(); i != e; ++i) {
       if (i == BestSucc) continue;
       BasicBlock *Succ = BBTerm->getSuccessor(i);
       Succ->removePredecessor(BB, true);
       Updates.push_back({DominatorTree::Delete, BB, Succ});
     }
 
     LLVM_DEBUG(dbgs() << "  In block '" << BB->getName()
                       << "' folding undef terminator: " << *BBTerm << '\n');
     BranchInst::Create(BBTerm->getSuccessor(BestSucc), BBTerm);
     BBTerm->eraseFromParent();
     DTU->applyUpdatesPermissive(Updates);
+    if (FI)
+      FI->eraseFromParent();
     return true;
   }
 
   // If the terminator of this block is branching on a constant, simplify the
   // terminator to an unconditional branch.  This can occur due to threading in
   // other blocks.
   if (getKnownConstant(Condition, Preference)) {
     LLVM_DEBUG(dbgs() << "  In block '" << BB->getName()

llvm/test/Transforms/JumpThreading/freeze.ll

T2:
   ret i32 %B
 
 F2:
   ret i32 %B
 }
 
 define i32 @test1_undef(i1 %cond) {
 ; CHECK-LABEL: @test1_undef(
-; CHECK-NEXT:    br i1 [[COND:%.*]], label [[MERGE_THREAD:%.*]], label [[MERGE:%.*]]
-; CHECK:       Merge.thread:
-; CHECK-NEXT:    [[V1:%.*]] = call i32 @f1()
-; CHECK-NEXT:    br label [[T2:%.*]]
-; CHECK:       Merge:
-; CHECK-NEXT:    [[V2:%.*]] = call i32 @f2()
-; CHECK-NEXT:    [[A_FR:%.*]] = freeze i1 undef
-; CHECK-NEXT:    br i1 [[A_FR]], label [[T2]], label [[F2:%.*]]
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[T2:%.*]], label [[F2:%.*]]
 ; CHECK:       T2:
-; CHECK-NEXT:    [[B4:%.*]] = phi i32 [ [[V1]], [[MERGE_THREAD]] ], [ [[V2]], [[MERGE]] ]
+; CHECK-NEXT:    [[V1:%.*]] = call i32 @f1()
 ; CHECK-NEXT:    call void @f3()
-; CHECK-NEXT:    ret i32 [[B4]]
+; CHECK-NEXT:    ret i32 [[V1]]
 ; CHECK:       F2:
+; CHECK-NEXT:    [[V2:%.*]] = call i32 @f2()
 ; CHECK-NEXT:    ret i32 [[V2]]
 ;
   br i1 %cond, label %T1, label %F1
 
 T1:
   %v1 = call i32 @f1()
   br label %Merge