[InstCombine] Simplify select operand based on equality condition

[nikic]

Summary

For selects of the type X == Y ? A : B, check if we can simplify A by using the X == Y equality and replace the operand if that's possible. We already try to do this in InstSimplify, but will only fold if the result of the simplification is the same as B, in which case the select can be dropped entirely. Here the select will be retained, just one operand simplified.

As we are performing an actual replacement here, we don't have problems with refinement / poison values.

[nikic]

This partially subsumes the foldSelectBinOpIdentity() fold at https://github.com/llvm/llvm-project/blob/cb19e8c6d192a108b72ab07362921864a9e244f9/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp#L57, which handles a special case of this (where the equality comparison is to the neutral element of a binary operator). However, that fold also handles floating point equalities, which are tricky due to signed zero, so it can't simply be dropped.

[reames]

LGTM

[spatel]

llvm/test/Transforms/InstCombine/select-binop-cmp.ll

566–567

Remove stale comment and rename test since it's no longer "bad",

604–605

Rename test - it's not "bad" now.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

 ///
 /// Consider:
 ///   %cmp = icmp eq i32 %x, 2147483647
 ///   %add = add nsw i32 %x, 1
 ///   %sel = select i1 %cmp, i32 -2147483648, i32 %add
 ///
 /// We can't replace %sel with %add unless we strip away the flags.
 /// TODO: Wrapping flags could be preserved in some cases with better analysis.
