[llvm] [InstCombine] Fold its select user into select (PR #83405)

[via llvm-commits]

https://github.com/vfdff updated https://github.com/llvm/llvm-project/pull/83405

>From 2b792fecf90d91080040cbc6c3eef6b5829ec05b Mon Sep 17 00:00:00 2001
From: zhongyunde 00443407 <zhongyunde at huawei.com>
Date: Wed, 28 Feb 2024 23:21:30 -0500
Subject: [PATCH] [InstCombine] Fold its select user into select

Fold select if the user of its select user indicates the condition

Fix https://github.com/llvm/llvm-project/issues/83225
---
 .../InstCombine/InstCombineSelect.cpp         | 39 +++++++++++++++++++
 llvm/test/Transforms/InstCombine/select.ll    | 29 ++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 71fa9b9ba41ebb..be3463d6a49cac 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -500,6 +500,42 @@ static bool isSelect01(const APInt &C1I, const APInt &C2I) {
   return C1I.isOne() || C1I.isAllOnes() || C2I.isOne() || C2I.isAllOnes();
 }
 
+/// Try to simplify a select instruction when the user of its select user
+/// indicates the condition.
+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)))) {
+    // If the type of select is not an integer type or if the condition and
+    // the selection type are not both scalar nor both vector types, there is no
+    // point in attempting to match these patterns.
+    if (!isa<Constant>(CondVal) && SelType->isIntOrIntVectorTy() &&
+        CondType->isVectorTy() == SelType->isVectorTy())
+      if (CondNext == CondVal) {
+        if (Value *S = simplifyWithOpReplaced(
+                FalseVal, CondVal, ConstantInt::getFalse(CondType), SQ,
+                /* AllowRefinement */ true)) {
+          IC.replaceOperand(*SINext, 2, S);
+          return true;
+        }
+      }
+
+    SINext = cast<SelectInst>(FalseVal);
+    FalseVal = SINext->getFalseValue();
+  }
+
+  return false;
+}
+
 /// Try to fold the select into one of the operands to allow further
 /// optimization.
 Instruction *InstCombinerImpl::foldSelectIntoOp(SelectInst &SI, Value *TrueVal,
@@ -557,6 +593,9 @@ Instruction *InstCombinerImpl::foldSelectIntoOp(SelectInst &SI, Value *TrueVal,
   if (Instruction *R = TryFoldSelectIntoOp(SI, FalseVal, TrueVal, true))
     return R;
 
+  if (simplifySeqSelectWithSameCond(SI, SQ, *this))
+    return &SI;
+
   return nullptr;
 }
 
diff --git a/llvm/test/Transforms/InstCombine/select.ll b/llvm/test/Transforms/InstCombine/select.ll
index c5f1b77c6d7404..81b84341e2d8f7 100644
--- a/llvm/test/Transforms/InstCombine/select.ll
+++ b/llvm/test/Transforms/InstCombine/select.ll
@@ -3672,3 +3672,32 @@ define i32 @src_select_xxory_eq0_xorxy_y(i32 %x, i32 %y) {
   %cond = select i1 %xor0, i32 %xor, i32 %y
   ret i32 %cond
 }
+
+define i32 @sequence_select_with_same_cond (i1 %c1, i1 %c2){
+; CHECK-LABEL: @sequence_select_with_same_cond(
+; CHECK-NEXT:    [[S2:%.*]] = select i1 [[C2:%.*]], i32 666, i32 45
+; CHECK-NEXT:    [[S3:%.*]] = select i1 [[C1:%.*]], i32 789, i32 [[S2]]
+; CHECK-NEXT:    ret i32 [[S3]]
+;
+  %s1 = select i1 %c1, i32 23, i32 45
+  %s2 = select i1 %c2, i32 666, i32 %s1
+  %s3 = select i1 %c1, i32 789, i32 %s2
+  ret i32 %s3
+}
+
+declare void @use32(i32)
+
+define i32 @sequence_select_with_same_cond_extra_use (i1 %c1, i1 %c2){
+; CHECK-LABEL: @sequence_select_with_same_cond_extra_use(
+; CHECK-NEXT:    [[S1:%.*]] = select i1 [[C1:%.*]], i32 23, i32 45
+; CHECK-NEXT:    call void @use32(i32 [[S1]])
+; CHECK-NEXT:    [[S2:%.*]] = select i1 [[C2:%.*]], i32 666, i32 45
+; CHECK-NEXT:    [[S3:%.*]] = select i1 [[C1]], i32 789, i32 [[S2]]
+; CHECK-NEXT:    ret i32 [[S3]]
+;
+  %s1 = select i1 %c1, i32 23, i32 45
+  call void @use32(i32 %s1)
+  %s2 = select i1 %c2, i32 666, i32 %s1
+  %s3 = select i1 %c1, i32 789, i32 %s2
+  ret i32 %s3
+}