[CodeGenPrepare] Eliminate llvm.expect before removing empty blocks

[thejh]

Summary

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.

[thejh]

bug report for context: https://bugs.llvm.org/show_bug.cgi?id=49403

[bkramer]

This makes sense to me.

llvm/lib/CodeGen/CodeGenPrepare.cpp

627–628

We now have a make_early_inc_range thing to make this pattern less ugly. It can be used in a for range loop.

[thejh]

llvm/lib/CodeGen/CodeGenPrepare.cpp

627–628

I don't think that fits well with resetIteratorIfInvalidatedWhileCalling? The code that I'm moving isn't just deleting the current instruction, but also other instructions that it depends on; my understanding is that the old code is explicitly resetting the iterator in case RecursivelyDeleteTriviallyDeadInstructions ends up throwing away the following instruction because it's recursing through a phi node, or something like that; and I'm preserving that behavior.

By the way, once you think this change is ready to land, please commit it for me - I don't have commit access.

[fhahn]

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

12

can you also test cases with blocks that contain multiple assumes, as well as tests where assume calls and other instructions are mixed. Also cases where the assumption is a value defined in the function.

[thejh]

added another test

Added another test that should cover fhahn's request.

[thejh]

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

12

Done - I think the test I've added covers all your requests.

llvm/lib/CodeGen/CodeGenPrepare.cpp

private:
     // Get the DominatorTree, building if necessary.
     DominatorTree &getDT(Function &F) {
       if (!DT)
         DT = std::make_unique<DominatorTree>(F);
       return *DT;
     }
 
     void removeAllAssertingVHReferences(Value *V);
+    bool eliminateAssumptions(Function &F);
     bool eliminateFallThrough(Function &F);
     bool eliminateMostlyEmptyBlocks(Function &F);
     BasicBlock *findDestBlockOfMergeableEmptyBlock(BasicBlock *BB);
     bool canMergeBlocks(const BasicBlock *BB, const BasicBlock *DestBB) const;
     void eliminateMostlyEmptyBlock(BasicBlock *BB);
     bool isMergingEmptyBlockProfitable(BasicBlock *BB, BasicBlock *DestBB,
                                        bool isPreheader);
     bool makeBitReverse(Instruction &I);
while (BB != nullptr) {
       BasicBlock* Next = BB->getNextNode();
       // F.hasOptSize is already checked in the outer if statement.
       if (!llvm::shouldOptimizeForSize(BB, PSI, BFI.get()))
         EverMadeChange |= bypassSlowDivision(BB, BypassWidths);
       BB = Next;
     }
   }
 
+  // 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);
 
   bool ModifiedDT = false;
   if (!DisableBranchOpts)
     EverMadeChange |= splitBranchCondition(F, ModifiedDT);
 

 #ifndef NDEBUG
   if (VerifyBFIUpdates)
     verifyBFIUpdates(F);
 #endif
 
   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) {
   LargeOffsetGEPMap.erase(V);
   NewGEPBases.erase(V);
 
   auto GEP = dyn_cast<GetElementPtrInst>(V);
   if (!GEP)
for (auto &Arg : CI->arg_operands()) {
       unsigned AS = Arg->getType()->getPointerAddressSpace();
       return optimizeMemoryInst(CI, Arg, Arg->getType(), AS);
     }
 
   IntrinsicInst *II = dyn_cast<IntrinsicInst>(CI);
   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
       // selection.
       if (II->use_empty()) {
         II->eraseFromParent();
         return true;
       }

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

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -codegenprepare -S -mtriple=x86_64-linux < %s | FileCheck %s
 
 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]]
 ;
 entry:
   %l = load i8, i8* %d
   %cmp = icmp eq i8 %l, 0
   br i1 %cmp, label %exit, label %if.end

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

+; 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;
+