[InstCombine] Preserve NSW/NUW flags when folding const BOp with min/max (#143471)

When folding `X Pred C2 ? X BOp C1 : C2 BOp C1` to `min/max(X, C2) BOp
C1`, if NUW/NSW flags are present on `X BOp C1` and could be safely
applied to `C2 BOp C1`, then they may be added on the BOp after the fold
is complete. https://alive2.llvm.org/ce/z/n_3aNJ

Preserving these flags can allow subsequent transforms to re-order the
min/max and BOp, which in the case of NVPTX would allow for some
potential future transformations which would improve
instruction-selection.

Author: AlexMaclean

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

@@ -771,6 +771,8 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
                             Value *A, Value *B, Instruction &Outer,
                             SelectPatternFlavor SPF2, Value *C);
   Instruction *foldSelectInstWithICmp(SelectInst &SI, ICmpInst *ICI);
+  Value *foldSelectWithConstOpToBinOp(ICmpInst *Cmp, Value *TrueVal,
+                                      Value *FalseVal);
   Instruction *foldSelectValueEquivalence(SelectInst &SI, CmpInst &CI);
   bool replaceInInstruction(Value *V, Value *Old, Value *New,
                             unsigned Depth = 0);

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

@@ -1879,9 +1879,9 @@ static Instruction *foldSelectICmpEq(SelectInst &SI, ICmpInst *ICI,
 
 /// Fold `X Pred C1 ? X BOp C2 : C1 BOp C2` to `min/max(X, C1) BOp C2`.
 /// This allows for better canonicalization.
-static Value *foldSelectWithConstOpToBinOp(ICmpInst *Cmp, Value *TrueVal,
-                                           Value *FalseVal,
-                                           IRBuilderBase &Builder) {
+Value *InstCombinerImpl::foldSelectWithConstOpToBinOp(ICmpInst *Cmp,
+                                                      Value *TrueVal,
+                                                      Value *FalseVal) {
   Constant *C1, *C2, *C3;
   Value *X;
   CmpPredicate Predicate;
@@ -1945,11 +1945,29 @@ static Value *foldSelectWithConstOpToBinOp(ICmpInst *Cmp, Value *TrueVal,
     return nullptr;
   }
 
-  Intrinsic::ID IntrinsicID = getMinMaxIntrinsic(SPF);
-  Value *Intrinsic = Builder.CreateBinaryIntrinsic(IntrinsicID, X, RHS);
-  return IsIntrinsic ? Builder.CreateBinaryIntrinsic(Opcode, Intrinsic, C2)
-                     : Builder.CreateBinOp(Instruction::BinaryOps(Opcode),
-                                           Intrinsic, C2);
+  Intrinsic::ID MinMaxID = getMinMaxIntrinsic(SPF);
+  Value *MinMax = Builder.CreateBinaryIntrinsic(MinMaxID, X, RHS);
+  if (IsIntrinsic)
+    return Builder.CreateBinaryIntrinsic(Opcode, MinMax, C2);
+
+  const auto BinOpc = Instruction::BinaryOps(Opcode);
+  Value *BinOp = Builder.CreateBinOp(BinOpc, MinMax, C2);
+
+  // If we can attach no-wrap flags to the new instruction, do so if the
+  // old instruction had them and C1 BinOp C2 does not overflow.
+  if (Instruction *BinOpInst = dyn_cast<Instruction>(BinOp)) {
+    if (BinOpc == Instruction::Add || BinOpc == Instruction::Sub ||
+        BinOpc == Instruction::Mul) {
+      Instruction *OldBinOp = cast<BinaryOperator>(TrueVal);
+      if (OldBinOp->hasNoSignedWrap() &&
+          willNotOverflow(BinOpc, RHS, C2, *BinOpInst, /*IsSigned=*/true))
+        BinOpInst->setHasNoSignedWrap();
+      if (OldBinOp->hasNoUnsignedWrap() &&
+          willNotOverflow(BinOpc, RHS, C2, *BinOpInst, /*IsSigned=*/false))
+        BinOpInst->setHasNoUnsignedWrap();
+    }
+  }
+  return BinOp;
 }
 
 /// Visit a SelectInst that has an ICmpInst as its first operand.
@@ -2027,7 +2045,7 @@ Instruction *InstCombinerImpl::foldSelectInstWithICmp(SelectInst &SI,
   if (Value *V = foldAbsDiff(ICI, TrueVal, FalseVal, Builder))
     return replaceInstUsesWith(SI, V);
 
-  if (Value *V = foldSelectWithConstOpToBinOp(ICI, TrueVal, FalseVal, Builder))
+  if (Value *V = foldSelectWithConstOpToBinOp(ICI, TrueVal, FalseVal))
     return replaceInstUsesWith(SI, V);
 
   return Changed ? &SI : nullptr;

llvm/test/Transforms/InstCombine/canonicalize-const-to-bop.ll

@@ -5,7 +5,7 @@ define i8 @add_and_sgt(i8 %x) {
 ; CHECK-LABEL: define i8 @add_and_sgt(
 ; CHECK-SAME: i8 [[X:%.*]]) {
 ; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.smax.i8(i8 [[X]], i8 8)
-; CHECK-NEXT:    [[S:%.*]] = add nuw i8 [[TMP1]], 16
+; CHECK-NEXT:    [[S:%.*]] = add nuw nsw i8 [[TMP1]], 16
 ; CHECK-NEXT:    ret i8 [[S]]
 ;
   %add = add nsw i8 %x, 16
@@ -155,7 +155,7 @@ define i8 @multi_use_cond_and_sel(i8 %x) {
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i8 [[X]], 8
 ; CHECK-NEXT:    call void @use(i1 [[CMP]])
 ; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.smax.i8(i8 [[X]], i8 8)
-; CHECK-NEXT:    [[S:%.*]] = add nuw i8 [[TMP1]], 16
+; CHECK-NEXT:    [[S:%.*]] = add nuw nsw i8 [[TMP1]], 16
 ; CHECK-NEXT:    call void @use_byte(i8 [[S]])
 ; CHECK-NEXT:    ret i8 [[S]]
 ;
@@ -450,3 +450,82 @@ define i8 @umax_sgt(i8 %x) {
   %s = select i1 %cmp, i8 100, i8 %umax
   ret i8 %s
 }
+
+define i8 @add_sgt_nuw_nsw_safe(i8 %x) {
+; CHECK-LABEL: define i8 @add_sgt_nuw_nsw_safe(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.smin.i8(i8 [[X]], i8 100)
+; CHECK-NEXT:    [[S:%.*]] = add nuw nsw i8 [[TMP1]], 1
+; CHECK-NEXT:    ret i8 [[S]]
+;
+  %add = add nuw nsw i8 %x, 1
+  %cmp = icmp sgt i8 %x, 100
+  %s = select i1 %cmp, i8 101, i8 %add
+  ret i8 %s
+}
+
+define i8 @add_sgt_nuw_only(i8 %x) {
+; CHECK-LABEL: define i8 @add_sgt_nuw_only(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.smin.i8(i8 [[X]], i8 100)
+; CHECK-NEXT:    [[S:%.*]] = add nuw i8 [[TMP1]], 50
+; CHECK-NEXT:    ret i8 [[S]]
+;
+  %add = add nuw nsw i8 %x, 50
+  %cmp = icmp sgt i8 %x, 100
+  %s = select i1 %cmp, i8 150, i8 %add
+  ret i8 %s
+}
+
+define i8 @add_sgt_nsw_only(i8 %x) {
+; CHECK-LABEL: define i8 @add_sgt_nsw_only(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.smin.i8(i8 [[X]], i8 100)
+; CHECK-NEXT:    [[S:%.*]] = add nsw i8 [[TMP1]], -99
+; CHECK-NEXT:    ret i8 [[S]]
+;
+  %add = add nuw nsw i8 %x, -99
+  %cmp = icmp sgt i8 %x, 100
+  %s = select i1 %cmp, i8 1, i8 %add
+  ret i8 %s
+}
+
+
+define i8 @mul_ult_nuw_nsw_safe(i8 %x) {
+; CHECK-LABEL: define i8 @mul_ult_nuw_nsw_safe(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.umax.i8(i8 [[X]], i8 10)
+; CHECK-NEXT:    [[S:%.*]] = mul nuw nsw i8 [[TMP1]], 3
+; CHECK-NEXT:    ret i8 [[S]]
+;
+  %mul = mul nuw nsw i8 %x, 3
+  %cmp = icmp ult i8 %x, 10
+  %s = select i1 %cmp, i8 30, i8 %mul
+  ret i8 %s
+}
+
+define i8 @mul_ult_nuw_only(i8 %x) {
+; CHECK-LABEL: define i8 @mul_ult_nuw_only(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.umax.i8(i8 [[X]], i8 10)
+; CHECK-NEXT:    [[S:%.*]] = mul nuw i8 [[TMP1]], 25
+; CHECK-NEXT:    ret i8 [[S]]
+;
+  %mul = mul nuw nsw i8 %x, 25
+  %cmp = icmp ult i8 %x, 10
+  %s = select i1 %cmp, i8 250, i8 %mul
+  ret i8 %s
+}
+
+define i8 @mul_ult_nsw_only(i8 %x) {
+; CHECK-LABEL: define i8 @mul_ult_nsw_only(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.umax.i8(i8 [[X]], i8 40)
+; CHECK-NEXT:    [[S:%.*]] = mul nsw i8 [[TMP1]], -2
+; CHECK-NEXT:    ret i8 [[S]]
+;
+  %mul = mul nuw nsw i8 %x, -2
+  %cmp = icmp ult i8 %x, 40
+  %s = select i1 %cmp, i8 -80, i8 %mul
+  ret i8 %s
+}