-static Value *foldSelectValueEquivalence(SelectInst &Sel, ICmpInst &Cmp,
-                                         const SimplifyQuery &Q) {
+static Instruction *foldSelectValueEquivalence(SelectInst &Sel, ICmpInst &Cmp,
+                                               const SimplifyQuery &Q,
+                                               InstCombiner &IC) {
   if (!Cmp.isEquality())
     return nullptr;
 
   // Canonicalize the pattern to ICMP_EQ by swapping the select operands.
   Value *TrueVal = Sel.getTrueValue(), *FalseVal = Sel.getFalseValue();
-  if (Cmp.getPredicate() == ICmpInst::ICMP_NE)
+  bool Swapped = false;
+  if (Cmp.getPredicate() == ICmpInst::ICMP_NE) {
     std::swap(TrueVal, FalseVal);
+    Swapped = true;
+  }
+
+  // In X == Y ? f(X) : Z, try to evaluate f(X) and replace the operand.
+  // Take care to avoid replacing X == Y ? X : Z with X == Y ? Y : Z, as that
+  // would lead to an infinite replacement cycle.
+  Value *CmpLHS = Cmp.getOperand(0), *CmpRHS = Cmp.getOperand(1);
+  if (TrueVal != CmpLHS)
+    if (Value *V = SimplifyWithOpReplaced(TrueVal, CmpLHS, CmpRHS, Q,
+                                          /* AllowRefinement */ true))
+      return IC.replaceOperand(Sel, Swapped ? 2 : 1, V);
+  if (TrueVal != CmpRHS)
+    if (Value *V = SimplifyWithOpReplaced(TrueVal, CmpRHS, CmpLHS, Q,
+                                          /* AllowRefinement */ true))
+      return IC.replaceOperand(Sel, Swapped ? 2 : 1, V);
 
   auto *FalseInst = dyn_cast<Instruction>(FalseVal);
   if (!FalseInst)
     return nullptr;
 
   // InstSimplify already performed this fold if it was possible subject to
   // current poison-generating flags. Try the transform again with
   // poison-generating flags temporarily dropped.
   bool WasNUW = false, WasNSW = false, WasExact = false;
   if (auto *OBO = dyn_cast<OverflowingBinaryOperator>(FalseVal)) {
     WasNUW = OBO->hasNoUnsignedWrap();
     WasNSW = OBO->hasNoSignedWrap();
     FalseInst->setHasNoUnsignedWrap(false);
     FalseInst->setHasNoSignedWrap(false);
   }
   if (auto *PEO = dyn_cast<PossiblyExactOperator>(FalseVal)) {
     WasExact = PEO->isExact();
     FalseInst->setIsExact(false);
   }
 
   // Try each equivalence substitution possibility.
   // We have an 'EQ' comparison, so the select's false value will propagate.
   // Example:
   // (X == 42) ? 43 : (X + 1) --> (X == 42) ? (X + 1) : (X + 1) --> X + 1
-  Value *CmpLHS = Cmp.getOperand(0), *CmpRHS = Cmp.getOperand(1);
   if (SimplifyWithOpReplaced(FalseVal, CmpLHS, CmpRHS, Q,
                              /* AllowRefinement */ false) == TrueVal ||
       SimplifyWithOpReplaced(FalseVal, CmpRHS, CmpLHS, Q,
                              /* AllowRefinement */ false) == TrueVal) {
-    return FalseVal;
+    return IC.replaceInstUsesWith(Sel, FalseVal);
   }
 
   // Restore poison-generating flags if the transform did not apply.
   if (WasNUW)
     FalseInst->setHasNoUnsignedWrap();
   if (WasNSW)
     FalseInst->setHasNoSignedWrap();
   if (WasExact)
tryToReuseConstantFromSelectInComparison(SelectInst &Sel, ICmpInst &Cmp,
   Sel.swapProfMetadata();
 
   return &Sel;
 }
 
 /// Visit a SelectInst that has an ICmpInst as its first operand.
 Instruction *InstCombinerImpl::foldSelectInstWithICmp(SelectInst &SI,
                                                       ICmpInst *ICI) {
-  if (Value *V = foldSelectValueEquivalence(SI, *ICI, SQ))
-    return replaceInstUsesWith(SI, V);
+  if (Instruction *NewSel = foldSelectValueEquivalence(SI, *ICI, SQ, *this))
+    return NewSel;
 
   if (Instruction *NewSel = canonicalizeMinMaxWithConstant(SI, *ICI, *this))
     return NewSel;
 
   if (Instruction *NewAbs = canonicalizeAbsNabs(SI, *ICI, *this))
     return NewAbs;
 
   if (Instruction *NewAbs = canonicalizeClampLike(SI, *ICI, Builder))

llvm/test/Transforms/InstCombine/rem.ll

 ;
   %rem = urem i8 %x, 129
   ret i8 %rem
 }
 
 define i5 @biggest_divisor(i5 %x) {
 ; CHECK-LABEL: @biggest_divisor(
 ; CHECK-NEXT:    [[DOTNOT:%.*]] = icmp eq i5 [[X:%.*]], -1
-; CHECK-NEXT:    [[TMP1:%.*]] = zext i1 [[DOTNOT]] to i5
-; CHECK-NEXT:    [[REM:%.*]] = add i5 [[TMP1]], [[X]]
+; CHECK-NEXT:    [[REM:%.*]] = select i1 [[DOTNOT]], i5 0, i5 [[X]]
 ; CHECK-NEXT:    ret i5 [[REM]]
 ;
   %rem = urem i5 %x, -1
   ret i5 %rem
 }
 
 define i8 @urem_with_sext_bool_divisor(i1 %x, i8 %y) {
 ; CHECK-LABEL: @urem_with_sext_bool_divisor(

llvm/test/Transforms/InstCombine/select-binop-cmp.ll

 ; CHECK-NEXT:    [[C:%.*]] = select <2 x i1> [[A]], <2 x i8> [[B]], <2 x i8> [[Y:%.*]]
 ; CHECK-NEXT:    ret <2 x i8> [[C]]
 ;
   %A = icmp eq <2 x i8>  %x, <i8 5, i8 3>
   %B = xor <2 x i8>  %x, %z
   %C = select <2 x i1>  %A, <2 x i8>  %B, <2 x i8>  %y
   ret <2 x i8>  %C
 }
 
-; TODO: support for undefs, check for an identity constant does not handle them yet
-define <2 x i8> @select_xor_icmp_vec_bad_2(<2 x i8> %x, <2 x i8> %y, <2 x i8> %z) {
-; CHECK-LABEL: @select_xor_icmp_vec_bad_2(
+define <2 x i8> @select_xor_icmp_vec_undef(<2 x i8> %x, <2 x i8> %y, <2 x i8> %z) {
+; CHECK-LABEL: @select_xor_icmp_vec_undef(
 ; CHECK-NEXT:    [[A:%.*]] = icmp eq <2 x i8> [[X:%.*]], <i8 0, i8 undef>
-; CHECK-NEXT:    [[B:%.*]] = xor <2 x i8> [[X]], [[Z:%.*]]
-; CHECK-NEXT:    [[C:%.*]] = select <2 x i1> [[A]], <2 x i8> [[B]], <2 x i8> [[Y:%.*]]
+; CHECK-NEXT:    [[C:%.*]] = select <2 x i1> [[A]], <2 x i8> [[Z:%.*]], <2 x i8> [[Y:%.*]]
 ; CHECK-NEXT:    ret <2 x i8> [[C]]
 ;
   %A = icmp eq <2 x i8>  %x, <i8 0, i8 undef>
   %B = xor <2 x i8>  %x, %z
   %C = select <2 x i1>  %A, <2 x i8>  %B, <2 x i8>  %y
   ret <2 x i8>  %C
 }
 

 ; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i32 [[B]], i32 [[Y:%.*]]
 ; CHECK-NEXT:    ret i32 [[C]]
 ;
   %A = icmp eq i32 %x, 1
   %B = add i32 %x, %z
   %C = select i1 %A, i32 %B, i32 %y
   ret i32 %C
 }
 
-define i32 @select_and_icmp_bad(i32 %x, i32 %y, i32 %z) {
-; CHECK-LABEL: @select_and_icmp_bad(
+define i32 @select_and_icmp_zero(i32 %x, i32 %y, i32 %z) {
+; CHECK-LABEL: @select_and_icmp_zero(
 ; CHECK-NEXT:    [[A:%.*]] = icmp eq i32 [[X:%.*]], 0
-; CHECK-NEXT:    [[B:%.*]] = and i32 [[X]], [[Z:%.*]]
-; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i32 [[B]], i32 [[Y:%.*]]
+; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i32 0, i32 [[Y:%.*]]
 ; CHECK-NEXT:    ret i32 [[C]]
 ;
   %A = icmp eq i32 %x, 0
   %B = and i32 %x, %z
   %C = select i1 %A, i32 %B, i32 %y
   ret i32 %C
 }
 

llvm/test/Transforms/InstCombine/select.ll

;
   %shifted = lshr i32 %arg, %trailing
   %r1.sroa.0.1 = select i1 %cmp, i32 0, i32 %shifted
   ret i32 %r1.sroa.0.1
 }
 
 define i8 @select_replacement_add_eq(i8 %x, i8 %y) {
 ; CHECK-LABEL: @select_replacement_add_eq(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X:%.*]], 1
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], 1
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[ADD]], i8 [[Y:%.*]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 2, i8 [[Y:%.*]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %cmp = icmp eq i8 %x, 1
   %add = add i8 %x, 1
   %sel = select i1 %cmp, i8 %add, i8 %y
   ret i8 %sel
 }
 
 define i8 @select_replacement_add_ne(i8 %x, i8 %y) {
 ; CHECK-LABEL: @select_replacement_add_ne(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i8 [[X:%.*]], 1
 ; CHECK-NEXT:    call void @use(i1 [[CMP]])
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], 1
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[Y:%.*]], i8 [[ADD]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[Y:%.*]], i8 2
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %cmp = icmp ne i8 %x, 1
   call void @use(i1 %cmp)
   %add = add i8 %x, 1
   %sel = select i1 %cmp, i8 %y, i8 %add
   ret i8 %sel
 }
 
 define i8 @select_replacement_add_nuw(i8 %x, i8 %y) {
 ; CHECK-LABEL: @select_replacement_add_nuw(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X:%.*]], 1
-; CHECK-NEXT:    [[ADD:%.*]] = add nuw i8 [[X]], 1
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[ADD]], i8 [[Y:%.*]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 2, i8 [[Y:%.*]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %cmp = icmp eq i8 %x, 1
   %add = add nuw i8 %x, 1
   %sel = select i1 %cmp, i8 %add, i8 %y
   ret i8 %sel
 }
 
 define i8 @select_replacement_sub(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @select_replacement_sub(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[SUB:%.*]] = sub i8 [[X]], [[Y]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[SUB]], i8 [[Z:%.*]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 0, i8 [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %cmp = icmp eq i8 %x, %y
   %sub = sub i8 %x, %y
   %sel = select i1 %cmp, i8 %sub, i8 %z
   ret i8 %sel
 }
 
 define i8 @select_replacement_shift(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @select_replacement_shift(
 ; CHECK-NEXT:    [[SHR:%.*]] = lshr exact i8 [[X:%.*]], 1
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[SHR]], [[Y:%.*]]
-; CHECK-NEXT:    [[SHL:%.*]] = shl i8 [[Y]], 1
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[SHL]], i8 [[Z:%.*]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[X]], i8 [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %shr = lshr exact i8 %x, 1
   %cmp = icmp eq i8 %shr, %y
   %shl = shl i8 %y, 1
   %sel = select i1 %cmp, i8 %shl, i8 %z
   ret i8 %sel
 }