[InstCombine] Optimize select (X == C), (icmp X Y), false

[aqjune]

Summary

This patch adds transformation select (X == C), (icmp X Y), false ->
select (X == C), (icmp C Y), false.

After the select i1 -> and/or transformation is removed, this optimization
should be explicitly added as a select transformation.
This pattern seems quite common in Transforms/InstCombine.

Alive2 proof:
And: https://alive2.llvm.org/ce/z/dVoKnd
Or: https://alive2.llvm.org/ce/z/3hb9qS

[nikic]

I think this is a case where being in select form favors us, and we can handle this through a more general transform, rather than the special case that is needed for and/or. The general transform is something along these lines: https://gist.github.com/nikic/4bad7e7e7ecda8d9c9e9633d86836b4d I've wanted to do this for a while, and this looks like a good excuse to make it happen :)

[aqjune]

The change makes sense..!
Since the patch is more succinct, I'm in favor of applying the patch than this one.

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

 //===- InstCombineAndOrXor.cpp --------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
if (UnsignedPred == ICmpInst::ICMP_UGT && EqPred == ICmpInst::ICMP_EQ &&
     return Builder.CreateICmpUGE(Base, Offset);
 
   return nullptr;
 }
 
 /// Reduce logic-of-compares with equality to a constant by substituting a
 /// common operand with the constant. Callers are expected to call this with
 /// Cmp0/Cmp1 switched to handle logic op commutativity.
