[InstCombine] Fold select of clamped shifts

If we are feeding a shift into a select conditioned by an inbounds check for the shift amount, then we can strip any mask/clamp limit that has been put on the shift amount

Fold (select (icmp_ugt A, BW-1), (shift X, (and A, C)), Y) --> (select (icmp_ugt A, BW-1), (shift X, A), T)
Fold (select (icmp_ugt A, BW-1), Y, (shift X, (umin A, C))) --> (select (icmp_ugt A, BW-1), Y, (shift X, A))

Alive2: https://alive2.llvm.org/ce/z/xC6FwD

Fixes #109888

Author: RKSimon

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

@@ -2761,6 +2761,67 @@ static Instruction *foldSelectWithSRem(SelectInst &SI, InstCombinerImpl &IC,
   return nullptr;
 }
 
+static Instruction *foldSelectWithClampedShift(SelectInst &SI,
+                                               InstCombinerImpl &IC,
+                                               IRBuilderBase &Builder) {
+  Value *CondVal = SI.getCondition();
+  Value *TrueVal = SI.getTrueValue();
+  Value *FalseVal = SI.getFalseValue();
+  Type *SelType = SI.getType();
+  uint64_t BW = SelType->getScalarSizeInBits();
+
+  auto MatchClampedShift = [&](Value *V, Value *Amt) -> BinaryOperator * {
+    Value *X, *Limit;
+
+    // Fold (select (icmp_ugt A, BW-1), TrueVal, (shift X, (umin A, C)))
+    //  --> (select (icmp_ugt A, BW-1), TrueVal, (shift X, A))
+    // Fold (select (icmp_ult A, BW), (shift X, (umin A, C)), FalseVal)
+    //  --> (select (icmp_ult A, BW), (shift X, A), FalseVal)
+    // iff C >= BW-1
+    if (match(V, m_OneUse(m_Shift(m_Value(X),
+                                  m_UMin(m_Specific(Amt), m_Value(Limit)))))) {
+      KnownBits KnownLimit = IC.computeKnownBits(Limit, 0, &SI);
+      if (KnownLimit.getMinValue().uge(BW - 1))
+        return cast<BinaryOperator>(V);
+    }
+
+    // Fold (select (icmp_ugt A, BW-1), (shift X, (and A, C)), FalseVal)
+    //  --> (select (icmp_ugt A, BW-1), (shift X, A), FalseVal)
+    // Fold (select (icmp_ult A, BW), (shift X, (and A, C)), FalseVal)
+    //  --> (select (icmp_ult A, BW), (shift X, A), FalseVal)
+    // iff Pow2 element width and C masks all amt bits.
+    if (isPowerOf2_64(BW) &&
+        match(V, m_OneUse(m_Shift(m_Value(X),
+                                  m_And(m_Specific(Amt), m_Value(Limit)))))) {
+      KnownBits KnownLimit = IC.computeKnownBits(Limit, 0, &SI);
+      if (KnownLimit.countMinTrailingOnes() >= Log2_64(BW))
+        return cast<BinaryOperator>(V);
+    }
+
+    return nullptr;
+  };
+
+  Value *Amt;
+  if (match(CondVal, m_SpecificICmp(ICmpInst::ICMP_UGT, m_Value(Amt),
+                                    m_SpecificInt(BW - 1)))) {
+    if (BinaryOperator *ShiftI = MatchClampedShift(FalseVal, Amt))
+      return SelectInst::Create(
+          CondVal, TrueVal,
+          Builder.CreateBinOp(ShiftI->getOpcode(), ShiftI->getOperand(0), Amt));
+  }
+
+  if (match(CondVal, m_SpecificICmp(ICmpInst::ICMP_ULT, m_Value(Amt),
+                                    m_SpecificInt(BW)))) {
+    if (BinaryOperator *ShiftI = MatchClampedShift(TrueVal, Amt))
+      return SelectInst::Create(
+          CondVal,
+          Builder.CreateBinOp(ShiftI->getOpcode(), ShiftI->getOperand(0), Amt),
+          FalseVal);
+  }
+
+  return nullptr;
+}
+
 static Value *foldSelectWithFrozenICmp(SelectInst &Sel, InstCombiner::BuilderTy &Builder) {
   FreezeInst *FI = dyn_cast<FreezeInst>(Sel.getCondition());
   if (!FI)
@@ -3871,6 +3932,9 @@ Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
   if (Instruction *I = foldSelectExtConst(SI))
     return I;
 
+  if (Instruction *I = foldSelectWithClampedShift(SI, *this, Builder))
+    return I;
+
   if (Instruction *I = foldSelectWithSRem(SI, *this, Builder))
     return I;
 

llvm/test/Transforms/InstCombine/select-shift-clamp.ll

@@ -13,8 +13,7 @@ declare void @use_i32(i32)
 define i32 @select_ult_shl_clamp_and_i32(i32 %a0, i32 %a1, i32 %a2) {
 ; CHECK-LABEL: @select_ult_shl_clamp_and_i32(
 ; CHECK-NEXT:    [[C:%.*]] = icmp ult i32 [[A1:%.*]], 32
-; CHECK-NEXT:    [[M:%.*]] = and i32 [[A1]], 31
-; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[A0:%.*]], [[M]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[A0:%.*]], [[A1]]
 ; CHECK-NEXT:    [[R:%.*]] = select i1 [[C]], i32 [[TMP1]], i32 [[A2:%.*]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
@@ -28,8 +27,7 @@ define i32 @select_ult_shl_clamp_and_i32(i32 %a0, i32 %a1, i32 %a2) {
 define i32 @select_ule_ashr_clamp_and_i32(i32 %a0, i32 %a1, i32 %a2) {
 ; CHECK-LABEL: @select_ule_ashr_clamp_and_i32(
 ; CHECK-NEXT:    [[C:%.*]] = icmp ult i32 [[A1:%.*]], 32
-; CHECK-NEXT:    [[M:%.*]] = and i32 [[A1]], 127
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[A0:%.*]], [[M]]
+; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[A0:%.*]], [[A1]]
 ; CHECK-NEXT:    [[R:%.*]] = select i1 [[C]], i32 [[TMP1]], i32 [[A2:%.*]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
@@ -43,8 +41,7 @@ define i32 @select_ule_ashr_clamp_and_i32(i32 %a0, i32 %a1, i32 %a2) {
 define i32 @select_ugt_lshr_clamp_and_i32(i32 %a0, i32 %a1, i32 %a2) {
 ; CHECK-LABEL: @select_ugt_lshr_clamp_and_i32(
 ; CHECK-NEXT:    [[C:%.*]] = icmp ugt i32 [[A1:%.*]], 31
-; CHECK-NEXT:    [[M:%.*]] = and i32 [[A1]], 31
-; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[A0:%.*]], [[M]]
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[A0:%.*]], [[A1]]
 ; CHECK-NEXT:    [[R:%.*]] = select i1 [[C]], i32 [[A2:%.*]], i32 [[TMP1]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
@@ -58,8 +55,7 @@ define i32 @select_ugt_lshr_clamp_and_i32(i32 %a0, i32 %a1, i32 %a2) {
 define i32 @select_uge_shl_clamp_and_i32(i32 %a0, i32 %a1, i32 %a2) {
 ; CHECK-LABEL: @select_uge_shl_clamp_and_i32(
 ; CHECK-NEXT:    [[C:%.*]] = icmp ugt i32 [[A1:%.*]], 31
-; CHECK-NEXT:    [[M:%.*]] = and i32 [[A1]], 63
-; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[A0:%.*]], [[M]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[A0:%.*]], [[A1]]
 ; CHECK-NEXT:    [[R:%.*]] = select i1 [[C]], i32 [[A2:%.*]], i32 [[TMP1]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
@@ -129,8 +125,7 @@ define i17 @select_uge_lshr_clamp_and_i17_nonpow2(i17 %a0, i17 %a1, i17 %a2) {
 define i32 @select_ult_shl_clamp_umin_i32(i32 %a0, i32 %a1, i32 %a2) {
 ; CHECK-LABEL: @select_ult_shl_clamp_umin_i32(
 ; CHECK-NEXT:    [[C:%.*]] = icmp ult i32 [[A1:%.*]], 32
-; CHECK-NEXT:    [[M:%.*]] = call i32 @llvm.umin.i32(i32 [[A1]], i32 31)
-; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[A0:%.*]], [[M]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[A0:%.*]], [[A1]]
 ; CHECK-NEXT:    [[R:%.*]] = select i1 [[C]], i32 [[TMP1]], i32 [[A2:%.*]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
@@ -144,8 +139,7 @@ define i32 @select_ult_shl_clamp_umin_i32(i32 %a0, i32 %a1, i32 %a2) {
 define i17 @select_ule_ashr_clamp_umin_i17(i17 %a0, i17 %a1, i17 %a2) {
 ; CHECK-LABEL: @select_ule_ashr_clamp_umin_i17(
 ; CHECK-NEXT:    [[C:%.*]] = icmp ult i17 [[A1:%.*]], 17
-; CHECK-NEXT:    [[M:%.*]] = call i17 @llvm.umin.i17(i17 [[A1]], i17 17)
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr i17 [[A0:%.*]], [[M]]
+; CHECK-NEXT:    [[TMP1:%.*]] = ashr i17 [[A0:%.*]], [[A1]]
 ; CHECK-NEXT:    [[R:%.*]] = select i1 [[C]], i17 [[TMP1]], i17 [[A2:%.*]]
 ; CHECK-NEXT:    ret i17 [[R]]
 ;
@@ -159,8 +153,7 @@ define i17 @select_ule_ashr_clamp_umin_i17(i17 %a0, i17 %a1, i17 %a2) {
 define i32 @select_ugt_shl_clamp_umin_i32(i32 %a0, i32 %a1, i32 %a2) {
 ; CHECK-LABEL: @select_ugt_shl_clamp_umin_i32(
 ; CHECK-NEXT:    [[C:%.*]] = icmp ugt i32 [[A1:%.*]], 31
-; CHECK-NEXT:    [[M:%.*]] = call i32 @llvm.umin.i32(i32 [[A1]], i32 128)
-; CHECK-NEXT:    [[S:%.*]] = shl i32 [[A0:%.*]], [[M]]
+; CHECK-NEXT:    [[S:%.*]] = shl i32 [[A0:%.*]], [[A1]]
 ; CHECK-NEXT:    [[R:%.*]] = select i1 [[C]], i32 [[A2:%.*]], i32 [[S]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
@@ -174,8 +167,7 @@ define i32 @select_ugt_shl_clamp_umin_i32(i32 %a0, i32 %a1, i32 %a2) {
 define <2 x i32> @select_uge_lshr_clamp_umin_v2i32(<2 x i32> %a0, <2 x i32> %a1, <2 x i32> %a2) {
 ; CHECK-LABEL: @select_uge_lshr_clamp_umin_v2i32(
 ; CHECK-NEXT:    [[C:%.*]] = icmp ugt <2 x i32> [[A1:%.*]], <i32 31, i32 31>
-; CHECK-NEXT:    [[M:%.*]] = call <2 x i32> @llvm.umin.v2i32(<2 x i32> [[A1]], <2 x i32> <i32 63, i32 31>)
-; CHECK-NEXT:    [[S:%.*]] = lshr <2 x i32> [[A0:%.*]], [[M]]
+; CHECK-NEXT:    [[S:%.*]] = lshr <2 x i32> [[A0:%.*]], [[A1]]
 ; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[C]], <2 x i32> [[A2:%.*]], <2 x i32> [[S]]
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
@@ -223,8 +215,7 @@ define i17 @select_uge_lshr_clamp_umin_i17_badlimit(i17 %a0, i17 %a1, i17 %a2) {
 define range(i64 0, -9223372036854775807) <4 x i64> @PR109888(<4 x i64> %0) {
 ; CHECK-LABEL: @PR109888(
 ; CHECK-NEXT:    [[C:%.*]] = icmp ult <4 x i64> [[TMP0:%.*]], <i64 64, i64 64, i64 64, i64 64>
-; CHECK-NEXT:    [[M:%.*]] = and <4 x i64> [[TMP0]], <i64 63, i64 63, i64 63, i64 63>
-; CHECK-NEXT:    [[TMP2:%.*]] = shl nuw <4 x i64> <i64 1, i64 1, i64 1, i64 1>, [[M]]
+; CHECK-NEXT:    [[TMP2:%.*]] = shl nuw <4 x i64> <i64 1, i64 1, i64 1, i64 1>, [[TMP0]]
 ; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[C]], <4 x i64> [[TMP2]], <4 x i64> zeroinitializer
 ; CHECK-NEXT:    ret <4 x i64> [[R]]
 ;