[llvm] [InstCombine] Fold "extract (select (cond, insert(agg, elem), FV))" (PR #115969)

[via llvm-commits]

https://github.com/weiguozhi created https://github.com/llvm/llvm-project/pull/115969

Do the transform
```
     extract (select (cond, insert(agg, elem), FV))
 ->  select (cond, elem, extract(FV))
```

>From a9999a7696b90b7468197f309f796354ce7fa4d5 Mon Sep 17 00:00:00 2001
From: Guozhi Wei <carrot at google.com>
Date: Tue, 12 Nov 2024 16:39:22 -0800
Subject: [PATCH] [InstCombine] Fold "extract (select (cond, insert(agg, elem),
 FV))"

Do the transform
     extract (select (cond, insert(agg, elem), FV))
 ->  select (cond, elem, extract(FV))
---
 .../InstCombine/InstructionCombining.cpp      | 20 +++++++++++-
 .../InstCombine/extract-select-agg.ll         | 31 ++++++++++++++-----
 2 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 2a54390c0f1882d..5914a7aa422e11f 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -1692,7 +1692,25 @@ Instruction *InstCombinerImpl::FoldOpIntoSelect(Instruction &Op, SelectInst *SI,
 
   Value *TV = SI->getTrueValue();
   Value *FV = SI->getFalseValue();
-  if (!(isa<Constant>(TV) || isa<Constant>(FV)))
+  if ((Op.getOpcode() == Instruction::ExtractValue) &&
+      (isa<InsertValueInst>(TV) || isa<InsertValueInst>(FV))) {
+    // extract (select (cond, insert(agg, elem), FV))
+    // -> select (cond, elem, extract(FV))
+    ExtractValueInst *EV = dyn_cast<ExtractValueInst>(&Op);
+    Value *NewTV = simplifyExtractValueInst(TV, EV->getIndices(),
+                                            SQ.getWithInstruction(&Op));
+    Value *NewFV = simplifyExtractValueInst(FV, EV->getIndices(),
+                                            SQ.getWithInstruction(&Op));
+    if (!NewTV && !NewFV)
+      return nullptr;
+    Builder.SetInsertPoint(SI);
+    if (!NewTV)
+      NewTV = Builder.CreateExtractValue(TV, EV->getIndices());
+    if (!NewFV)
+      NewFV = Builder.CreateExtractValue(FV, EV->getIndices());
+    return SelectInst::Create(SI->getCondition(), NewTV, NewFV, "", nullptr,
+                              SI);
+  } else if (!(isa<Constant>(TV) || isa<Constant>(FV)))
     return nullptr;
 
   // Bool selects with constant operands can be folded to logical ops.
diff --git a/llvm/test/Transforms/InstCombine/extract-select-agg.ll b/llvm/test/Transforms/InstCombine/extract-select-agg.ll
index 6ba6b1a575601d0..2813870feeb9319 100644
--- a/llvm/test/Transforms/InstCombine/extract-select-agg.ll
+++ b/llvm/test/Transforms/InstCombine/extract-select-agg.ll
@@ -56,14 +56,9 @@ define void @test_select_agg_multiuse(i1 %cond, i64 %v1, i64 %v2, i64 %v3, i64 %
 ; CHECK-LABEL: define void @test_select_agg_multiuse(
 ; CHECK-SAME: i1 [[COND:%.*]], i64 [[V1:%.*]], i64 [[V2:%.*]], i64 [[V3:%.*]], i64 [[V4:%.*]]) {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[A0:%.*]] = insertvalue { i64, i64 } poison, i64 [[V1]], 0
-; CHECK-NEXT:    [[A1:%.*]] = insertvalue { i64, i64 } [[A0]], i64 [[V2]], 1
-; CHECK-NEXT:    [[B0:%.*]] = insertvalue { i64, i64 } poison, i64 [[V3]], 0
-; CHECK-NEXT:    [[B1:%.*]] = insertvalue { i64, i64 } [[B0]], i64 [[V4]], 1
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], { i64, i64 } [[A1]], { i64, i64 } [[B1]]
-; CHECK-NEXT:    [[X:%.*]] = extractvalue { i64, i64 } [[SEL]], 0
+; CHECK-NEXT:    [[X:%.*]] = select i1 [[COND]], i64 [[V1]], i64 [[V3]]
 ; CHECK-NEXT:    call void @use(i64 [[X]])
-; CHECK-NEXT:    [[Y:%.*]] = extractvalue { i64, i64 } [[SEL]], 1
+; CHECK-NEXT:    [[Y:%.*]] = select i1 [[COND]], i64 [[V2]], i64 [[V4]]
 ; CHECK-NEXT:    call void @use(i64 [[Y]])
 ; CHECK-NEXT:    ret void
 ;
@@ -81,3 +76,25 @@ entry:
 }
 
 declare void @use(i64)
+
+define i64 @test_extract_select_insert(ptr %p1, i64 %v) {
+; CHECK-LABEL: define i64 @test_extract_select_insert(
+; CHECK-SAME: ptr [[P1:%.*]], i64 [[V:%.*]]) {
+; CHECK-NEXT:    [[CALL:%.*]] = call { ptr, i64 } @foo()
+; CHECK-NEXT:    [[ELM1:%.*]] = extractvalue { ptr, i64 } [[CALL]], 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[ELM1]], [[V]]
+; CHECK-NEXT:    [[TMP1:%.*]] = extractvalue { ptr, i64 } [[CALL]], 1
+; CHECK-NEXT:    [[RES:%.*]] = select i1 [[CMP]], i64 4294967294, i64 [[TMP1]]
+; CHECK-NEXT:    ret i64 [[RES]]
+;
+  %call = call { ptr, i64 } @foo()
+  %elm1 = extractvalue { ptr, i64 } %call, 1
+  %cmp = icmp eq i64 %elm1, %v
+  %fca0 = insertvalue { ptr, i64 } poison, ptr %p1, 0
+  %fca1 = insertvalue { ptr, i64 } %fca0, i64 4294967294, 1
+  %select = select i1 %cmp, { ptr, i64 } %fca1, { ptr, i64 } %call
+  %res = extractvalue { ptr, i64 } %select, 1
+  ret i64 %res
+}
+
+declare { ptr, i64 } @foo()