-static Value *foldAndOrOfICmpsWithConstEq(ICmpInst *Cmp0, ICmpInst *Cmp1,
-                                          BinaryOperator &Logic,
-                                          InstCombiner::BuilderTy &Builder,
-                                          const SimplifyQuery &Q) {
-  bool IsAnd = Logic.getOpcode() == Instruction::And;
-  assert((IsAnd || Logic.getOpcode() == Instruction::Or) && "Wrong logic op");
+Value *InstCombinerImpl::foldAndOrOfICmpsWithConstEq(
+    ICmpInst *Cmp0, ICmpInst *Cmp1, const Instruction::BinaryOps Logic,
+    const SimplifyQuery &Q, bool SelectForm) {
+  bool IsAnd = Logic == Instruction::And;
+  assert((IsAnd || Logic == Instruction::Or) && "Wrong logic op");
 
   // Match an equality compare with a non-poison constant as Cmp0.
   // Also, give up if the compare can be constant-folded to avoid looping.
   ICmpInst::Predicate Pred0;
   Value *X;
   Constant *C;
   if (!match(Cmp0, m_ICmp(Pred0, m_Value(X), m_Constant(C))) ||
       !isGuaranteedNotToBeUndefOrPoison(C) || isa<Constant>(X))
Value *InstCombinerImpl::foldAndOrOfICmpsWithConstEq(
   Value *SubstituteCmp = SimplifyICmpInst(Pred1, Y, C, Q);
   if (!SubstituteCmp) {
     // If we need to create a new instruction, require that the old compare can
     // be removed.
     if (!Cmp1->hasOneUse())
       return nullptr;
     SubstituteCmp = Builder.CreateICmp(Pred1, Y, C);
   }
-  return Builder.CreateBinOp(Logic.getOpcode(), Cmp0, SubstituteCmp);
+
+  if (SelectForm) {
+    auto *True = ConstantInt::getTrue(SubstituteCmp->getType());
+    auto *False = ConstantInt::getFalse(SubstituteCmp->getType());
+    return Builder.CreateSelect(
+        Cmp0, Logic == Instruction::And ? SubstituteCmp : True,
+        Logic == Instruction::And ? False : SubstituteCmp);
+  }
+  return Builder.CreateBinOp(Logic, Cmp0, SubstituteCmp);
 }
 
 /// Fold (icmp)&(icmp) if possible.
 Value *InstCombinerImpl::foldAndOfICmps(ICmpInst *LHS, ICmpInst *RHS,
                                         BinaryOperator &And) {
   const SimplifyQuery Q = SQ.getWithInstruction(&And);
 
   // Fold (!iszero(A & K1) & !iszero(A & K2)) ->  (A & (K1 | K2)) == (K1 | K2)
if (LHS->getOperand(0) == RHS->getOperand(0) &&
       return getNewICmpValue(Code, IsSigned, Op0, Op1, Builder);
     }
   }
 
   // handle (roughly):  (icmp eq (A & B), C) & (icmp eq (A & D), E)
   if (Value *V = foldLogOpOfMaskedICmps(LHS, RHS, true, Builder))
     return V;
 
-  if (Value *V = foldAndOrOfICmpsWithConstEq(LHS, RHS, And, Builder, Q))
+  if (Value *V = foldAndOrOfICmpsWithConstEq(LHS, RHS, Instruction::And, Q))
     return V;
-  if (Value *V = foldAndOrOfICmpsWithConstEq(RHS, LHS, And, Builder, Q))
+  if (Value *V = foldAndOrOfICmpsWithConstEq(RHS, LHS, Instruction::And, Q))
     return V;
 
   // E.g. (icmp sge x, 0) & (icmp slt x, n) --> icmp ult x, n
   if (Value *V = simplifyRangeCheck(LHS, RHS, /*Inverted=*/false))
     return V;
 
   // E.g. (icmp slt x, n) & (icmp sge x, 0) --> icmp ult x, n
   if (Value *V = simplifyRangeCheck(RHS, LHS, /*Inverted=*/false))
else if (PredR == ICmpInst::ICMP_EQ && match(RHS1, m_Zero())) {
         A = LHS1;
     }
     if (A && B && B->getType()->isIntOrIntVectorTy())
       return Builder.CreateICmp(
           ICmpInst::ICMP_UGE,
           Builder.CreateAdd(B, Constant::getAllOnesValue(B->getType())), A);
   }
 
-  if (Value *V = foldAndOrOfICmpsWithConstEq(LHS, RHS, Or, Builder, Q))
+  if (Value *V = foldAndOrOfICmpsWithConstEq(LHS, RHS, Instruction::Or, Q))
     return V;
-  if (Value *V = foldAndOrOfICmpsWithConstEq(RHS, LHS, Or, Builder, Q))
+  if (Value *V = foldAndOrOfICmpsWithConstEq(RHS, LHS, Instruction::Or, Q))
     return V;
 
   // E.g. (icmp slt x, 0) | (icmp sgt x, n) --> icmp ugt x, n
   if (Value *V = simplifyRangeCheck(LHS, RHS, /*Inverted=*/true))
     return V;
 
   // E.g. (icmp sgt x, n) | (icmp slt x, 0) --> icmp ugt x, n
   if (Value *V = simplifyRangeCheck(RHS, LHS, /*Inverted=*/true))

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

 //===- InstCombineInternal.h - InstCombine pass internals -------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
private:
   /// \see CastInst::isEliminableCastPair
   Instruction::CastOps isEliminableCastPair(const CastInst *CI1,
                                             const CastInst *CI2);
 
   Value *foldAndOfICmps(ICmpInst *LHS, ICmpInst *RHS, BinaryOperator &And);
   Value *foldOrOfICmps(ICmpInst *LHS, ICmpInst *RHS, BinaryOperator &Or);
   Value *foldXorOfICmps(ICmpInst *LHS, ICmpInst *RHS, BinaryOperator &Xor);
 
+  /// Reduce logic-of-compares with equality to a constant by substituting a
+  /// common operand with the constant. Callers are expected to call this with
+  /// Cmp0/Cmp1 switched to handle logic op commutativity.
+  /// If SelectForm is true, create select i1 instead of and/or.
+  Value *foldAndOrOfICmpsWithConstEq(ICmpInst *Cmp0, ICmpInst *Cmp1,
+                                     const Instruction::BinaryOps Logic,
+                                     const SimplifyQuery &SQ,
+                                     bool SelectForm = false);
+
   /// Optimize (fcmp)&(fcmp) or (fcmp)|(fcmp).
   /// NOTE: Unlike most of instcombine, this returns a Value which should
   /// already be inserted into the function.
   Value *foldLogicOfFCmps(FCmpInst *LHS, FCmpInst *RHS, bool IsAnd);
 
   Value *foldAndOrOfICmpsOfAndWithPow2(ICmpInst *LHS, ICmpInst *RHS,
                                        BinaryOperator &Logic);
   Value *matchSelectFromAndOr(Value *A, Value *B, Value *C, Value *D);

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

 //===- InstCombineSelect.cpp ----------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
if (SelType->isIntOrIntVectorTy(1) &&
     // select a, !a, b -> select !a, b, false
     if (match(TrueVal, m_Not(m_Specific(CondVal))))
       return SelectInst::Create(TrueVal, FalseVal,
                                 ConstantInt::getFalse(SelType));
     // select a, b, !a -> select !a, true, b
     if (match(FalseVal, m_Not(m_Specific(CondVal))))
       return SelectInst::Create(FalseVal, ConstantInt::getTrue(SelType),
                                 TrueVal);
+
+    // select (X == C), (X pred Y), false -> select (X == C), (C pred Y), false
+    // select (X != C), true, (X pred Y)  -> select (X != C), true, (C pred Y)
+    if (auto *CondICmp = dyn_cast<ICmpInst>(CondVal)) {
+      auto *TrueICmp = dyn_cast<ICmpInst>(TrueVal),
+           *FalseICmp = dyn_cast<ICmpInst>(FalseVal);
+      if (match(TrueVal, m_One()) && FalseICmp) {
+        // or case
+        if (Value *V = foldAndOrOfICmpsWithConstEq(CondICmp, FalseICmp,
+                                                   Instruction::Or, SQ, true))
+          return replaceInstUsesWith(SI, V);
+      }
+      if (TrueICmp && match(FalseVal, m_Zero())) {
+        // and case
+        if (Value *V = foldAndOrOfICmpsWithConstEq(CondICmp, TrueICmp,
+                                                   Instruction::And, SQ, true))
+          return replaceInstUsesWith(SI, V);
+      }
+    }
   }
 
   // Selecting between two integer or vector splat integer constants?
   //
   // Note that we don't handle a scalar select of vectors:
   // select i1 %c, <2 x i8> <1, 1>, <2 x i8> <0, 0>
   // because that may need 3 instructions to splat the condition value:
   // extend, insertelement, shufflevector.

llvm/test/Transforms/InstCombine/select-safe-transforms.ll

;
   %lhs = icmp ult i8 %X, %Y
   %res = select i1 %cond, i1 %lhs, i1 false
   ret i1 %res
 }
 
 define i1 @cond_eq_and_const(i8 %X, i8 %Y) {
 ; CHECK-LABEL: @cond_eq_and_const(
 ; CHECK-NEXT:    [[COND:%.*]] = icmp eq i8 [[X:%.*]], 10
-; CHECK-NEXT:    [[LHS:%.*]] = icmp ult i8 [[X]], [[Y:%.*]]
-; CHECK-NEXT:    [[RES:%.*]] = select i1 [[COND]], i1 [[LHS]], i1 false
-; CHECK-NEXT:    ret i1 [[RES]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ugt i8 [[Y:%.*]], 10
+; CHECK-NEXT:    [[TMP2:%.*]] = select i1 [[COND]], i1 [[TMP1]], i1 false
+; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %cond = icmp eq i8 %X, 10
   %lhs = icmp ult i8 %X, %Y
   %res = select i1 %cond, i1 %lhs, i1 false
   ret i1 %res
 }
 
 define i1 @cond_eq_or(i8 %X, i8 %Y, i8 noundef %C) {
 ; CHECK-LABEL: @cond_eq_or(
 ; CHECK-NEXT:    [[COND:%.*]] = icmp ne i8 [[X:%.*]], [[C:%.*]]
 ; CHECK-NEXT:    [[LHS:%.*]] = icmp ult i8 [[X]], [[Y:%.*]]
 ; CHECK-NEXT:    [[RES:%.*]] = select i1 [[COND]], i1 true, i1 [[LHS]]
 ; CHECK-NEXT:    ret i1 [[RES]]
 ;
   %cond = icmp ne i8 %X, %C
   %lhs = icmp ult i8 %X, %Y
   %res = select i1 %cond, i1 true, i1 %lhs
   ret i1 %res
 }
 
 define i1 @cond_eq_or_const(i8 %X, i8 %Y) {
 ; CHECK-LABEL: @cond_eq_or_const(
 ; CHECK-NEXT:    [[COND:%.*]] = icmp ne i8 [[X:%.*]], 10
-; CHECK-NEXT:    [[LHS:%.*]] = icmp ult i8 [[X]], [[Y:%.*]]
-; CHECK-NEXT:    [[RES:%.*]] = select i1 [[COND]], i1 true, i1 [[LHS]]
-; CHECK-NEXT:    ret i1 [[RES]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ugt i8 [[Y:%.*]], 10
+; CHECK-NEXT:    [[TMP2:%.*]] = select i1 [[COND]], i1 true, i1 [[TMP1]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %cond = icmp ne i8 %X, 10
   %lhs = icmp ult i8 %X, %Y
   %res = select i1 %cond, i1 true, i1 %lhs
   ret i1 %res
 }
 
 define i1 @merge_and(i1 %X, i1 %Y) {