Fold Select with a Const Operand to BinOp

veera-sivarajan · veera-sivarajan · commit 2a572a832ed1 · 2024-11-22T02:52:09.000Z
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -1898,6 +1898,69 @@ static Instruction *foldSelectICmpEq(SelectInst &SI, ICmpInst *ICI,
   return nullptr;
 }
 
+// Turn select (Cmp X C1) (BOp X C2) C3
+//   -> BOp (min/max X C1) C2
+// iff C3 == BOp C1 C2
+// This allows for better canonicalization.
+static Value *canonicalizeConstToBOp(ICmpInst *Cmp, Value *TrueVal,
+                                     Value *FalseVal,
+                                     InstCombiner::BuilderTy &Builder) {
+  BinaryOperator *BOp;
+  Constant *C1, *C2, *C3;
+  Value *X;
+  ICmpInst::Predicate Predicate;
+
+  if (!match(Cmp, m_ICmp(Predicate, m_Value(X), m_Constant(C1))))
+    return nullptr;
+
+  if (!((match(TrueVal, m_BinOp(BOp)) && match(FalseVal, m_Constant(C3))) ||
+        (match(FalseVal, m_BinOp(BOp)) && match(TrueVal, m_Constant(C3)))))
+    return nullptr;
+
+  if (!match(BOp, m_OneUse(m_c_BinOp(m_Specific(X), m_Constant(C2)))))
+    return nullptr;
+
+  // `select (Cmp X C1) (sub C2 X) (sub C1 C2)` cannot be transformed
+  // into something like `sub (select (Cmp X C1) X C1) C2`
+  if (!BOp->isCommutative() && BOp->getOperand(0) != X)
+    return nullptr;
+
+  if (!(ICmpInst::isRelational(Predicate) &&
+        C3 == ConstantFoldBinaryOpOperands(BOp->getOpcode(), C1, C2,
+                                           BOp->getDataLayout())))
+    return nullptr;
+
+  Value *NewTrueVal = isa<Constant>(TrueVal) ? C1 : X;
+  Value *NewFalseVal = isa<Constant>(FalseVal) ? C1 : X;
+  Value *LHS, *RHS;
+  SelectPatternFlavor SPF =
+      matchDecomposedSelectPattern(Cmp, NewTrueVal, NewFalseVal, LHS, RHS)
+          .Flavor;
+
+  Intrinsic::ID IntrinsicID;
+  switch (SPF) {
+  case SelectPatternFlavor::SPF_UMIN:
+    IntrinsicID = Intrinsic::umin;
+    break;
+  case SelectPatternFlavor::SPF_UMAX:
+    IntrinsicID = Intrinsic::umax;
+    break;
+  case SelectPatternFlavor::SPF_SMIN:
+    IntrinsicID = Intrinsic::smin;
+    break;
+  case SelectPatternFlavor::SPF_SMAX:
+    IntrinsicID = Intrinsic::smax;
+    break;
+  default:
+    llvm_unreachable("Unexpected SPF");
+  }
+  Value *Intrinsic = Builder.CreateBinaryIntrinsic(IntrinsicID, LHS, RHS);
+  BinaryOperator *NewBO = cast<BinaryOperator>(
+      Builder.CreateBinOp(BOp->getOpcode(), Intrinsic, C2));
+  NewBO->copyIRFlags(BOp);
+  return NewBO;
+}
+
 /// Visit a SelectInst that has an ICmpInst as its first operand.
 Instruction *InstCombinerImpl::foldSelectInstWithICmp(SelectInst &SI,
                                                       ICmpInst *ICI) {
@@ -1987,6 +2050,9 @@ Instruction *InstCombinerImpl::foldSelectInstWithICmp(SelectInst &SI,
   if (Value *V = foldAbsDiff(ICI, TrueVal, FalseVal, Builder))
     return replaceInstUsesWith(SI, V);
 
+  if (Value *V = canonicalizeConstToBOp(ICI, TrueVal, FalseVal, Builder))
+    return replaceInstUsesWith(SI, V);
+
   return Changed ? &SI : nullptr;
 }
 
diff --git a/llvm/test/Transforms/InstCombine/canonicalize-const-to-bop.ll b/llvm/test/Transforms/InstCombine/canonicalize-const-to-bop.ll
@@ -4,9 +4,8 @@
 define i8 @add_and_sgt(i8 %x) {
 ; CHECK-LABEL: define i8 @add_and_sgt(
 ; CHECK-SAME: i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[ADD:%.*]] = add nsw i8 [[X]], 16
-; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i8 [[X]], 8
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[CMP]], i8 [[ADD]], i8 24
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.smax.i8(i8 [[X]], i8 8)
+; CHECK-NEXT:    [[S:%.*]] = add nuw nsw i8 [[TMP1]], 16
 ; CHECK-NEXT:    ret i8 [[S]]
 ;
   %add = add nsw i8 %x, 16
@@ -18,9 +17,8 @@ define i8 @add_and_sgt(i8 %x) {
 define i8 @add_sgt_nuw(i8 %x) {
 ; CHECK-LABEL: define i8 @add_sgt_nuw(
 ; CHECK-SAME: i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[ADD:%.*]] = add nuw i8 [[X]], 16
-; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i8 [[X]], 8
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[CMP]], i8 [[ADD]], i8 24
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.smax.i8(i8 [[X]], i8 8)
+; CHECK-NEXT:    [[S:%.*]] = add nuw i8 [[TMP1]], 16
 ; CHECK-NEXT:    ret i8 [[S]]
 ;
   %add = add nuw i8 %x, 16
@@ -32,9 +30,8 @@ define i8 @add_sgt_nuw(i8 %x) {
 define i8 @sub_and_ugt(i8 %x) {
 ; CHECK-LABEL: define i8 @sub_and_ugt(
 ; CHECK-SAME: i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[SUB:%.*]] = add nsw i8 [[X]], -50
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i8 [[X]], 100
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[CMP]], i8 50, i8 [[SUB]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.umin.i8(i8 [[X]], i8 100)
+; CHECK-NEXT:    [[S:%.*]] = add nsw i8 [[TMP1]], -50
 ; CHECK-NEXT:    ret i8 [[S]]
 ;
   %sub = sub nsw i8 %x, 50
@@ -46,9 +43,8 @@ define i8 @sub_and_ugt(i8 %x) {
 define i8 @sub_ugt_nuw_nsw(i8 %x) {
 ; CHECK-LABEL: define i8 @sub_ugt_nuw_nsw(
 ; CHECK-SAME: i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[SUB:%.*]] = add nsw i8 [[X]], -50
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i8 [[X]], 100
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[CMP]], i8 50, i8 [[SUB]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.umin.i8(i8 [[X]], i8 100)
+; CHECK-NEXT:    [[S:%.*]] = add nsw i8 [[TMP1]], -50
 ; CHECK-NEXT:    ret i8 [[S]]
 ;
   %sub = sub nuw nsw i8 %x, 50
@@ -60,9 +56,8 @@ define i8 @sub_ugt_nuw_nsw(i8 %x) {
 define i8 @mul_and_ult(i8 %x) {
 ; CHECK-LABEL: define i8 @mul_and_ult(
 ; CHECK-SAME: i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[ADD:%.*]] = mul nsw i8 [[X]], 10
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i8 [[X]], 10
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[CMP]], i8 100, i8 [[ADD]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.umin.i8(i8 [[X]], i8 10)
+; CHECK-NEXT:    [[S:%.*]] = mul nuw nsw i8 [[TMP1]], 10
 ; CHECK-NEXT:    ret i8 [[S]]
 ;
   %add = mul nsw i8 %x, 10
@@ -74,9 +69,8 @@ define i8 @mul_and_ult(i8 %x) {
 define i8 @mul_ult_noflags(i8 %x) {
 ; CHECK-LABEL: define i8 @mul_ult_noflags(
 ; CHECK-SAME: i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[ADD:%.*]] = mul i8 [[X]], 10
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i8 [[X]], 10
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[CMP]], i8 100, i8 [[ADD]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.umin.i8(i8 [[X]], i8 10)
+; CHECK-NEXT:    [[S:%.*]] = mul nuw i8 [[TMP1]], 10
 ; CHECK-NEXT:    ret i8 [[S]]
 ;
   %add = mul i8 %x, 10
@@ -88,9 +82,8 @@ define i8 @mul_ult_noflags(i8 %x) {
 define i8 @udiv_and_slt(i8 %x) {
 ; CHECK-LABEL: define i8 @udiv_and_slt(
 ; CHECK-SAME: i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[SUB:%.*]] = udiv i8 [[X]], 10
-; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i8 [[X]], 100
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[CMP]], i8 10, i8 [[SUB]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.smax.i8(i8 [[X]], i8 100)
+; CHECK-NEXT:    [[S:%.*]] = udiv i8 [[TMP1]], 10
 ; CHECK-NEXT:    ret i8 [[S]]
 ;
   %sub = udiv i8 %x, 10
@@ -102,9 +95,8 @@ define i8 @udiv_and_slt(i8 %x) {
 define i8 @udiv_slt_exact(i8 %x) {
 ; CHECK-LABEL: define i8 @udiv_slt_exact(
 ; CHECK-SAME: i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[SUB:%.*]] = udiv exact i8 [[X]], 10
-; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i8 [[X]], 100
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[CMP]], i8 10, i8 [[SUB]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.smax.i8(i8 [[X]], i8 100)
+; CHECK-NEXT:    [[S:%.*]] = udiv exact i8 [[TMP1]], 10
 ; CHECK-NEXT:    ret i8 [[S]]
 ;
   %sub = udiv exact i8 %x, 10
@@ -116,9 +108,8 @@ define i8 @udiv_slt_exact(i8 %x) {
 define i8 @canonicalize_icmp_operands(i8 %x) {
 ; CHECK-LABEL: define i8 @canonicalize_icmp_operands(
 ; CHECK-SAME: i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[ADD:%.*]] = add nsw i8 [[X]], 8
-; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i8 [[X]], 119
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[CMP]], i8 127, i8 [[ADD]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.smin.i8(i8 [[X]], i8 119)
+; CHECK-NEXT:    [[S:%.*]] = add nsw i8 [[TMP1]], 8
 ; CHECK-NEXT:    ret i8 [[S]]
 ;
   %add = add nsw i8 %x, 8
@@ -133,10 +124,10 @@ declare void @use_byte(i8)
 define i8 @multi_use_cond_and_sel(i8 %x) {
 ; CHECK-LABEL: define i8 @multi_use_cond_and_sel(
 ; CHECK-SAME: i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[ADD:%.*]] = add nsw i8 [[X]], 16
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i8 [[X]], 8
 ; CHECK-NEXT:    call void @use(i1 [[CMP]])
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[CMP]], i8 [[ADD]], i8 24
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.smax.i8(i8 [[X]], i8 8)
+; CHECK-NEXT:    [[S:%.*]] = add nuw nsw i8 [[TMP1]], 16
 ; CHECK-NEXT:    call void @use_byte(i8 [[S]])
 ; CHECK-NEXT:    ret i8 [[S]]
 ;
@@ -155,11 +146,9 @@ define void @rust_noop_loop() {
 ; CHECK:       [[BB2_I]]:
 ; CHECK-NEXT:    [[ITER_SROA_0_07:%.*]] = phi i32 [ 0, %[[START]] ], [ [[SPEC_SELECT5:%.*]], %[[BB2_I]] ]
 ; CHECK-NEXT:    [[_0_I3_I:%.*]] = icmp sgt i32 [[ITER_SROA_0_07]], 99
-; CHECK-NEXT:    [[TMP0:%.*]] = add nsw i32 [[ITER_SROA_0_07]], 1
-; CHECK-NEXT:    [[SPEC_SELECT5]] = select i1 [[_0_I3_I]], i32 100, i32 [[TMP0]]
-; CHECK-NEXT:    [[_0_I_NOT_I:%.*]] = icmp sgt i32 [[SPEC_SELECT5]], 100
-; CHECK-NEXT:    [[OR_COND:%.*]] = select i1 [[_0_I3_I]], i1 true, i1 [[_0_I_NOT_I]]
-; CHECK-NEXT:    br i1 [[OR_COND]], label %[[BASICBLOCK4:.*]], label %[[BB2_I]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call i32 @llvm.smin.i32(i32 [[ITER_SROA_0_07]], i32 99)
+; CHECK-NEXT:    [[SPEC_SELECT5]] = add nsw i32 [[TMP0]], 1
+; CHECK-NEXT:    br i1 [[_0_I3_I]], label %[[BASICBLOCK4:.*]], label %[[BB2_I]]
 ; CHECK:       [[BASICBLOCK4]]:
 ; CHECK-NEXT:    ret void
 ;
@@ -182,9 +171,8 @@ basicblock4:
 define <2 x i8> @add_non_splat_vector(<2 x i8> %x) {
 ; CHECK-LABEL: define <2 x i8> @add_non_splat_vector(
 ; CHECK-SAME: <2 x i8> [[X:%.*]]) {
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i8> [[X]], <i8 1, i8 0>
-; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt <2 x i8> [[X]], <i8 0, i8 1>
-; CHECK-NEXT:    [[S:%.*]] = select <2 x i1> [[CMP]], <2 x i8> [[ADD]], <2 x i8> splat (i8 1)
+; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x i8> @llvm.smax.v2i8(<2 x i8> [[X]], <2 x i8> <i8 0, i8 1>)
+; CHECK-NEXT:    [[S:%.*]] = add nuw <2 x i8> [[TMP1]], <i8 1, i8 0>
 ; CHECK-NEXT:    ret <2 x i8> [[S]]
 ;
   %add = add <2 x i8> %x, <i8 1, i8 0>
@@ -196,9 +184,8 @@ define <2 x i8> @add_non_splat_vector(<2 x i8> %x) {
 define <2 x i8> @or_splat_vector(<2 x i8> %x) {
 ; CHECK-LABEL: define <2 x i8> @or_splat_vector(
 ; CHECK-SAME: <2 x i8> [[X:%.*]]) {
-; CHECK-NEXT:    [[ADD:%.*]] = or <2 x i8> [[X]], splat (i8 1)
-; CHECK-NEXT:    [[CMP_INV:%.*]] = icmp slt <2 x i8> [[X]], splat (i8 1)
-; CHECK-NEXT:    [[S:%.*]] = select <2 x i1> [[CMP_INV]], <2 x i8> splat (i8 1), <2 x i8> [[ADD]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x i8> @llvm.smax.v2i8(<2 x i8> [[X]], <2 x i8> splat (i8 1))
+; CHECK-NEXT:    [[S:%.*]] = or <2 x i8> [[TMP1]], splat (i8 1)
 ; CHECK-NEXT:    ret <2 x i8> [[S]]
 ;
   %add = or <2 x i8> %x, <i8 1, i8 1>
diff --git a/llvm/test/Transforms/InstCombine/saturating-add-sub.ll b/llvm/test/Transforms/InstCombine/saturating-add-sub.ll
@@ -1793,10 +1793,9 @@ define i32 @not_uadd_sat(i32 %x, i32 %y) {
 
 define i32 @not_uadd_sat2(i32 %x, i32 %y) {
 ; CHECK-LABEL: @not_uadd_sat2(
-; CHECK-NEXT:    [[A:%.*]] = add i32 [[X:%.*]], -2
-; CHECK-NEXT:    [[C:%.*]] = icmp ugt i32 [[X]], 1
-; CHECK-NEXT:    [[R:%.*]] = select i1 [[C]], i32 [[A]], i32 -1
-; CHECK-NEXT:    ret i32 [[R]]
+; CHECK-NEXT:    [[X:%.*]] = call i32 @llvm.umax.i32(i32 [[X1:%.*]], i32 1)
+; CHECK-NEXT:    [[A:%.*]] = add i32 [[X]], -2
+; CHECK-NEXT:    ret i32 [[A]]
 ;
   %a = add i32 %x, -2
   %c = icmp ugt i32 %x, 1
diff --git a/llvm/test/Transforms/InstCombine/select.ll b/llvm/test/Transforms/InstCombine/select.ll
@@ -2989,10 +2989,9 @@ define i8 @select_replacement_loop3(i32 noundef %x) {
 
 define i16 @select_replacement_loop4(i16 noundef %p_12) {
 ; CHECK-LABEL: @select_replacement_loop4(
-; CHECK-NEXT:    [[AND1:%.*]] = and i16 [[P_12:%.*]], 1
-; CHECK-NEXT:    [[CMP21:%.*]] = icmp ult i16 [[P_12]], 2
-; CHECK-NEXT:    [[AND3:%.*]] = select i1 [[CMP21]], i16 [[AND1]], i16 0
-; CHECK-NEXT:    ret i16 [[AND3]]
+; CHECK-NEXT:    [[P_12:%.*]] = call i16 @llvm.umin.i16(i16 [[P_13:%.*]], i16 2)
+; CHECK-NEXT:    [[AND1:%.*]] = and i16 [[P_12]], 1
+; CHECK-NEXT:    ret i16 [[AND1]]
 ;
   %cmp1 = icmp ult i16 %p_12, 2
   %and1 = and i16 %p_12, 1
