[CodeGenPrepare] Eliminate llvm.expect before removing empty blocks

CodeGenPrepare currently first removes empty blocks, then in a loop
performs other optimizations. One of those optimizations is the removal
of call instructions that invoke @llvm.assume, which can create new
empty blocks.

This means that when a branch only contains a call to __builtin_assume(),
the empty branch will survive into MIR, and will then only be
half-removed by MIR-level optimizations (e.g. removing the branch but
leaving the condition intact).

Fix it by eliminating @llvm.expect builtin calls before removing empty
blocks.

Reviewed By: bkramer

Differential Revision: https://reviews.llvm.org/D97848

llvm/lib/CodeGen/CodeGenPrepare.cpp

@@ -377,6 +377,7 @@ class TypePromotionTransaction;
     }
 
     void removeAllAssertingVHReferences(Value *V);
+    bool eliminateAssumptions(Function &F);
     bool eliminateFallThrough(Function &F);
     bool eliminateMostlyEmptyBlocks(Function &F);
     BasicBlock *findDestBlockOfMergeableEmptyBlock(BasicBlock *BB);
@@ -506,6 +507,11 @@ bool CodeGenPrepare::runOnFunction(Function &F) {
     }
   }
 
+  // Get rid of @llvm.assume builtins before attempting to eliminate empty
+  // blocks, since there might be blocks that only contain @llvm.assume calls
+  // (plus arguments that we can get rid of).
+  EverMadeChange |= eliminateAssumptions(F);
+
   // Eliminate blocks that contain only PHI nodes and an
   // unconditional branch.
   EverMadeChange |= eliminateMostlyEmptyBlocks(F);
@@ -614,6 +620,28 @@ bool CodeGenPrepare::runOnFunction(Function &F) {
   return EverMadeChange;
 }
 
+bool CodeGenPrepare::eliminateAssumptions(Function &F) {
+  bool MadeChange = false;
+  for (BasicBlock &BB : F) {
+    CurInstIterator = BB.begin();
+    while (CurInstIterator != BB.end()) {
+      Instruction *I = &*(CurInstIterator++);
+      if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
+        if (II->getIntrinsicID() != Intrinsic::assume)
+          continue;
+        MadeChange = true;
+        Value *Operand = II->getOperand(0);
+        II->eraseFromParent();
+
+        resetIteratorIfInvalidatedWhileCalling(&BB, [&]() {
+          RecursivelyDeleteTriviallyDeadInstructions(Operand, TLInfo, nullptr);
+        });
+      }
+    }
+  }
+  return MadeChange;
+}
+
 /// An instruction is about to be deleted, so remove all references to it in our
 /// GEP-tracking data strcutures.
 void CodeGenPrepare::removeAllAssertingVHReferences(Value *V) {
@@ -2118,18 +2146,8 @@ bool CodeGenPrepare::optimizeCallInst(CallInst *CI, bool &ModifiedDT) {
   if (II) {
     switch (II->getIntrinsicID()) {
     default: break;
-    case Intrinsic::assume: {
-      Value *Operand = II->getOperand(0);
-      II->eraseFromParent();
-      // Prune the operand, it's most likely dead.
-      resetIteratorIfInvalidatedWhileCalling(BB, [&]() {
-        RecursivelyDeleteTriviallyDeadInstructions(
-            Operand, TLInfo, nullptr,
-            [&](Value *V) { removeAllAssertingVHReferences(V); });
-      });
-      return true;
-    }
-
+    case Intrinsic::assume:
+      llvm_unreachable("llvm.assume should have been removed already");
     case Intrinsic::experimental_widenable_condition: {
       // Give up on future widening oppurtunties so that we can fold away dead
       // paths and merge blocks before going into block-local instruction

llvm/test/Transforms/CodeGenPrepare/X86/delete-assume-dead-code.ll

@@ -5,11 +5,6 @@ define i32 @test1(i8* %d) nounwind {
 ; CHECK-LABEL: @test1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[L:%.*]] = load i8, i8* [[D:%.*]], align 1
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[L]], 0
-; CHECK-NEXT:    br i1 [[CMP]], label [[EXIT:%.*]], label [[IF_END:%.*]]
-; CHECK:       if.end:
-; CHECK-NEXT:    br label [[EXIT]]
-; CHECK:       exit:
 ; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i8 [[L]], 0
 ; CHECK-NEXT:    [[CONV:%.*]] = zext i1 [[TMP0]] to i32
 ; CHECK-NEXT:    ret i32 [[CONV]]

llvm/test/Transforms/CodeGenPrepare/remove-assume-block.ll

@@ -0,0 +1,48 @@
+; RUN: opt -S -codegenprepare < %s | FileCheck %s
+;
+; Ensure that blocks that only contain @llvm.assume are removed completely
+; during CodeGenPrepare.
+
+target triple = "x86_64-unknown-linux-gnu"
+
+; CHECK-LABEL: @simple(
+; CHECK-NEXT: end:
+; CHECK-NEXT: ret void
+define void @simple(i64 %addr, i1 %assumption) {
+  %cmp1 = icmp eq i64 %addr, 0
+  br i1 %cmp1, label %do_assume, label %end
+
+do_assume:
+  tail call void @llvm.assume(i1 %assumption)
+  br label %end
+
+end:
+  ret void
+}
+
+; CHECK-LABEL: @complex_assume(
+; CHECK-NEXT: end:
+; CHECK-NEXT: ret void
+define void @complex_assume(i64 %addr, i1 %assumption_a, i1 %assumption_b,
+                          i64 %val_a, i64 %val_b) {
+  %cmp1 = icmp eq i64 %addr, 0
+  br i1 %cmp1, label %do_assume, label %end
+
+do_assume:
+  call void @llvm.assume(i1 %assumption_a)
+  call void @llvm.assume(i1 %assumption_b)
+  %val_xor = xor i64 %val_a, %val_b
+  %val_shifted = lshr i64 %val_xor, 7
+  %assumption_c = trunc i64 %val_shifted to i1
+  call void @llvm.assume(i1 %assumption_c)
+  %assumption_d = call i1 @readonly_func(i64 %val_b)
+  call void @llvm.assume(i1 %assumption_d)
+  br label %end
+
+end:
+  ret void
+}
+
+declare void @llvm.assume(i1 noundef)
+declare i1 @readonly_func(i64) nounwind readonly willreturn;
+