[CodeGenPrepare] Fix zapping dead operands of assume

[yrouban]

Summary

This patch fixes a problem of the commit
52cc97a0 Zap the argument of llvm.assume when deleting it
A test case is created to demonstrate the crash caused by the instruction iterator invalidated by the recursive removal of dead operands of assume. The solution restarts from the first instruction in case CurInstIterator is invalidated by RecursivelyDeleteTriviallyDeadInstructions().

[bkramer]

Thanks. Can we just wrap it in resetIteratorIfInvalidatedWhileCalling like the call a few lines later? That seems less error prone.

[yrouban]

In D87434#2265247, @bkramer wrote:

Thanks. Can we just wrap it in resetIteratorIfInvalidatedWhileCalling like the call a few lines later? That seems less error prone.

Ok, but it will be less efficient.

[yrouban]

RecursivelyDeleteTriviallyDeadInstructions() call is wrapped with resetIteratorIfInvalidatedWhileCalling().

[bkramer]

Thanks, I prefer simplicity over efficiency here.

llvm/lib/CodeGen/CodeGenPrepare.cpp

if (CI->hasFnAttr(Attribute::Cold) &&
     }
 
   IntrinsicInst *II = dyn_cast<IntrinsicInst>(CI);
   if (II) {
     switch (II->getIntrinsicID()) {
     default: break;
     case Intrinsic::assume: {
       Value *Operand = II->getOperand(0);
-      II->eraseFromParent();
+      II->setOperand(0, nullptr);
       // Prune the operand, it's most likely dead.
       RecursivelyDeleteTriviallyDeadInstructions(
-          Operand, TLInfo, nullptr,
-          [&](Value *V) { removeAllAssertingVHReferences(V); });
+          Operand, TLInfo, nullptr, [&](Value *V) {
+            assert(V != II && "initial assume must persist");
+            removeAllAssertingVHReferences(V);
+          });
+      CurInstIterator = II->eraseFromParent();
       return true;
     }
 
     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()) {

llvm/test/Transforms/CodeGenPrepare/X86/recursively-delete-dead-instructions.ll

+; RUN: opt -codegenprepare -S -mtriple=x86_64-linux < %s | FileCheck %s
+
+declare void @llvm.assume(i1 noundef) nounwind willreturn
+
+; Recursively deleting dead operands of assume() may result in its next
+; instruction deleted and the iterator pointing to the next instruction
+; invalidated. This prevents the following simple loop in
+; CodeGenPrepare::optimizeBlock() unless CurInstIterator is fixed:
+;
+;   CurInstIterator = BB.begin();
+;   while (CurInstIterator != BB.end())
+;     optimizeInst(&*CurInstIterator++, ModifiedDT);
+;
+define i32 @test_assume_in_loop(i1 %cond1, i1 %cond2) {
+; CHECK-LABEL: @test_assume_in_loop(
+; CHECK-NEXT:  entry:
+entry:
+  br label %loop
+
+; CHECK: loop:
+; CHECK-NEXT:  br label %loop
+loop:
+  %cond3 = phi i1 [%cond1, %entry], [%cond4, %loop]
+  call void @llvm.assume(i1 %cond3)
+  %cond4 = icmp ult i1 %cond1, %cond2
+  br label %loop
+}