[InstCombine] Fold select of and/or(icmp, other_cond) to select of other_cond

XChy · XChy · commit a56467d18d35 · 2023-12-29T21:51:20.000+08:00
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -442,7 +442,7 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
   // into simplier select instruction using isImpliedCondition.
   Instruction *foldAndOrOfSelectUsingImpliedCond(Value *Op, SelectInst &SI,
                                                  bool IsAnd);
-
+  Instruction *foldSelectOfAndOr(SelectInst &SI);
   Instruction *hoistFNegAboveFMulFDiv(Value *FNegOp, Instruction &FMFSource);
 
 public:
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -2683,6 +2683,56 @@ Instruction *InstCombinerImpl::foldAndOrOfSelectUsingImpliedCond(Value *Op,
   }
 }
 
+Instruction *InstCombinerImpl::foldSelectOfAndOr(SelectInst &SI) {
+  Value *CondVal = SI.getCondition();
+  Value *CondLHS, *CondRHS;
+  bool IsBitwise = !isa<SelectInst>(CondVal);
+  bool IsAnd;
+
+  if (match(CondVal, m_LogicalOr(m_Value(CondLHS), m_Value(CondRHS))))
+    IsAnd = false;
+  else if (match(CondVal, m_LogicalAnd(m_Value(CondLHS), m_Value(CondRHS))))
+    IsAnd = true;
+  else
+    return nullptr;
+
+  Value *TrueVal = SI.getTrueValue(), *FalseVal = SI.getFalseValue();
+
+  // %cmp = icmp eq i32 %a, %b
+  // %cond = or i1 %other_cond, %cmp
+  // %select = select i1 %cond, i32 %a, i32 %b
+  // =>
+  // %select = select i1 %other_cond, i32 %a, i32 %b
+
+  // Or for an inverted version, we fold it like:
+  // %cmp = icmp ne i32 %a, %b
+  // %cond = and i1 %other_cond, %cmp
+  // %select = select i1 %cond, i32 %a, i32 %b
+  // =>
+  // %select = select i1 %other_cond, i32 %a, i32 %b
+
+  auto FoldAndOrOfEquality = [&](Value *Cmp, Value *OtherCond,
+                                 bool NeedToFreeze) -> Instruction * {
+    ICmpInst::Predicate Pred;
+    if (!match(Cmp, m_c_ICmp(Pred, m_Specific(TrueVal), m_Specific(FalseVal))))
+      return nullptr;
+    if (Pred != (IsAnd ? ICmpInst::ICMP_NE : ICmpInst::ICMP_EQ))
+      return nullptr;
+    if (NeedToFreeze)
+      OtherCond = Builder.CreateFreeze(OtherCond);
+    return SelectInst::Create(OtherCond, TrueVal, FalseVal);
+  };
+
+  // Order doesn't matter for bitwise and/or. We care about FreezeInst only when
+  // and/or is logical.
+  if (Instruction *Ret = FoldAndOrOfEquality(CondLHS, CondRHS, !IsBitwise))
+    return Ret;
+  if (Instruction *Ret = FoldAndOrOfEquality(CondRHS, CondLHS, false))
+    return Ret;
+
+  return nullptr;
+}
+
 // Canonicalize select with fcmp to fabs(). -0.0 makes this tricky. We need
 // fast-math-flags (nsz) or fsub with +0.0 (not fneg) for this to work.
 static Instruction *foldSelectWithFCmpToFabs(SelectInst &SI,
@@ -3409,6 +3459,9 @@ Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
     }
   }
 
