[InstCombine] Fold "extract (select (cond, insert(agg, elem), FV))"

[weiguozhi]

Do the transform

     extract (select (cond, insert(agg, elem), FV))
 ->  select (cond, elem, extract(FV))

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: None (weiguozhi)

Changes

Do the transform

     extract (select (cond, insert(agg, elem), FV))
 ->  select (cond, elem, extract(FV))

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstructionCombining.cpp (+19-1)
• (modified) llvm/test/Transforms/InstCombine/extract-select-agg.ll (+24-7)

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 2a54390c0f1882..5914a7aa422e11 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 6ba6b1a575601d..2813870feeb931 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()

[dtcxzyw]

LG. Please wait for additional approval from other reviewers.

We may generalize to use simplifyInstruction for other instructions in the future.

[nikic]

We may generalize to use simplifyInstruction for other instructions in the future.

I'd rather directly do that than add a special case for extractelement. It's easier to generalize this than not. I did a quick test at #116073 and it doesn't seem to affect compile-time, which would have been my main concern.

[dtcxzyw]

We may generalize to use simplifyInstruction for other instructions in the future.

I'd rather directly do that than add a special case for extractelement. It's easier to generalize this than not. I did a quick test at #116073 and it doesn't seem to affect compile-time, which would have been my main concern.

@nikic Is it okay to just keep these negative tests and make this patch NFC? Positive tests have been added in 2ca25ab.

[nikic]

We may generalize to use simplifyInstruction for other instructions in the future.

I'd rather directly do that than add a special case for extractelement. It's easier to generalize this than not. I did a quick test at #116073 and it doesn't seem to affect compile-time, which would have been my main concern.

@nikic Is it okay to just keep these negative tests and make this patch NFC? Positive tests have been added in 2ca25ab.

Sure

[weiguozhi]

Deprecate this pr because more general version has been committed in 9a844a3.