[llvm] [VectorCombine] scalarizeLoadExtract - don't create scalar loads if any extract is waiting to be erased (PR #129375)

[Simon Pilgrim via llvm-commits]

https://github.com/RKSimon updated https://github.com/llvm/llvm-project/pull/129375

>From 93eacbcfeb70e281fdbec3dd7be6343d1b9c2f15 Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev at redking.me.uk>
Date: Sat, 1 Mar 2025 10:01:42 +0000
Subject: [PATCH 1/3] [VectorCombine] scalarizeLoadExtract - don't create
 scalar loads is any extract is waiting to be erased

If any extract is waiting to be erased, then bail out as this will distort the cost calculation and possibly lead to infinite loops.

Fixes #129373
---
 llvm/lib/Transforms/Vectorize/VectorCombine.cpp      |  5 +++++
 .../X86/load-extractelement-scalarization.ll         | 12 ++++++++++++
 2 files changed, 17 insertions(+)

diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index bb4e1c0f67e79..84cd08cd690d2 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -1579,6 +1579,11 @@ bool VectorCombine::scalarizeLoadExtract(Instruction &I) {
     if (!UI || UI->getParent() != LI->getParent())
       return false;
 
+    // If any extract is waiting to be erased, then bail out as this will
+    // distort the cost calculation and possibly lead to infinite loops.
+    if (isInstructionTriviallyDead(UI))
+      return false;
+
     // Check if any instruction between the load and the extract may modify
     // memory.
     if (LastCheckedInst->comesBefore(UI)) {
diff --git a/llvm/test/Transforms/VectorCombine/X86/load-extractelement-scalarization.ll b/llvm/test/Transforms/VectorCombine/X86/load-extractelement-scalarization.ll
index 0acfeccb92ef7..4e95950591cbc 100644
--- a/llvm/test/Transforms/VectorCombine/X86/load-extractelement-scalarization.ll
+++ b/llvm/test/Transforms/VectorCombine/X86/load-extractelement-scalarization.ll
@@ -24,3 +24,15 @@ define void @multiple_extract(ptr %p) {
   store i32 %e1, ptr %p1, align 4
   ret void
 }
+
+; infinite loop if we fold an extract that is waiting to be erased
+define void @unsued_extract(ptr %p) {
+; CHECK-LABEL: @unsued_extract(
+; CHECK-NEXT:    ret void
+;
+  %load = load <4 x float>, ptr %p, align 8
+  %shuffle0 = shufflevector <4 x float> zeroinitializer, <4 x float> %load, <4 x i32> <i32 0, i32 4, i32 1, i32 5>
+  %shuffle1 = shufflevector <4 x float> %shuffle0, <4 x float> zeroinitializer, <4 x i32> <i32 0, i32 4, i32 poison, i32 poison>
+  %extract = extractelement <4 x float> %load, i64 1
+  ret void
+}

>From ba9ba1958e37ec20abd4432d8c659450f768ccb0 Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev at redking.me.uk>
Date: Sat, 1 Mar 2025 12:36:00 +0000
Subject: [PATCH 2/3] fix typo in test case name

---
 .../VectorCombine/X86/load-extractelement-scalarization.ll    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/test/Transforms/VectorCombine/X86/load-extractelement-scalarization.ll b/llvm/test/Transforms/VectorCombine/X86/load-extractelement-scalarization.ll
index 4e95950591cbc..d46c8c0de4037 100644
--- a/llvm/test/Transforms/VectorCombine/X86/load-extractelement-scalarization.ll
+++ b/llvm/test/Transforms/VectorCombine/X86/load-extractelement-scalarization.ll
@@ -26,8 +26,8 @@ define void @multiple_extract(ptr %p) {
 }
 
 ; infinite loop if we fold an extract that is waiting to be erased
-define void @unsued_extract(ptr %p) {
-; CHECK-LABEL: @unsued_extract(
+define void @unused_extract(ptr %p) {
+; CHECK-LABEL: @unused_extract(
 ; CHECK-NEXT:    ret void
 ;
   %load = load <4 x float>, ptr %p, align 8

>From 52a651178ad1f2ea04a9782525a5494bcea79a74 Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev at redking.me.uk>
Date: Sat, 1 Mar 2025 12:38:21 +0000
Subject: [PATCH 3/3] Just use use_empty as we're only handling
 ExtractElementInst - isInstructionTriviallyDead is only needed for more
 complex general cases

---
 llvm/lib/Transforms/Vectorize/VectorCombine.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index 84cd08cd690d2..2e8458e2cd03d 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -1581,7 +1581,7 @@ bool VectorCombine::scalarizeLoadExtract(Instruction &I) {
 
     // If any extract is waiting to be erased, then bail out as this will
     // distort the cost calculation and possibly lead to infinite loops.
-    if (isInstructionTriviallyDead(UI))
+    if (UI->use_empty())
       return false;
 
     // Check if any instruction between the load and the extract may modify