+  if (Instruction *I = foldSelectOfAndOr(SI))
+    return I;
+
   auto *SIFPOp = dyn_cast<FPMathOperator>(&SI);
 
   if (auto *FCmp = dyn_cast<FCmpInst>(CondVal)) {
diff --git a/llvm/test/Transforms/InstCombine/select-and-or.ll b/llvm/test/Transforms/InstCombine/select-and-or.ll
@@ -968,9 +968,7 @@ define i1 @or_and3_wrong_operand(i1 %a, i1 %b, i32 %x, i32 %y, i1 %d) {
 
 define i32 @or_eq_a_b(i1 %other_cond, i32 %a, i32 %b)  {
 ; CHECK-LABEL: @or_eq_a_b(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    [[COND:%.*]] = or i1 [[CMP]], [[OTHER_COND:%.*]]
-; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[COND]], i32 [[A]], i32 [[B]]
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[OTHER_COND:%.*]], i32 [[A:%.*]], i32 [[B:%.*]]
 ; CHECK-NEXT:    ret i32 [[SELECT]]
 ;
   %cmp = icmp eq i32 %a, %b
@@ -981,9 +979,7 @@ define i32 @or_eq_a_b(i1 %other_cond, i32 %a, i32 %b)  {
 
 define i32 @and_ne_a_b(i1 %other_cond, i32 %a, i32 %b)  {
 ; CHECK-LABEL: @and_ne_a_b(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    [[COND:%.*]] = and i1 [[CMP]], [[OTHER_COND:%.*]]
-; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[COND]], i32 [[A]], i32 [[B]]
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[OTHER_COND:%.*]], i32 [[A:%.*]], i32 [[B:%.*]]
 ; CHECK-NEXT:    ret i32 [[SELECT]]
 ;
   %cmp = icmp ne i32 %a, %b
@@ -994,9 +990,7 @@ define i32 @and_ne_a_b(i1 %other_cond, i32 %a, i32 %b)  {
 
 define i32 @or_eq_a_b_commuted(i1 %other_cond, i32 %a, i32 %b)  {
 ; CHECK-LABEL: @or_eq_a_b_commuted(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    [[COND:%.*]] = or i1 [[CMP]], [[OTHER_COND:%.*]]
-; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[COND]], i32 [[B]], i32 [[A]]
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[OTHER_COND:%.*]], i32 [[B:%.*]], i32 [[A:%.*]]
 ; CHECK-NEXT:    ret i32 [[SELECT]]
 ;
   %cmp = icmp eq i32 %a, %b
@@ -1007,9 +1001,7 @@ define i32 @or_eq_a_b_commuted(i1 %other_cond, i32 %a, i32 %b)  {
 
 define i32 @and_ne_a_b_commuted(i1 %other_cond, i32 %a, i32 %b)  {
 ; CHECK-LABEL: @and_ne_a_b_commuted(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    [[COND:%.*]] = and i1 [[CMP]], [[OTHER_COND:%.*]]
-; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[COND]], i32 [[B]], i32 [[A]]
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[OTHER_COND:%.*]], i32 [[B:%.*]], i32 [[A:%.*]]
 ; CHECK-NEXT:    ret i32 [[SELECT]]
 ;
   %cmp = icmp ne i32 %a, %b
@@ -1022,9 +1014,7 @@ define i32 @or_eq_different_operands(i32 %a, i32 %b, i32 %c)  {
 ; CHECK-LABEL: @or_eq_different_operands(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[A:%.*]], [[C:%.*]]
-; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq i32 [[B:%.*]], [[A]]
-; CHECK-NEXT:    [[COND:%.*]] = or i1 [[CMP]], [[CMP1]]
-; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[COND]], i32 [[A]], i32 [[B]]
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[CMP]], i32 [[A]], i32 [[B:%.*]]
 ; CHECK-NEXT:    ret i32 [[SELECT]]
 ;
 entry:
@@ -1041,7 +1031,7 @@ define i32 @or_eq_a_b_multi_use(i1 %other_cond, i32 %a, i32 %b)  {
 ; CHECK-NEXT:    [[COND:%.*]] = or i1 [[CMP]], [[OTHER_COND:%.*]]
 ; CHECK-NEXT:    call void @use(i1 [[CMP]])
 ; CHECK-NEXT:    call void @use(i1 [[COND]])
-; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[COND]], i32 [[A]], i32 [[B]]
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[OTHER_COND]], i32 [[A]], i32 [[B]]
 ; CHECK-NEXT:    ret i32 [[SELECT]]
 ;
   %cmp = icmp eq i32 %a, %b
@@ -1055,9 +1045,7 @@ define i32 @or_eq_a_b_multi_use(i1 %other_cond, i32 %a, i32 %b)  {
 define <2 x i32> @or_eq_a_b_vec(<2 x i1> %other_cond ,<2 x i32> %a, <2 x i32> %b)  {
 ; CHECK-LABEL: @or_eq_a_b_vec(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <2 x i32> [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    [[COND:%.*]] = or <2 x i1> [[CMP]], [[OTHER_COND:%.*]]
-; CHECK-NEXT:    [[SELECT:%.*]] = select <2 x i1> [[COND]], <2 x i32> [[A]], <2 x i32> [[B]]
+; CHECK-NEXT:    [[SELECT:%.*]] = select <2 x i1> [[OTHER_COND:%.*]], <2 x i32> [[A:%.*]], <2 x i32> [[B:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[SELECT]]
 ;
 entry:
@@ -1116,9 +1104,7 @@ entry:
 
 define i32 @logical_or_eq_a_b(i1 %other_cond, i32 %a, i32 %b)  {
 ; CHECK-LABEL: @logical_or_eq_a_b(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    [[OR_COND:%.*]] = select i1 [[OTHER_COND:%.*]], i1 true, i1 [[CMP]]
-; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[OR_COND]], i32 [[A]], i32 [[B]]
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[OTHER_COND:%.*]], i32 [[A:%.*]], i32 [[B:%.*]]
 ; CHECK-NEXT:    ret i32 [[SELECT]]
 ;
   %cmp = icmp eq i32 %a, %b
@@ -1129,9 +1115,8 @@ define i32 @logical_or_eq_a_b(i1 %other_cond, i32 %a, i32 %b)  {
 
 define i32 @logical_commuted_or_eq_a_b(i1 %other_cond, i32 %a, i32 %b)  {
 ; CHECK-LABEL: @logical_commuted_or_eq_a_b(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    [[OR_COND:%.*]] = select i1 [[CMP]], i1 true, i1 [[OTHER_COND:%.*]]
-; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[OR_COND]], i32 [[A]], i32 [[B]]
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i1 [[OTHER_COND:%.*]]
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[TMP1]], i32 [[A:%.*]], i32 [[B:%.*]]
 ; CHECK-NEXT:    ret i32 [[SELECT]]
 ;
   %cmp = icmp eq i32 %a, %b
@@ -1142,9 +1127,7 @@ define i32 @logical_commuted_or_eq_a_b(i1 %other_cond, i32 %a, i32 %b)  {
 
 define i32 @logical_and_ne_a_b(i1 %other_cond, i32 %a, i32 %b)  {
 ; CHECK-LABEL: @logical_and_ne_a_b(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    [[OR_COND:%.*]] = select i1 [[OTHER_COND:%.*]], i1 [[CMP]], i1 false
-; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[OR_COND]], i32 [[A]], i32 [[B]]
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[OTHER_COND:%.*]], i32 [[A:%.*]], i32 [[B:%.*]]
 ; CHECK-NEXT:    ret i32 [[SELECT]]
 ;
   %cmp = icmp ne i32 %a, %b
@@ -1155,9 +1138,8 @@ define i32 @logical_and_ne_a_b(i1 %other_cond, i32 %a, i32 %b)  {
 
 define i32 @logical_commuted_and_ne_a_b(i1 %other_cond, i32 %a, i32 %b)  {
 ; CHECK-LABEL: @logical_commuted_and_ne_a_b(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    [[OR_COND:%.*]] = select i1 [[CMP]], i1 [[OTHER_COND:%.*]], i1 false
-; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[OR_COND]], i32 [[A]], i32 [[B]]
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i1 [[OTHER_COND:%.*]]
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[TMP1]], i32 [[A:%.*]], i32 [[B:%.*]]
 ; CHECK-NEXT:    ret i32 [[SELECT]]
 ;
   %cmp = icmp ne i32 %a, %b
diff --git a/llvm/test/Transforms/Reassociate/basictest.ll b/llvm/test/Transforms/Reassociate/basictest.ll
@@ -239,10 +239,8 @@ define i32 @test14(i32 %X1, i32 %X2) {
 
 define i32 @test15(i32 %X1, i32 %X2, i32 %X3) {
 ; CHECK-LABEL: @test15(
-; CHECK-NEXT:    [[A:%.*]] = icmp ne i32 [[X1:%.*]], 0
 ; CHECK-NEXT:    [[B:%.*]] = icmp slt i32 [[X2:%.*]], [[X3:%.*]]
-; CHECK-NEXT:    [[C:%.*]] = and i1 [[A]], [[B]]
-; CHECK-NEXT:    [[D:%.*]] = select i1 [[C]], i32 [[X1]], i32 0
+; CHECK-NEXT:    [[D:%.*]] = select i1 [[B]], i32 [[X1:%.*]], i32 0
 ; CHECK-NEXT:    ret i32 [[D]]
 ;
   %A = icmp ne i32 %X1, 0
