[InstCombine] Optimize usub.sat pattern

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

@@ -50,7 +50,6 @@
 using namespace llvm;
 using namespace PatternMatch;
 
-
 /// Replace a select operand based on an equality comparison with the identity
 /// constant of a binop.
 static Instruction *foldSelectBinOpIdentity(SelectInst &Sel,
@@ -1713,7 +1712,6 @@ tryToReuseConstantFromSelectInComparison(SelectInst &Sel, ICmpInst &Cmp,
   if (Pred == CmpInst::ICMP_ULT && match(X, m_Add(m_Value(), m_Constant())))
     return nullptr;
 
-
   Value *SelVal0, *SelVal1; // We do not care which one is from where.
   match(&Sel, m_Select(m_Value(), m_Value(SelVal0), m_Value(SelVal1)));
   // At least one of these values we are selecting between must be a constant
@@ -1993,6 +1991,137 @@ Value *InstCombinerImpl::foldSelectWithConstOpToBinOp(ICmpInst *Cmp,
   return BinOp;
 }
 
+/// Folds:
+///   %a_sub = call @llvm.usub.sat(x, IntConst1)
+///   %b_sub = call @llvm.usub.sat(y, IntConst2)
+///   %or = or %a_sub, %b_sub
+///   %cmp = icmp eq %or, 0
+///   %sel = select %cmp, 0, MostSignificantBit
+/// into:
+///   %a_sub' = usub.sat(x, IntConst1 - MostSignificantBit)
+///   %b_sub' = usub.sat(y, IntConst2 - MostSignificantBit)
+///   %or = or %a_sub', %b_sub'
+///   %and = and %or, MostSignificantBit
+/// Likewise, for vector arguments as well.
+static Instruction *foldICmpUSubSatWithAndForMostSignificantBitCmp(
+    SelectInst &SI, ICmpInst *ICI, InstCombiner::BuilderTy &Builder) {
+  auto *CI = dyn_cast<ICmpInst>(SI.getCondition());
+  if (!CI) {
+    return nullptr;
+  }
+
+  Value *CmpLHS = CI->getOperand(0);
+  Value *CmpRHS = CI->getOperand(1);
+  if (!match(CmpRHS, m_Zero())) {
+    return nullptr;
+  }
+  auto Pred = CI->getPredicate();
+  auto *TrueVal = SI.getTrueValue();
+  auto *FalseVal = SI.getFalseValue();
+
+  if (Pred != ICmpInst::ICMP_EQ && Pred != llvm::ICmpInst::ICMP_NE)
+    return nullptr;
+
+  // Match: icmp eq (or (usub.sat A, IntConst1), (usub.sat B, IntConst2)), 0
+  Value *A, *B;
+  ConstantInt *IntConst1, *IntConst2, *PossibleMSBInt;
+
+  if (match(CmpLHS, m_Or(m_Intrinsic<Intrinsic::usub_sat>(
+                             m_Value(A), m_ConstantInt(IntConst1)),
+                         m_Intrinsic<Intrinsic::usub_sat>(
+                             m_Value(B), m_ConstantInt(IntConst2)))) &&
+      (match(TrueVal, m_Zero()) &&
+           match(FalseVal, m_ConstantInt(PossibleMSBInt)) ||
+       match(TrueVal, m_ConstantInt(PossibleMSBInt)) &&
+           match(FalseVal, m_Zero()))) {
+    auto *Ty = A->getType();
+    unsigned BW = Ty->getIntegerBitWidth();
+    APInt MostSignificantBit = APInt::getOneBitSet(BW, BW - 1);
+
+    if (PossibleMSBInt->getValue() != MostSignificantBit)
+      return nullptr;
+    // Ensure IntConst1 and IntConst2 are >= MostSignificantBit
+    if (IntConst1->getValue().ult(MostSignificantBit) ||
+        IntConst2->getValue().ult(MostSignificantBit))
+      return nullptr;
+
+    // Rewrite:
+    Value *NewA = Builder.CreateBinaryIntrinsic(
+        Intrinsic::usub_sat, A,
+        ConstantInt::get(Ty, IntConst1->getValue() - MostSignificantBit + 1));
+    Value *NewB = Builder.CreateBinaryIntrinsic(
+        Intrinsic::usub_sat, B,
+        ConstantInt::get(Ty, IntConst2->getValue() - MostSignificantBit + 1));
+    Value *Or = Builder.CreateOr(NewA, NewB);
+    Value *And =
+        Builder.CreateAnd(Or, ConstantInt::get(Ty, MostSignificantBit));
+    return cast<Instruction>(And);
+  }
+  Constant *Const1, *Const2, *PossibleMSB;
+  if (match(CmpLHS, m_Or(m_Intrinsic<Intrinsic::usub_sat>(m_Value(A),
+                                                          m_Constant(Const1)),
+                         m_Intrinsic<Intrinsic::usub_sat>(
+                             m_Value(B), m_Constant(Const2)))) &&
+      (match(TrueVal, m_Zero()) && match(FalseVal, m_Constant(PossibleMSB))
+    || match(TrueVal, m_Constant(PossibleMSB) ) && match(FalseVal, m_Zero()))) {
+    auto *VecTy1 = dyn_cast<FixedVectorType>(Const1->getType());
+    auto *VecTy2 = dyn_cast<FixedVectorType>(Const2->getType());
+    auto *VecTyMSB = dyn_cast<FixedVectorType>(PossibleMSB->getType());
+    if (!VecTy1 || !VecTy2 || !VecTyMSB) {
+      return nullptr;
+    }
+
+    unsigned NumElements = VecTy1->getNumElements();
+
+    if (NumElements != VecTy2->getNumElements() ||
+        NumElements != VecTyMSB->getNumElements() || NumElements == 0) {
+      return nullptr;
+    }
+    auto *SplatMSB =
+        dyn_cast<ConstantInt>(PossibleMSB->getAggregateElement(0u));
+    unsigned BW = SplatMSB->getValue().getBitWidth();
+    APInt MostSignificantBit = APInt::getOneBitSet(BW, BW - 1);
+    if (!SplatMSB || SplatMSB->getValue() != MostSignificantBit) {
+      return nullptr;
+    }
+    for (unsigned int i = 1; i < NumElements; ++i) {
+      auto *Element =
+          dyn_cast<ConstantInt>(PossibleMSB->getAggregateElement(i));
+      if (!Element || Element->getValue() != SplatMSB->getValue()) {
+        return nullptr;
+      }
+    }
+    SmallVector<Constant *, 16> Arg1, Arg2;
+    for (unsigned int i = 0; i < NumElements; ++i) {
+      auto *E1 = dyn_cast<ConstantInt>(Const1->getAggregateElement(i));
+      auto *E2 = dyn_cast<ConstantInt>(Const2->getAggregateElement(i));
+      if (!E1 || !E2) {
+        return nullptr;
+      }
+      if (E1->getValue().ult(SplatMSB->getValue()) ||
+          E2->getValue().ult(SplatMSB->getValue())) {
+        return nullptr;
+      }
+      Arg1.emplace_back(
+          ConstantInt::get(A->getType()->getScalarType(),
+                           E1->getValue() - MostSignificantBit + 1));
+      Arg2.emplace_back(
+          ConstantInt::get(B->getType()->getScalarType(),
+                           E2->getValue() - MostSignificantBit + 1));
+    }
+    Constant *ConstVec1 = ConstantVector::get(Arg1);
+    Constant *ConstVec2 = ConstantVector::get(Arg2);
+    Value *NewA =
+        Builder.CreateBinaryIntrinsic(Intrinsic::usub_sat, A, ConstVec1);
+    Value *NewB =
+        Builder.CreateBinaryIntrinsic(Intrinsic::usub_sat, B, ConstVec2);
+    Value *Or = Builder.CreateOr(NewA, NewB);
+    Value *And = Builder.CreateAnd(Or, PossibleMSB);
+    return cast<Instruction>(And);
+  }
+  return nullptr;
+}
+
 /// Visit a SelectInst that has an ICmpInst as its first operand.
 Instruction *InstCombinerImpl::foldSelectInstWithICmp(SelectInst &SI,
                                                       ICmpInst *ICI) {
@@ -2009,6 +2138,9 @@ Instruction *InstCombinerImpl::foldSelectInstWithICmp(SelectInst &SI,
   if (Instruction *NewSel =
           tryToReuseConstantFromSelectInComparison(SI, *ICI, *this))
     return NewSel;
+  if (Instruction *Folded =
+          foldICmpUSubSatWithAndForMostSignificantBitCmp(SI, ICI, Builder))
+    return replaceInstUsesWith(SI, Folded);
 
   // NOTE: if we wanted to, this is where to detect integer MIN/MAX
   bool Changed = false;

llvm/test/Transforms/InstCombine/usub_sat_to_msb_mask.ll

@@ -0,0 +1,206 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+
+; RUN: opt -passes=instcombine -S < %s 2>&1 | FileCheck %s
+
+declare i8 @llvm.usub.sat.i8(i8, i8)
+declare i16 @llvm.usub.sat.i16(i16, i16)
+declare i32 @llvm.usub.sat.i32(i32, i32)
+declare i64 @llvm.usub.sat.i64(i64, i64)
+
+define i8 @test_i8(i8 %a, i8 %b) {
+; CHECK-LABEL: define i8 @test_i8(
+; CHECK-SAME: i8 [[A:%.*]], i8 [[B:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.usub.sat.i8(i8 [[A]], i8 96)
+; CHECK-NEXT:    [[TMP2:%.*]] = call i8 @llvm.usub.sat.i8(i8 [[B]], i8 112)
+; CHECK-NEXT:    [[TMP3:%.*]] = or i8 [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[RES:%.*]] = and i8 [[TMP3]], -128
+; CHECK-NEXT:    ret i8 [[RES]]
+;
+
+  %a_sub = call i8 @llvm.usub.sat.i8(i8 %a, i8 223)
+  %b_sub = call i8 @llvm.usub.sat.i8(i8 %b, i8 239)
+  %or = or i8 %a_sub, %b_sub
+  %cmp = icmp eq i8 %or, 0
+  %res = select i1 %cmp, i8 0, i8 128
+  ret i8 %res
+}
+
+define i8 @test_i8_ne(i8 %a, i8 %b) {
+; CHECK-LABEL: define i8 @test_i8_ne(
+; CHECK-SAME: i8 [[A:%.*]], i8 [[B:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.usub.sat.i8(i8 [[A]], i8 96)
+; CHECK-NEXT:    [[TMP2:%.*]] = call i8 @llvm.usub.sat.i8(i8 [[B]], i8 112)
+; CHECK-NEXT:    [[TMP3:%.*]] = or i8 [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[RES:%.*]] = and i8 [[TMP3]], -128
+; CHECK-NEXT:    ret i8 [[RES]]
+;
+
+  %a_sub = call i8 @llvm.usub.sat.i8(i8 %a, i8 223)
+  %b_sub = call i8 @llvm.usub.sat.i8(i8 %b, i8 239)
+  %or = or i8 %a_sub, %b_sub
+  %cmp = icmp ne i8 %or, 0
+  %res = select i1 %cmp, i8 128, i8 0
+  ret i8 %res
+}
+
+define i16 @test_i16(i16 %a, i16 %b) {
+; CHECK-LABEL: define i16 @test_i16(
+; CHECK-SAME: i16 [[A:%.*]], i16 [[B:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i16 @llvm.usub.sat.i16(i16 [[A]], i16 32642)
+; CHECK-NEXT:    [[TMP2:%.*]] = call i16 @llvm.usub.sat.i16(i16 [[B]], i16 32656)
+; CHECK-NEXT:    [[TMP3:%.*]] = or i16 [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[RES:%.*]] = and i16 [[TMP3]], -32768
+; CHECK-NEXT:    ret i16 [[RES]]
+;
+
+  %a_sub = call i16 @llvm.usub.sat.i16(i16 %a, i16 65409)
+  %b_sub = call i16 @llvm.usub.sat.i16(i16 %b, i16 65423)
+  %or = or i16 %a_sub, %b_sub
+  %cmp = icmp eq i16 %or, 0
+  %res = select i1 %cmp, i16 0, i16 32768
+  ret i16 %res
+}
+
+define i32 @test_i32(i32 %a, i32 %b) {
+; CHECK-LABEL: define i32 @test_i32(
+; CHECK-SAME: i32 [[A:%.*]], i32 [[B:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.usub.sat.i32(i32 [[A]], i32 224)
+; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @llvm.usub.sat.i32(i32 [[B]], i32 240)
+; CHECK-NEXT:    [[TMP3:%.*]] = or i32 [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[RES:%.*]] = and i32 [[TMP3]], -2147483648
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+
+  %a_sub = call i32 @llvm.usub.sat.i32(i32 %a, i32 2147483871)
+  %b_sub = call i32 @llvm.usub.sat.i32(i32 %b, i32 2147483887)
+  %or = or i32 %a_sub, %b_sub
+  %cmp = icmp eq i32 %or, 0
+  %res = select i1 %cmp, i32 0, i32 2147483648
+  ret i32 %res
+}
+
+define i64 @test_i64(i64 %a, i64 %b) {
+; CHECK-LABEL: define i64 @test_i64(
+; CHECK-SAME: i64 [[A:%.*]], i64 [[B:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i64 @llvm.usub.sat.i64(i64 [[A]], i64 224)
+; CHECK-NEXT:    [[TMP2:%.*]] = call i64 @llvm.usub.sat.i64(i64 [[B]], i64 240)
+; CHECK-NEXT:    [[TMP3:%.*]] = or i64 [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[RES:%.*]] = and i64 [[TMP3]], -9223372036854775808
+; CHECK-NEXT:    ret i64 [[RES]]
+;
+
+  %a_sub = call i64 @llvm.usub.sat.i64(i64 %a, i64 9223372036854776031)
+  %b_sub = call i64 @llvm.usub.sat.i64(i64 %b, i64 9223372036854776047)
+  %or = or i64 %a_sub, %b_sub
+  %cmp = icmp eq i64 %or, 0
+  %res = select i1 %cmp, i64 0, i64 9223372036854775808
+  ret i64 %res
+}
+
+define i32 @no_fold_due_to_small_K(i32 %a, i32 %b) {
+; CHECK-LABEL: define i32 @no_fold_due_to_small_K(
+; CHECK-SAME: i32 [[A:%.*]], i32 [[B:%.*]]) {
+; CHECK-NEXT:    [[A_SUB:%.*]] = call i32 @llvm.usub.sat.i32(i32 [[A]], i32 100)
+; CHECK-NEXT:    [[B_SUB:%.*]] = call i32 @llvm.usub.sat.i32(i32 [[B]], i32 239)
+; CHECK-NEXT:    [[OR:%.*]] = or i32 [[A_SUB]], [[B_SUB]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[OR]], 0
+; CHECK-NEXT:    [[RES:%.*]] = select i1 [[CMP]], i32 0, i32 -2147483648
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+
+  %a_sub = call i32 @llvm.usub.sat.i32(i32 %a, i32 100)
+  %b_sub = call i32 @llvm.usub.sat.i32(i32 %b, i32 239)
+  %or = or i32 %a_sub, %b_sub
+  %cmp = icmp eq i32 %or, 0
+  %res = select i1 %cmp, i32 0, i32 2147483648
+  ret i32 %res
+}
+
+define i32 @commuted_test_neg(i32 %a, i32 %b) {
+; CHECK-LABEL: define i32 @commuted_test_neg(
+; CHECK-SAME: i32 [[A:%.*]], i32 [[B:%.*]]) {
+; CHECK-NEXT:    [[B_SUB:%.*]] = call i32 @llvm.usub.sat.i32(i32 [[B]], i32 239)
+; CHECK-NEXT:    [[A_SUB:%.*]] = call i32 @llvm.usub.sat.i32(i32 [[A]], i32 223)
+; CHECK-NEXT:    [[OR:%.*]] = or i32 [[B_SUB]], [[A_SUB]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[OR]], 0
+; CHECK-NEXT:    [[RES:%.*]] = select i1 [[CMP]], i32 0, i32 -2147483648
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+
+  %b_sub = call i32 @llvm.usub.sat.i32(i32 %b, i32 239)
+  %a_sub = call i32 @llvm.usub.sat.i32(i32 %a, i32 223)
+  %or = or i32 %b_sub, %a_sub
+  %cmp = icmp eq i32 %or, 0
+  %res = select i1 %cmp, i32 0, i32 2147483648
+  ret i32 %res
+}
+define <4 x i32> @vector_test(<4 x i32> %a, <4 x i32> %b) {
+; CHECK-LABEL: define <4 x i32> @vector_test(
+; CHECK-SAME: <4 x i32> [[A:%.*]], <4 x i32> [[B:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x i32> @llvm.usub.sat.v4i32(<4 x i32> [[A]], <4 x i32> splat (i32 224))
+; CHECK-NEXT:    [[TMP2:%.*]] = call <4 x i32> @llvm.usub.sat.v4i32(<4 x i32> [[B]], <4 x i32> splat (i32 240))
+; CHECK-NEXT:    [[TMP3:%.*]] = or <4 x i32> [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[RES:%.*]] = and <4 x i32> [[TMP3]], splat (i32 -2147483648)
+; CHECK-NEXT:    ret <4 x i32> [[RES]]
+;
+
+
+  %a_sub = call <4 x i32> @llvm.usub.sat.v4i32(
+  <4 x i32> %a,
+  <4 x i32> <i32 2147483871, i32 2147483871, i32 2147483871, i32 2147483871>)
+  %b_sub = call <4 x i32> @llvm.usub.sat.v4i32(
+  <4 x i32> %b,
+  <4 x i32> <i32 2147483887, i32 2147483887, i32 2147483887, i32 2147483887>)
+  %or = or <4 x i32> %a_sub, %b_sub
+  %cmp = icmp eq <4 x i32> %or, zeroinitializer
+  %res = select <4 x i1> %cmp, <4 x i32> zeroinitializer,
+  <4 x i32> <i32 -2147483648, i32 -2147483648, i32 -2147483648, i32 -2147483648>
+  ret <4 x i32> %res
+}
+
+define <4 x i32> @vector_negative_test(<4 x i32> %a, <4 x i32> %b) {
+; CHECK-LABEL: define <4 x i32> @vector_negative_test(
+; CHECK-SAME: <4 x i32> [[A:%.*]], <4 x i32> [[B:%.*]]) {
+; CHECK-NEXT:    [[A_SUB:%.*]] = call <4 x i32> @llvm.usub.sat.v4i32(<4 x i32> [[A]], <4 x i32> <i32 -2147483425, i32 0, i32 -2147483425, i32 -2147483425>)
+; CHECK-NEXT:    [[B_SUB:%.*]] = call <4 x i32> @llvm.usub.sat.v4i32(<4 x i32> [[B]], <4 x i32> splat (i32 -2147483409))
+; CHECK-NEXT:    [[OR:%.*]] = or <4 x i32> [[A_SUB]], [[B_SUB]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <4 x i32> [[OR]], zeroinitializer
+; CHECK-NEXT:    [[RES:%.*]] = select <4 x i1> [[CMP]], <4 x i32> zeroinitializer, <4 x i32> splat (i32 -2147483648)
+; CHECK-NEXT:    ret <4 x i32> [[RES]]
+;
+  %a_sub = call <4 x i32> @llvm.usub.sat.v4i32(
+  <4 x i32> %a,
+  <4 x i32> <i32 2147483871, i32 0, i32 2147483871, i32 2147483871>)
+  %b_sub = call <4 x i32> @llvm.usub.sat.v4i32(
+  <4 x i32> %b,
+  <4 x i32> <i32 2147483887, i32 2147483887, i32 2147483887, i32 2147483887>)
+  %or = or <4 x i32> %a_sub, %b_sub
+  %cmp = icmp eq <4 x i32> %or, zeroinitializer
+  %res = select <4 x i1> %cmp, <4 x i32> zeroinitializer,
+  <4 x i32> <i32 -2147483648, i32 -2147483648, i32 -2147483648, i32 -2147483648>
+  ret <4 x i32> %res
+}
+
+define <4 x i32> @vector_ne_test(<4 x i32> %a, <4 x i32> %b) {
+; CHECK-LABEL: define <4 x i32> @vector_ne_test(
+; CHECK-SAME: <4 x i32> [[A:%.*]], <4 x i32> [[B:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x i32> @llvm.usub.sat.v4i32(<4 x i32> [[A]], <4 x i32> splat (i32 224))
+; CHECK-NEXT:    [[TMP2:%.*]] = call <4 x i32> @llvm.usub.sat.v4i32(<4 x i32> [[B]], <4 x i32> splat (i32 240))
+; CHECK-NEXT:    [[TMP3:%.*]] = or <4 x i32> [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[RES:%.*]] = and <4 x i32> [[TMP3]], splat (i32 -2147483648)
+; CHECK-NEXT:    ret <4 x i32> [[RES]]
+;
+
+
+  %a_sub = call <4 x i32> @llvm.usub.sat.v4i32(
+  <4 x i32> %a,
+  <4 x i32> <i32 2147483871, i32 2147483871, i32 2147483871, i32 2147483871>)
+  %b_sub = call <4 x i32> @llvm.usub.sat.v4i32(
+  <4 x i32> %b,
+  <4 x i32> <i32 2147483887, i32 2147483887, i32 2147483887, i32 2147483887>)
+  %or = or <4 x i32> %a_sub, %b_sub
+  %cmp = icmp eq <4 x i32> %or, zeroinitializer
+  %res = select <4 x i1> %cmp, <4 x i32> zeroinitializer,
+  <4 x i32> <i32 -2147483648, i32 -2147483648, i32 -2147483648, i32 -2147483648>
+  ret <4 x i32> %res
+}