[InstCombine] Improve select equiv fold for true condition value

Refactor with a lambda function to address both false and true value.

Author: vfdff

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

@@ -511,26 +511,36 @@ static bool simplifySeqSelectWithSameCond(SelectInst &SI,
                                           const SimplifyQuery &SQ,
                                           InstCombinerImpl &IC) {
   Value *CondVal = SI.getCondition();
-  Type *CondType = CondVal->getType();
-  SelectInst *SINext = &SI;
-  Type *SelType = SINext->getType();
-  Value *FalseVal = SINext->getFalseValue();
-  Value *CondNext;
-  Value *FalseNext;
-  while (match(FalseVal,
-               m_Select(m_Value(CondNext), m_Value(), m_Value(FalseNext)))) {
+  auto trySimplifySeqSelect = [=, &SI, &IC](unsigned OpIndex) {
+    assert((OpIndex == 1 || OpIndex == 2) && "Unexpected operand index");
+    SelectInst *SINext = &SI;
+    Type *SelType = SINext->getType();
+    Value *ValOp = SINext->getOperand(OpIndex);
+    Value *CondNext;
     // Only support the type of select is an integer type as float type need
     // address FMF flag.
-    if (CondNext == CondVal && SelType->isIntOrIntVectorTy() &&
-        SINext->hasOneUse()) {
-      IC.replaceOperand(*SINext, 2, FalseNext);
-      return true;
+    while (match(ValOp, m_Select(m_Value(CondNext), m_Value(), m_Value()))) {
+      if (CondNext == CondVal && SelType->isIntOrIntVectorTy() &&
+          SINext->hasOneUse()) {
+        IC.replaceOperand(*SINext, OpIndex,
+                          cast<SelectInst>(ValOp)->getOperand(OpIndex));
+        return true;
+      }
+
+      SINext = cast<SelectInst>(ValOp);
+      SelType = SINext->getType();
+      ValOp = SINext->getOperand(OpIndex);
     }
+    return false;
+  };
 
-    SINext = cast<SelectInst>(FalseVal);
-    SelType = SINext->getType();
-    FalseVal = SINext->getFalseValue();
-  }
+  // Try to simplify the true value of select.
+  if (trySimplifySeqSelect(/*OpIndex=*/1))
+    return true;
+
+  // Try to simplify the false value of select.
+  if (trySimplifySeqSelect(/*OpIndex=*/2))
+    return true;
 
   return false;
 }

llvm/test/Transforms/InstCombine/select.ll

@@ -4539,9 +4539,8 @@ define i32 @sequence_select_with_same_cond_false(i1 %c1, i1 %c2){
 
 define i32 @sequence_select_with_same_cond_true(i1 %c1, i1 %c2){
 ; CHECK-LABEL: @sequence_select_with_same_cond_true(
-; CHECK-NEXT:    [[S1:%.*]] = select i1 [[C1:%.*]], i32 45, i32 23
-; CHECK-NEXT:    [[S2:%.*]] = select i1 [[C2:%.*]], i32 [[S1]], i32 666
-; CHECK-NEXT:    [[S3:%.*]] = select i1 [[C1]], i32 [[S2]], i32 789
+; CHECK-NEXT:    [[S2:%.*]] = select i1 [[C2:%.*]], i32 45, i32 666
+; CHECK-NEXT:    [[S3:%.*]] = select i1 [[C1:%.*]], i32 [[S2]], i32 789
 ; CHECK-NEXT:    ret i32 [[S3]]
 ;
   %s1 = select i1 %c1, i32 45, i32 23
@@ -4639,10 +4638,12 @@ define i8 @test_replace_freeze_oneuse(i1 %x, i8 %y) {
   ret i8 %sel
 }
 
+; first, %sel2 change into select i1 %cond1, i8 %sel0, i8 3, the the %sel1 is OneUse
+; second, %sel1 change into select i1 %cond1, i8 %a, i8 2
 define i8 @sequence_select_with_same_cond_multi_arms(i1 %cond0, i1 %cond1, i8 %a, i8 %b) {
 ; CHECK-LABEL: @sequence_select_with_same_cond_multi_arms(
 ; CHECK-NEXT:    [[SEL0:%.*]] = select i1 [[COND0:%.*]], i8 [[A:%.*]], i8 [[B:%.*]]
-; CHECK-NEXT:    [[SEL1:%.*]] = select i1 [[COND1:%.*]], i8 [[SEL0]], i8 2
+; CHECK-NEXT:    [[SEL1:%.*]] = select i1 [[COND1:%.*]], i8 [[A]], i8 2
 ; CHECK-NEXT:    [[SEL2:%.*]] = select i1 [[COND1]], i8 [[SEL0]], i8 3
 ; CHECK-NEXT:    [[SEL3:%.*]] = select i1 [[COND0]], i8 [[SEL1]], i8 [[SEL2]]
 ; CHECK-NEXT:    ret i8 [[SEL3]]