[InstCombine] Fold icmp(X) ? f(X) : C

This diff extends foldSelectInstWithICmp to handle the case icmp(X) ? f(X) : C
when f(X) is guaranteed to be equal to C for all X in the exact range of the inverse predicate.
This addresses the issue #54089.

Differential revision: https://reviews.llvm.org/D123159

Test plan: make check-all

Author: alexander-shaposhnikov

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

@@ -22,6 +22,7 @@
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constant.h"
+#include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
@@ -1529,6 +1530,36 @@ static Instruction *foldSelectZeroOrOnes(ICmpInst *Cmp, Value *TVal,
   return nullptr;
 }
 
+static Value *foldSelectInstWithICmpConst(SelectInst &SI, ICmpInst *ICI) {
+  const APInt *CmpC;
+  Value *V;
+  CmpInst::Predicate Pred;
+  if (!match(ICI, m_ICmp(Pred, m_Value(V), m_APInt(CmpC))))
+    return nullptr;
+
+  BinaryOperator *BO;
+  const APInt *C;
+  CmpInst::Predicate CPred;
+  if (match(&SI, m_Select(m_Specific(ICI), m_APInt(C), m_BinOp(BO))))
+    CPred = ICI->getPredicate();
+  else if (match(&SI, m_Select(m_Specific(ICI), m_BinOp(BO), m_APInt(C))))
+    CPred = ICI->getInversePredicate();
+  else
+    return nullptr;
+
+  const APInt *BinOpC;
+  if (!match(BO, m_BinOp(m_Specific(V), m_APInt(BinOpC))))
+    return nullptr;
+
+  ConstantRange R = ConstantRange::makeExactICmpRegion(CPred, *CmpC)
+                        .binaryOp(BO->getOpcode(), *BinOpC);
+  if (R == *C) {
+    BO->dropPoisonGeneratingFlags();
+    return BO;
+  }
+  return nullptr;
+}
+
 /// Visit a SelectInst that has an ICmpInst as its first operand.
 Instruction *InstCombinerImpl::foldSelectInstWithICmp(SelectInst &SI,
                                                       ICmpInst *ICI) {
@@ -1538,6 +1569,9 @@ Instruction *InstCombinerImpl::foldSelectInstWithICmp(SelectInst &SI,
   if (Instruction *NewSPF = canonicalizeSPF(SI, *ICI, *this))
     return NewSPF;
 
+  if (Value *V = foldSelectInstWithICmpConst(SI, ICI))
+    return replaceInstUsesWith(SI, V);
+
   if (Value *V = canonicalizeClampLike(SI, *ICI, Builder))
     return replaceInstUsesWith(SI, V);
 

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

@@ -648,10 +648,8 @@ define i32 @select_or_icmp_bad(i32 %x, i32 %y, i32 %z) {
 
 define i32 @select_lshr_icmp_const(i32 %x) {
 ; CHECK-LABEL: @select_lshr_icmp_const(
-; CHECK-NEXT:    [[A:%.*]] = icmp ugt i32 %x, 31
 ; CHECK-NEXT:    [[B:%.*]] = lshr i32 %x, 5
-; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i32 [[B]], i32 0
-; CHECK-NEXT:    ret i32 [[C]]
+; CHECK-NEXT:    ret i32 [[B]]
 ;
   %A = icmp ugt i32 %x, 31
   %B = lshr i32 %x, 5
@@ -661,10 +659,8 @@ define i32 @select_lshr_icmp_const(i32 %x) {
 
 define i32 @select_lshr_icmp_const_reordered(i32 %x) {
 ; CHECK-LABEL: @select_lshr_icmp_const_reordered(
-; CHECK-NEXT:    [[A:%.*]] = icmp ult i32 %x, 32
 ; CHECK-NEXT:    [[B:%.*]] = lshr i32 %x, 5
-; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i32 0, i32 [[B]]
-; CHECK-NEXT:    ret i32 [[C]]
+; CHECK-NEXT:    ret i32 [[B]]
 ;
   %A = icmp ult i32 %x, 32
   %B = lshr i32 %x, 5
@@ -674,10 +670,8 @@ define i32 @select_lshr_icmp_const_reordered(i32 %x) {
 
 define i32 @select_exact_lshr_icmp_const(i32 %x) {
 ; CHECK-LABEL: @select_exact_lshr_icmp_const(
-; CHECK-NEXT:    [[A:%.*]] = icmp ugt i32 %x, 31
-; CHECK-NEXT:    [[B:%.*]] = lshr exact i32 %x, 5
-; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i32 [[B]], i32 0
-; CHECK-NEXT:    ret i32 [[C]]
+; CHECK-NEXT:    [[B:%.*]] = lshr i32 %x, 5
+; CHECK-NEXT:    ret i32 [[B]]
 ;
   %A = icmp ugt i32 %x, 31
   %B = lshr exact i32 %x, 5