[VectorCombine] Avoid double deletion in eraseInstruction

[dtcxzyw]

Consider the following pattern:

C = op A B
D = op C
E = op D, C

As E is dead, we call eraseInstruction(E) and see if its operands become dead. RecursivelyDeleteTriviallyDeadInstructions(D) also erases C, which causes a UAF crash in the subsequent call RecursivelyDeleteTriviallyDeadInstructions(C).

This patch also adds deleted ops into the visit list to avoid double deletion. As an alternative, we can use WeakVH for the operand list.

Closes #155543.

[llvmbot]

@llvm/pr-subscribers-vectorizers

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

Consider the following pattern:

C = op A B
D = op C
E = op D, C

As E is dead, we call eraseInstruction(E) and see if its operands become dead. RecursivelyDeleteTriviallyDeadInstructions(D) also erases C, which causes a UAF crash in the subsequent call RecursivelyDeleteTriviallyDeadInstructions(C).

This patch also adds deleted ops into the visit list to avoid double deletion. As an alternative, we can use WeakVH for the operand list.

Closes #155543.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/VectorCombine.cpp (+3-2)
• (added) llvm/test/Transforms/VectorCombine/X86/pr155543.ll (+15)

diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index 56a08b8438718..c88ed95de2946 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -173,15 +173,16 @@ class VectorCombine {
     // further folds that were hindered by OneUse limits.
     SmallPtrSet<Value *, 4> Visited;
     for (Value *Op : Ops) {
-      if (Visited.insert(Op).second) {
+      if (!Visited.contains(Op)) {
         if (auto *OpI = dyn_cast<Instruction>(Op)) {
           if (RecursivelyDeleteTriviallyDeadInstructions(
-                  OpI, nullptr, nullptr, [this](Value *V) {
+                  OpI, nullptr, nullptr, [&](Value *V) {
                     if (auto *I = dyn_cast<Instruction>(V)) {
                       LLVM_DEBUG(dbgs() << "VC: Erased: " << *I << '\n');
                       Worklist.remove(I);
                       if (I == NextInst)
                         NextInst = NextInst->getNextNode();
+                      Visited.insert(I);
                     }
                   }))
             continue;
diff --git a/llvm/test/Transforms/VectorCombine/X86/pr155543.ll b/llvm/test/Transforms/VectorCombine/X86/pr155543.ll
new file mode 100644
index 0000000000000..b161c0aaf031c
--- /dev/null
+++ b/llvm/test/Transforms/VectorCombine/X86/pr155543.ll
@@ -0,0 +1,15 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes=vector-combine -S -mtriple=x86_64-- | FileCheck %s
+
+; Make sure we don't double delete a dead instruction.
+
+define void @pr155543() {
+; CHECK-LABEL: define void @pr155543() {
+; CHECK-NEXT:    ret void
+;
+  %shuffle1 = shufflevector <4 x double> poison, <4 x double> poison, <8 x i32> <i32 poison, i32 poison, i32 poison, i32 poison, i32 0, i32 1, i32 2, i32 3>
+  %shuffle2 = shufflevector <8 x double> poison, <8 x double> %shuffle1, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 12, i32 13, i32 14, i32 15>
+  %fadd = fadd <8 x double> %shuffle1, zeroinitializer
+  %dead = shufflevector <8 x double> %fadd, <8 x double> %shuffle2, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 12, i32 13, i32 14, i32 15>
+  ret void
+}

[nikic]

LGTM