[VectorCombine] Avoid crash when the next node is deleted.

[dtcxzyw]

RecursivelyDeleteTriviallyDeadInstructions is introduced by #149047 to immediately drop dead instructions. However, it may invalidate the next iterator in make_early_inc_range in some edge cases, which leads to a crash. This patch manually maintains the next iterator and updates it when the next instruction is about to be deleted.

Closes #155110.

[llvmbot]

@llvm/pr-subscribers-vectorizers

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

RecursivelyDeleteTriviallyDeadInstructions is introduced by #149047 to immediately drop dead instructions. However, it may invalidate the next iterator in make_early_inc_range in some edge cases, which leads to a crash. This patch manually maintains the next iterator and updates it when the next instruction is about to be deleted.

Closes #155110.

---

Full diff: https://github.com/llvm/llvm-project/pull/155115.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/VectorCombine.cpp (+16-5)
• (modified) llvm/test/Transforms/VectorCombine/X86/insert-binop-inseltpoison.ll (+20)

diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index 1275d53a075b5..dbc4f57537821 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -99,6 +99,10 @@ class VectorCombine {
 
   InstructionWorklist Worklist;
 
+  /// Next instruction to iterate. It will be updated when it is erased by
+  /// RecursivelyDeleteTriviallyDeadInstructions.
+  Instruction *NextInst;
+
   // TODO: Direct calls from the top-level "run" loop use a plain "Instruction"
   //       parameter. That should be updated to specific sub-classes because the
   //       run loop was changed to dispatch on opcode.
@@ -172,9 +176,11 @@ class VectorCombine {
         if (auto *OpI = dyn_cast<Instruction>(Op)) {
           if (RecursivelyDeleteTriviallyDeadInstructions(
                   OpI, nullptr, nullptr, [this](Value *V) {
-                    if (auto I = dyn_cast<Instruction>(V)) {
+                    if (auto *I = dyn_cast<Instruction>(V)) {
                       LLVM_DEBUG(dbgs() << "VC: Erased: " << *I << '\n');
                       Worklist.remove(I);
+                      if (I == NextInst)
+                        NextInst = NextInst->getNextNode();
                     }
                   }))
             continue;
@@ -4254,13 +4260,18 @@ bool VectorCombine::run() {
     if (!DT.isReachableFromEntry(&BB))
       continue;
     // Use early increment range so that we can erase instructions in loop.
-    for (Instruction &I : make_early_inc_range(BB)) {
-      if (I.isDebugOrPseudoInst())
-        continue;
-      MadeChange |= FoldInst(I);
+    Instruction *I = &BB.front();
+    while (I) {
+      NextInst = I->getNextNode();
+      if (!I->isDebugOrPseudoInst()) {
+        MadeChange |= FoldInst(*I);
+      }
+      I = NextInst;
     }
   }
 
+  NextInst = nullptr;
+
   while (!Worklist.isEmpty()) {
     Instruction *I = Worklist.removeOne();
     if (!I)
diff --git a/llvm/test/Transforms/VectorCombine/X86/insert-binop-inseltpoison.ll b/llvm/test/Transforms/VectorCombine/X86/insert-binop-inseltpoison.ll
index c1100780254c1..6228b09884c03 100644
--- a/llvm/test/Transforms/VectorCombine/X86/insert-binop-inseltpoison.ll
+++ b/llvm/test/Transforms/VectorCombine/X86/insert-binop-inseltpoison.ll
@@ -232,3 +232,23 @@ define <4 x float> @ins3_ins3_fdiv(float %x, float %y) {
   %r = fdiv <4 x float> %i0, %i1
   ret <4 x float> %r
 }
+
+; EEnsure we don't crash when erasing dead instructions.
+
+define i32 @pr155110(i32 %x) {
+; CHECK-LABEL: @pr155110(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[VECTOR_PH:%.*]]
+; CHECK:       vector.ph:
+; CHECK-NEXT:    br label [[VECTOR_PH]]
+;
+entry:
+  br label %vector.ph
+
+vector.ph:                                        ; preds = %vector.ph, %entry
+  %phi = phi i32 [ 0, %entry ], [ %reduce, %vector.ph ]
+  %inselt = insertelement <4 x i32> poison, i32 %phi, i64 0
+  %and = and <4 x i32> %inselt, zeroinitializer
+  %reduce = call i32 @llvm.vector.reduce.and.v4i32(<4 x i32> zeroinitializer)
+  br label %vector.ph
+}

[nikic]

LGTM

[davemgreen]

Thanks for the fix. I hadn't considered a phi being trivially removed around a loop edge, if that is what happened here.