[InstCombine] Implement folds of icmp of UCMP/SCMP call and a constant (#96118)

This patch handles various cases where an operation of the kind `icmp
(ucmp/scmp x, y), constant` folds to `icmp x, y`. Another patch with
cases where this operation folds to a constant (i.e. dumb cases like
`icmp eq (cmp x, y), 4` should be published in a couple of days.

I wasn't sure what negative tests should be added here, if any are
necessary at all. I'd love to hear your suggestions.

Proofs (ucmp): https://alive2.llvm.org/ce/z/qQ7ihz
Proofs (scmp): https://alive2.llvm.org/ce/z/cipKEn

---------

Co-authored-by: Nikita Popov <github@npopov.com>

3 files changed, 363 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index 522c31f..1ad6375 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -3933,6 +3933,52 @@ foldICmpUSubSatOrUAddSatWithConstant(ICmpInst::Predicate Pred,
       ConstantInt::get(Op1->getType(), EquivInt));
 }
 
+static Instruction *
+foldICmpOfCmpIntrinsicWithConstant(ICmpInst::Predicate Pred, IntrinsicInst *I,
+                                   const APInt &C,
+                                   InstCombiner::BuilderTy &Builder) {
+  std::optional<ICmpInst::Predicate> NewPredicate = std::nullopt;
+  switch (Pred) {
+  case ICmpInst::ICMP_EQ:
+  case ICmpInst::ICMP_NE:
+    if (C.isZero())
+      NewPredicate = Pred;
+    else if (C.isOne())
+      NewPredicate =
+          Pred == ICmpInst::ICMP_EQ ? ICmpInst::ICMP_UGT : ICmpInst::ICMP_ULE;
+    else if (C.isAllOnes())
+      NewPredicate =
+          Pred == ICmpInst::ICMP_EQ ? ICmpInst::ICMP_ULT : ICmpInst::ICMP_UGE;
+    break;
+
+  case ICmpInst::ICMP_SGT:
+    if (C.isAllOnes())
+      NewPredicate = ICmpInst::ICMP_UGE;
+    else if (C.isZero())
+      NewPredicate = ICmpInst::ICMP_UGT;
+    break;
+
+  case ICmpInst::ICMP_SLT:
+    if (C.isZero())
+      NewPredicate = ICmpInst::ICMP_ULT;
+    else if (C.isOne())
+      NewPredicate = ICmpInst::ICMP_ULE;
+    break;
+
+  default:
+    break;
+  }
+
+  if (!NewPredicate)
+    return nullptr;
+
+  if (I->getIntrinsicID() == Intrinsic::scmp)
+    NewPredicate = ICmpInst::getSignedPredicate(*NewPredicate);
+  Value *LHS = I->getOperand(0);
+  Value *RHS = I->getOperand(1);
+  return new ICmpInst(*NewPredicate, LHS, RHS);
+}
+
 /// Fold an icmp with LLVM intrinsic and constant operand: icmp Pred II, C.
 Instruction *InstCombinerImpl::foldICmpIntrinsicWithConstant(ICmpInst &Cmp,
                                                              IntrinsicInst *II,
@@ -3954,6 +4000,11 @@ Instruction *InstCombinerImpl::foldICmpIntrinsicWithConstant(ICmpInst &Cmp,
     if (Instruction *R = foldCtpopPow2Test(Cmp, II, C, Builder, Q))
       return R;
   } break;
+  case Intrinsic::scmp:
+  case Intrinsic::ucmp:
+    if (auto *Folded = foldICmpOfCmpIntrinsicWithConstant(Pred, II, C, Builder))
+      return Folded;
+    break;
   }
 
   if (Cmp.isEquality())
diff --git a/llvm/test/Transforms/InstCombine/scmp.ll b/llvm/test/Transforms/InstCombine/scmp.ll
new file mode 100644
index 0000000..4f903a7
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/scmp.ll
@@ -0,0 +1,156 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define i1 @scmp_eq_0(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @scmp_eq_0(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.scmp(i32 %x, i32 %y)
+  %2 = icmp eq i8 %1, 0
+  ret i1 %2
+}
+
+define i1 @scmp_ne_0(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @scmp_ne_0(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.scmp(i32 %x, i32 %y)
+  %2 = icmp ne i8 %1, 0
+  ret i1 %2
+}
+
+define i1 @scmp_eq_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @scmp_eq_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp sgt i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.scmp(i32 %x, i32 %y)
+  %2 = icmp eq i8 %1, 1
+  ret i1 %2
+}
+
+define i1 @scmp_ne_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @scmp_ne_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp sle i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.scmp(i32 %x, i32 %y)
+  %2 = icmp ne i8 %1, 1
+  ret i1 %2
+}
+
+define i1 @scmp_eq_negative_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @scmp_eq_negative_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp slt i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.scmp(i32 %x, i32 %y)
+  %2 = icmp eq i8 %1, -1
+  ret i1 %2
+}
+
+define i1 @scmp_ne_negative_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @scmp_ne_negative_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp sge i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.scmp(i32 %x, i32 %y)
+  %2 = icmp ne i8 %1, -1
+  ret i1 %2
+}
+
+define i1 @scmp_sgt_0(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @scmp_sgt_0(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp sgt i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.scmp(i32 %x, i32 %y)
+  %2 = icmp sgt i8 %1, 0
+  ret i1 %2
+}
+
+define i1 @scmp_sgt_neg_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @scmp_sgt_neg_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp sge i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.scmp(i32 %x, i32 %y)
+  %2 = icmp sgt i8 %1, -1
+  ret i1 %2
+}
+
+define i1 @scmp_sge_0(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @scmp_sge_0(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp sge i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.scmp(i32 %x, i32 %y)
+  %2 = icmp sge i8 %1, 0
+  ret i1 %2
+}
+
+define i1 @scmp_sge_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @scmp_sge_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp sgt i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.scmp(i32 %x, i32 %y)
+  %2 = icmp sge i8 %1, 1
+  ret i1 %2
+}
+
+define i1 @scmp_slt_0(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @scmp_slt_0(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp slt i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.scmp(i32 %x, i32 %y)
+  %2 = icmp slt i8 %1, 0
+  ret i1 %2
+}
+
+define i1 @scmp_slt_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @scmp_slt_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp sle i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.scmp(i32 %x, i32 %y)
+  %2 = icmp slt i8 %1, 1
+  ret i1 %2
+}
+
+define i1 @scmp_sle_0(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @scmp_sle_0(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp sle i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.scmp(i32 %x, i32 %y)
+  %2 = icmp sle i8 %1, 0
+  ret i1 %2
+}
+
+define i1 @scmp_sle_neg_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @scmp_sle_neg_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp slt i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.scmp(i32 %x, i32 %y)
+  %2 = icmp sle i8 %1, -1
+  ret i1 %2
+}
diff --git a/llvm/test/Transforms/InstCombine/ucmp.ll b/llvm/test/Transforms/InstCombine/ucmp.ll
new file mode 100644
index 0000000..9ab6756
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/ucmp.ll
@@ -0,0 +1,156 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define i1 @ucmp_eq_0(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @ucmp_eq_0(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.ucmp(i32 %x, i32 %y)
+  %2 = icmp eq i8 %1, 0
+  ret i1 %2
+}
+
+define i1 @ucmp_ne_0(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @ucmp_ne_0(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.ucmp(i32 %x, i32 %y)
+  %2 = icmp ne i8 %1, 0
+  ret i1 %2
+}
+
+define i1 @ucmp_eq_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @ucmp_eq_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ugt i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.ucmp(i32 %x, i32 %y)
+  %2 = icmp eq i8 %1, 1
+  ret i1 %2
+}
+
+define i1 @ucmp_ne_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @ucmp_ne_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ule i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.ucmp(i32 %x, i32 %y)
+  %2 = icmp ne i8 %1, 1
+  ret i1 %2
+}
+
+define i1 @ucmp_eq_negative_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @ucmp_eq_negative_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.ucmp(i32 %x, i32 %y)
+  %2 = icmp eq i8 %1, -1
+  ret i1 %2
+}
+
+define i1 @ucmp_ne_negative_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @ucmp_ne_negative_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp uge i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.ucmp(i32 %x, i32 %y)
+  %2 = icmp ne i8 %1, -1
+  ret i1 %2
+}
+
+define i1 @ucmp_sgt_0(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @ucmp_sgt_0(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ugt i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.ucmp(i32 %x, i32 %y)
+  %2 = icmp sgt i8 %1, 0
+  ret i1 %2
+}
+
+define i1 @ucmp_sgt_neg_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @ucmp_sgt_neg_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp uge i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.ucmp(i32 %x, i32 %y)
+  %2 = icmp sgt i8 %1, -1
+  ret i1 %2
+}
+
+define i1 @ucmp_sge_0(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @ucmp_sge_0(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp uge i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.ucmp(i32 %x, i32 %y)
+  %2 = icmp sge i8 %1, 0
+  ret i1 %2
+}
+
+define i1 @ucmp_sge_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @ucmp_sge_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ugt i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.ucmp(i32 %x, i32 %y)
+  %2 = icmp sge i8 %1, 1
+  ret i1 %2
+}
+
+define i1 @ucmp_slt_0(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @ucmp_slt_0(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.ucmp(i32 %x, i32 %y)
+  %2 = icmp slt i8 %1, 0
+  ret i1 %2
+}
+
+define i1 @ucmp_slt_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @ucmp_slt_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ule i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.ucmp(i32 %x, i32 %y)
+  %2 = icmp slt i8 %1, 1
+  ret i1 %2
+}
+
+define i1 @ucmp_sle_0(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @ucmp_sle_0(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ule i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.ucmp(i32 %x, i32 %y)
+  %2 = icmp sle i8 %1, 0
+  ret i1 %2
+}
+
+define i1 @ucmp_sle_neg_1(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @ucmp_sle_neg_1(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = call i8 @llvm.ucmp(i32 %x, i32 %y)
+  %2 = icmp sle i8 %1, -1
+  ret i1 %2
+}