[InstSimplify] Adding safety checks for 'shl', 'ashr' and 'lshr'

[opaparo]

Summary

As of today, InstSimplify does not transform patterns that include an 'shl', a 'ashr' or a 'lshr', as such a transformation may result in a different behavior in the case that that operand loses bits. However, in many cases these operators can be proven to not violate their restrictions and not lose any bits, and so we are missing safe transformation opportunities.
This patch introduces safety checks for the OverflowingBinaryOperators 'shl', and for the PossiblyExactOperators 'ashr' and 'lshr', enabling new kinds of transformations.

[opaparo]

Ping.

[spatel]

Moving the functions from instcombine to valuetracking is NFC, right? If so, can you do that now, so this patch is minimized? The similar functions should be in one place already.

[opaparo]

In D46380#1094593, @spatel wrote:

Moving the functions from instcombine to valuetracking is NFC, right? If so, can you do that now, so this patch is minimized? The similar functions should be in one place already.

Done. Separated the NFC to a different review and rebased on top of it.

[spatel]

I haven't looked at the details of everything that's going on in these folds, but a couple of higher-level points:

1. This is an instsimplify patch - the review title should be changed; there should be minimal tests for each transform under tests/Transforms/InstSimplify.
2. This is using ValueTracking which can be expensive in compile-time - do you have any stats for how often this fires, perf improvements, compile-time regressions? (adding Eli for thoughts about where we would draw that line because I really don't know)

[opaparo]

Removing safety checks for 'add', 'sub' and 'mul', since they require ValueTracking which proved to be too expensive in compile time.

[opaparo]

In D46380#1096108, @spatel wrote:

1. This is an instsimplify patch - the review title should be changed; there should be minimal tests for each transform under tests/Transforms/InstSimplify.

The title was changed. However, those changes do not affect InstSimplify directly, as they are not used by this pass. The only functional change is in InstCombine, which uses these parts of InstSimplify as a library.

In D46380#1096108, @spatel wrote:

2. This is using ValueTracking which can be expensive in compile-time - do you have any stats for how often this fires, perf improvements, compile-time regressions? (adding Eli for thoughts about where we would draw that line because I really don't know)

After a close look I discovered some compile-time regressions, as you predicted. I removed the safety checks for 'add', 'sub' and 'mul', which require ValueTracking (so now there are only safety checks for 'shl', 'ashr' and 'lshr'). That elimintated those regressions.

[spatel]

In D46380#1098045, @opaparo wrote:

In D46380#1096108, @spatel wrote:

1. This is an instsimplify patch - the review title should be changed; there should be minimal tests for each transform under tests/Transforms/InstSimplify.

The title was changed. However, those changes do not affect InstSimplify directly, as they are not used by this pass. The only functional change is in InstCombine, which uses these parts of InstSimplify as a library.

I don't understand this statement. The proposal affects SimplifyWithOpReplaced() which is only called from within instsimplify (simplifySelectWithICmpCond), so every test should be visible using only -instsimplify. The first test should be reduced to something like this:

define i64 @sel_false_val_is_a_masked_shl_of_true_val1(i32 %x) {
  %x15 = and i32 %x, 15
  %sh = shl nuw nsw i32 %x15, 2
  %z = zext i32 %sh to i64
  %cmp = icmp eq i32 %x15, 0
  %r = select i1 %cmp, i64 0, i64 %z
  ret i64 %r
}

But this example doesn't need nsw/nuw, so what is this test trying to demonstrate?
https://rise4fun.com/Alive/9zX

Also, this is still using ValueTracking, so I'm not comfortable continuing this review until we have more data about the cost and benefits.

[craig.topper]

lib/Analysis/InstructionSimplify.cpp

3497

dyn_cast_or_null allows OBO to be null, but it was dereferenced on the line above. If PEO can't be null, then this should just be dyn_cast.

3525

dyn_cast_or_null allows PEO to be null, but it was dereferenced on the line above. If PEO can't be null, then this should just be dyn_cast.

[opaparo]

In D46380#1103696, @spatel wrote:

In D46380#1098045, @opaparo wrote:

In D46380#1096108, @spatel wrote:

1. This is an instsimplify patch - the review title should be changed; there should be minimal tests for each transform under tests/Transforms/InstSimplify.

The title was changed. However, those changes do not affect InstSimplify directly, as they are not used by this pass. The only functional change is in InstCombine, which uses these parts of InstSimplify as a library.

I don't understand this statement. The proposal affects SimplifyWithOpReplaced() which is only called from within instsimplify (simplifySelectWithICmpCond), so every test should be visible using only -instsimplify. The first test should be reduced to something like this:

define i64 @sel_false_val_is_a_masked_shl_of_true_val1(i32 %x) {
  %x15 = and i32 %x, 15
  %sh = shl nuw nsw i32 %x15, 2
  %z = zext i32 %sh to i64
  %cmp = icmp eq i32 %x15, 0
  %r = select i1 %cmp, i64 0, i64 %z
  ret i64 %r
}

You're right. I've originally missed the direct effect on InstSimplify, and I will change the tests location and structure if this change is accepted, but please also note that lib/Transforms/InstCombine/InstCombineSelect.cpp calls SimplifySelectInst which calls simplifySelectWithICmpCond which in turn calls SimplifyWithOpReplaced. Hence, InstCombine also benefits from this functional change.

But this example doesn't need nsw/nuw, so what is this test trying to demonstrate?
https://rise4fun.com/Alive/9zX

Exactly that. In the current status, InstCombine and InstSimplify will not perform this transformation despite it's correctness. Furthermore, if you remove the nsw/nuw flags from the 'shl' in this example, InstCombine will first identify that the 'shl' really has no signed and unsigned wraps and will add the flags before the 'select' had the chance to optimize, which will result in the same un-transformed code. This is exactly one of the missed transformation opportunities I'm trying to cease.

Also, this is still using ValueTracking, so I'm not comfortable continuing this review until we have more data about the cost and benefits.

You're right that ValueTracking is still used, but running this change on a large set of benchmark did not yield any noticeable compile-time regressions. In addition, the two functions from ValueTracking that I'm using (ComputeNumSignBits and MaskedValueIsZero) are already in use in InstructionSimplify.cpp.

[opaparo]

Changing 'dyn_cast_or_null' to 'dyn_cast' in two places where the operator cannot be null.

[opaparo]

Rebasing the patch and adjusting it to the changes made in rL333689.
Removed special code for CastInst in InstructionSimplify::SimplifyWithOpReplaced as rL333689's changes handle that (in a different way).

[spatel]

I see the motivation better now, but I'm still worried about using ValueTracking.

Can you collect some stats as was done in D47891? I'm curious to know:

1. How often do we call ValueTracking?
2. How often does this transform occur?

Reasonable benchmarks are the same as in the other patch: test-suite or compiling clang/llvm.

[lebedev.ri]

(Removing from my review queue for the time being.)

[lebedev.ri]

This review may be stuck/dead, consider abandoning if no longer relevant.
Removing myself as reviewer in attempt to clean dashboard.

lib/Analysis/InstructionSimplify.cpp

static Value *SimplifyFCmpInst(unsigned Predicate, Value *LHS, Value *RHS,
   return nullptr;
 }
 
 Value *llvm::SimplifyFCmpInst(unsigned Predicate, Value *LHS, Value *RHS,
                               FastMathFlags FMF, const SimplifyQuery &Q) {
   return ::SimplifyFCmpInst(Predicate, LHS, RHS, FMF, Q, RecursionLimit);
 }
 
+static bool
+IsOverflowingBinaryOperatorSafe(const OverflowingBinaryOperator *OBO,
+                                const SimplifyQuery &Q) {
+  bool NUW = OBO->hasNoUnsignedWrap(), NSW = OBO->hasNoSignedWrap();
+  if (!NUW && !NSW)
+    return true;
+  const Instruction *I = dyn_cast<Instruction>(OBO);
+  if (!I)
+    return false;
+  Value *LHS = OBO->getOperand(0), *RHS = OBO->getOperand(1);
+  switch (OBO->getOpcode()) {
+  default:
+    return false;
+  case Instruction::Shl:
+    const APInt *ShAmtAPInt;
+    if (!match(RHS, m_APInt(ShAmtAPInt)))
+      return false;
+
+    unsigned ShAmt = ShAmtAPInt->getZExtValue();
+    unsigned BitWidth = OBO->getType()->getScalarSizeInBits();
+
+    if (NUW && !MaskedValueIsZero(LHS, APInt::getHighBitsSet(BitWidth, ShAmt),
+                                  Q.DL, 0, Q.AC, I, Q.DT))
+      return false;
+    if (NSW && ComputeNumSignBits(LHS, Q.DL, 0, Q.AC, I, Q.DT) <= ShAmt)
+      return false;
+    return true;
+  }
+}
+
+static bool IsPossiblyExactOperatorSafe(const PossiblyExactOperator *PEO,
+                                        const SimplifyQuery &Q) {
+  if (!PEO->isExact())
+    return true;
+  const Instruction *I = dyn_cast<Instruction>(PEO);
+  if (!I)
+    return false;
+  switch (PEO->getOpcode()) {
+  default:
+    return false;
+  case Instruction::LShr:
+  case Instruction::AShr:
+    Value *Op0 = I->getOperand(0), *Op1 = I->getOperand(1);
+
+    const APInt *ShAmtAPInt;
+    if (!match(Op1, m_APInt(ShAmtAPInt)))
+      return false;
+
+    unsigned ShAmt = ShAmtAPInt->getZExtValue();
+    unsigned BitWidth = I->getType()->getScalarSizeInBits();
+
+    return MaskedValueIsZero(Op0, APInt::getLowBitsSet(BitWidth, ShAmt), Q.DL,
+                             0, Q.AC, I, Q.DT);
+  }
+}
+
 /// See if V simplifies when its operand Op is replaced with RepOp.
 static const Value *SimplifyWithOpReplaced(Value *V, Value *Op, Value *RepOp,
                                            const SimplifyQuery &Q,
                                            unsigned MaxRecurse) {
   // Trivial replacement.
   if (V == Op)
     return RepOp;
 
   // We cannot replace a constant, and shouldn't even try.
   if (isa<Constant>(Op))
     return nullptr;
 
   auto *I = dyn_cast<Instruction>(V);
   if (!I)
     return nullptr;
 
   // If this is a binary operator, try to simplify it with the replaced op.
   if (auto *B = dyn_cast<BinaryOperator>(I)) {
     // Consider:
     //   %cmp = icmp eq i32 %x, 2147483647
     //   %add = add nsw i32 %x, 1
     //   %sel = select i1 %cmp, i32 -2147483648, i32 %add
     //
     // We can't replace %sel with %add unless we strip away the flags.
-    if (isa<OverflowingBinaryOperator>(B))
-      if (B->hasNoSignedWrap() || B->hasNoUnsignedWrap())
+    if (OverflowingBinaryOperator *OBO = dyn_cast<OverflowingBinaryOperator>(B))
+      if (!IsOverflowingBinaryOperatorSafe(OBO, Q))
         return nullptr;
-    if (isa<PossiblyExactOperator>(B))
-      if (B->isExact())
+    if (PossiblyExactOperator *PEO = dyn_cast<PossiblyExactOperator>(B))
+      if (!IsPossiblyExactOperatorSafe(PEO, Q))
         return nullptr;
 
     if (MaxRecurse) {
       if (B->getOperand(0) == Op)
         return SimplifyBinOp(B->getOpcode(), RepOp, B->getOperand(1), Q,
                              MaxRecurse - 1);
       if (B->getOperand(1) == Op)
         return SimplifyBinOp(B->getOpcode(), B->getOperand(0), RepOp, Q,

test/Transforms/InstCombine/select-bitext-bitwise-ops.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -instcombine < %s | FileCheck %s
 
 define i64 @sel_false_val_is_a_masked_shl_of_true_val1(i32 %x, i64 %y) {
 ; CHECK-LABEL: @sel_false_val_is_a_masked_shl_of_true_val1(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 15
-; CHECK-NEXT:    [[TMP2:%.*]] = shl nuw nsw i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP1]], 0
-; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[TMP3]], i32 0, i32 [[TMP2]]
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[NARROW]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = ashr i64 [[Y:%.*]], [[TMP4]]
-; CHECK-NEXT:    ret i64 [[TMP5]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 %x, 2
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], 60
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = ashr i64 %y, [[TMP3]]
+; CHECK-NEXT:    ret i64 [[TMP4]]
 ;
   %1 = and i32 %x, 15
   %2 = shl nuw nsw i32 %1, 2
   %3 = zext i32 %2 to i64
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
;
   %4 = icmp eq i32 %2, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
 }
 
 define i64 @sel_false_val_is_a_masked_lshr_of_true_val1(i32 %x, i64 %y) {
 ; CHECK-LABEL: @sel_false_val_is_a_masked_lshr_of_true_val1(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 60
-; CHECK-NEXT:    [[TMP2:%.*]] = lshr exact i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP1]], 0
-; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[TMP3]], i32 0, i32 [[TMP2]]
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[NARROW]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = ashr i64 [[Y:%.*]], [[TMP4]]
-; CHECK-NEXT:    ret i64 [[TMP5]]
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 %x, 2
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], 15
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = ashr i64 %y, [[TMP3]]
+; CHECK-NEXT:    ret i64 [[TMP4]]
 ;
   %1 = and i32 %x, 60
   %2 = lshr i32 %1, 2
   %3 = zext i32 %2 to i64
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
;
   %4 = icmp eq i32 %2, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
 }
 
 define i64 @sel_false_val_is_a_masked_ashr_of_true_val1(i32 %x, i64 %y) {
 ; CHECK-LABEL: @sel_false_val_is_a_masked_ashr_of_true_val1(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], -2147483588
-; CHECK-NEXT:    [[TMP2:%.*]] = ashr exact i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP1]], 0
-; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[TMP3]], i32 0, i32 [[TMP2]]
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[NARROW]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = ashr i64 [[Y:%.*]], [[TMP4]]
-; CHECK-NEXT:    ret i64 [[TMP5]]
+; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 %x, 2
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], -536870897
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = ashr i64 %y, [[TMP3]]
+; CHECK-NEXT:    ret i64 [[TMP4]]
 ;
   %1 = and i32 %x, -2147483588
   %2 = ashr i32 %1, 2
   %3 = zext i32 %2 to i64
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6

test/Transforms/InstCombine/select-obo-peo-ops.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -instcombine < %s | FileCheck %s
 
 define i64 @test_shl_nuw_nsw__all_are_safe(i32 %x, i64 %y) {
 ; CHECK-LABEL: @test_shl_nuw_nsw__all_are_safe(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 15
-; CHECK-NEXT:    [[TMP2:%.*]] = shl nuw nsw i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP1]], 0
-; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[TMP3]], i32 0, i32 [[TMP2]]
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[NARROW]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = ashr i64 [[Y:%.*]], [[TMP4]]
-; CHECK-NEXT:    ret i64 [[TMP5]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[X:%.*]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], 60
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = ashr i64 [[Y:%.*]], [[TMP3]]
+; CHECK-NEXT:    ret i64 [[TMP4]]
 ;
   %1 = and i32 %x, 15
   %2 = shl nuw nsw i32 %1, 2
   %3 = zext i32 %2 to i64
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
 }
 
 define i64 @test_shl_nuw__all_are_safe(i32 %x, i64 %y) {
 ; CHECK-LABEL: @test_shl_nuw__all_are_safe(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 15
-; CHECK-NEXT:    [[TMP2:%.*]] = shl nuw nsw i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP1]], 0
-; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[TMP3]], i32 0, i32 [[TMP2]]
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[NARROW]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = ashr i64 [[Y:%.*]], [[TMP4]]
-; CHECK-NEXT:    ret i64 [[TMP5]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[X:%.*]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], 60
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = ashr i64 [[Y:%.*]], [[TMP3]]
+; CHECK-NEXT:    ret i64 [[TMP4]]
 ;
   %1 = and i32 %x, 15
   %2 = shl nuw i32 %1, 2
   %3 = zext i32 %2 to i64
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
 }
 
 define i64 @test_shl_nsw__all_are_safe(i32 %x, i64 %y) {
 ; CHECK-LABEL: @test_shl_nsw__all_are_safe(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 15
-; CHECK-NEXT:    [[TMP2:%.*]] = shl nuw nsw i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP1]], 0
-; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[TMP3]], i32 0, i32 [[TMP2]]
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[NARROW]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = ashr i64 [[Y:%.*]], [[TMP4]]
-; CHECK-NEXT:    ret i64 [[TMP5]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[X:%.*]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], 60
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = ashr i64 [[Y:%.*]], [[TMP3]]
+; CHECK-NEXT:    ret i64 [[TMP4]]
 ;
   %1 = and i32 %x, 15
   %2 = shl nsw i32 %1, 2
   %3 = zext i32 %2 to i64
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
 }
 
 define i64 @test_shl__all_are_safe(i32 %x, i64 %y) {
 ; CHECK-LABEL: @test_shl__all_are_safe(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 15
-; CHECK-NEXT:    [[TMP2:%.*]] = shl nuw nsw i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP1]], 0
-; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[TMP3]], i32 0, i32 [[TMP2]]
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[NARROW]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = ashr i64 [[Y:%.*]], [[TMP4]]
-; CHECK-NEXT:    ret i64 [[TMP5]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[X:%.*]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], 60
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = ashr i64 [[Y:%.*]], [[TMP3]]
+; CHECK-NEXT:    ret i64 [[TMP4]]
 ;
   %1 = and i32 %x, 15
   %2 = shl i32 %1, 2
   %3 = zext i32 %2 to i64
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
;
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
 }
 
 define i64 @test_shl_nuw__nuw_is_safe(i32 %x, i64 %y) {
 ; CHECK-LABEL: @test_shl_nuw__nuw_is_safe(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 1073741822
-; CHECK-NEXT:    [[TMP2:%.*]] = shl nuw i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP1]], 0
-; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[TMP3]], i32 0, i32 [[TMP2]]
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[NARROW]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = ashr i64 [[Y:%.*]], [[TMP4]]
-; CHECK-NEXT:    ret i64 [[TMP5]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[X:%.*]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], -8
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = ashr i64 [[Y:%.*]], [[TMP3]]
+; CHECK-NEXT:    ret i64 [[TMP4]]
 ;
   %1 = and i32 %x, 1073741822
   %2 = shl nuw i32 %1, 2
   %3 = zext i32 %2 to i64
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
;
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
 }
 
 define i64 @test_shl__nuw_is_safe(i32 %x, i64 %y) {
 ; CHECK-LABEL: @test_shl__nuw_is_safe(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 1073741822
-; CHECK-NEXT:    [[TMP2:%.*]] = shl nuw i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP1]], 0
-; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[TMP3]], i32 0, i32 [[TMP2]]
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[NARROW]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = ashr i64 [[Y:%.*]], [[TMP4]]
-; CHECK-NEXT:    ret i64 [[TMP5]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[X:%.*]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], -8
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = ashr i64 [[Y:%.*]], [[TMP3]]
+; CHECK-NEXT:    ret i64 [[TMP4]]
 ;
   %1 = and i32 %x, 1073741822
   %2 = shl i32 %1, 2
   %3 = zext i32 %2 to i64
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
;
   %5 = mul i32 %4, %1
   %6 = mul i32 %5, %3
   ret i32 %6
 }
 
 define i32 @test_shl_nsw__nsw_is_safe(i32 %x) {
 ; CHECK-LABEL: @test_shl_nsw__nsw_is_safe(
 ; CHECK-NEXT:    [[TMP1:%.*]] = or i32 [[X:%.*]], -83886080
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i32 [[TMP1]], -83886079
-; CHECK-NEXT:    [[TMP3:%.*]] = shl nsw i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP4:%.*]] = select i1 [[TMP2]], i32 -335544316, i32 [[TMP3]]
-; CHECK-NEXT:    [[TMP5:%.*]] = mul i32 [[TMP4]], [[TMP1]]
-; CHECK-NEXT:    [[TMP6:%.*]] = mul i32 [[TMP5]], [[TMP3]]
-; CHECK-NEXT:    ret i32 [[TMP6]]
+; CHECK-NEXT:    [[TMP2:%.*]] = shl nsw i32 [[TMP1]], 2
+; CHECK-NEXT:    [[TMP3:%.*]] = mul i32 [[TMP2]], [[TMP1]]
+; CHECK-NEXT:    [[TMP4:%.*]] = mul i32 [[TMP3]], [[TMP2]]
+; CHECK-NEXT:    ret i32 [[TMP4]]
 ;
   %1 = or i32 %x, -83886080
   %2 = icmp eq i32 %1, -83886079
   %3 = shl nsw i32 %1, 2
   %4 = select i1 %2, i32 -335544316, i32 %3
   %5 = mul i32 %4, %1
   %6 = mul i32 %5, %3
   ret i32 %6
 }
 
 define i32 @test_shl__nsw_is_safe(i32 %x) {
 ; CHECK-LABEL: @test_shl__nsw_is_safe(
 ; CHECK-NEXT:    [[TMP1:%.*]] = or i32 [[X:%.*]], -83886080
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i32 [[TMP1]], -83886079
-; CHECK-NEXT:    [[TMP3:%.*]] = shl nsw i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP4:%.*]] = select i1 [[TMP2]], i32 -335544316, i32 [[TMP3]]
-; CHECK-NEXT:    [[TMP5:%.*]] = mul i32 [[TMP4]], [[TMP1]]
-; CHECK-NEXT:    [[TMP6:%.*]] = mul i32 [[TMP5]], [[TMP3]]
-; CHECK-NEXT:    ret i32 [[TMP6]]
+; CHECK-NEXT:    [[TMP2:%.*]] = shl nsw i32 [[TMP1]], 2
+; CHECK-NEXT:    [[TMP3:%.*]] = mul i32 [[TMP2]], [[TMP1]]
+; CHECK-NEXT:    [[TMP4:%.*]] = mul i32 [[TMP3]], [[TMP2]]
+; CHECK-NEXT:    ret i32 [[TMP4]]
 ;
   %1 = or i32 %x, -83886080
   %2 = icmp eq i32 %1, -83886079
   %3 = shl i32 %1, 2
   %4 = select i1 %2, i32 -335544316, i32 %3
   %5 = mul i32 %4, %1
   %6 = mul i32 %5, %3
   ret i32 %6
;
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
 }
 
 define i64 @test_lshr_exact__exact_is_safe(i32 %x, i64 %y) {
 ; CHECK-LABEL: @test_lshr_exact__exact_is_safe(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 60
-; CHECK-NEXT:    [[TMP2:%.*]] = lshr exact i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP1]], 0
-; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[TMP3]], i32 0, i32 [[TMP2]]
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[NARROW]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = ashr i64 [[Y:%.*]], [[TMP4]]
-; CHECK-NEXT:    ret i64 [[TMP5]]
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[X:%.*]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], 15
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = ashr i64 [[Y:%.*]], [[TMP3]]
+; CHECK-NEXT:    ret i64 [[TMP4]]
 ;
   %1 = and i32 %x, 60
   %2 = lshr exact i32 %1, 2
   %3 = zext i32 %2 to i64
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
 }
 
 define i64 @test_lshr__exact_is_safe(i32 %x, i64 %y) {
 ; CHECK-LABEL: @test_lshr__exact_is_safe(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 60
-; CHECK-NEXT:    [[TMP2:%.*]] = lshr exact i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP1]], 0
-; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[TMP3]], i32 0, i32 [[TMP2]]
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[NARROW]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = ashr i64 [[Y:%.*]], [[TMP4]]
-; CHECK-NEXT:    ret i64 [[TMP5]]
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[X:%.*]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], 15
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = ashr i64 [[Y:%.*]], [[TMP3]]
+; CHECK-NEXT:    ret i64 [[TMP4]]
 ;
   %1 = and i32 %x, 60
   %2 = lshr i32 %1, 2
   %3 = zext i32 %2 to i64
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
;
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
 }
 
 define i64 @test_ashr_exact__exact_is_safe(i32 %x, i64 %y) {
 ; CHECK-LABEL: @test_ashr_exact__exact_is_safe(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], -2147483588
-; CHECK-NEXT:    [[TMP2:%.*]] = ashr exact i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP1]], 0
-; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[TMP3]], i32 0, i32 [[TMP2]]
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[NARROW]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = ashr i64 [[Y:%.*]], [[TMP4]]
-; CHECK-NEXT:    ret i64 [[TMP5]]
+; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[X:%.*]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], -536870897
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = ashr i64 [[Y:%.*]], [[TMP3]]
+; CHECK-NEXT:    ret i64 [[TMP4]]
 ;
   %1 = and i32 %x, -2147483588
   %2 = ashr exact i32 %1, 2
   %3 = zext i32 %2 to i64
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6
 }
 
 define i64 @test_ashr__exact_is_safe(i32 %x, i64 %y) {
 ; CHECK-LABEL: @test_ashr__exact_is_safe(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], -2147483588
-; CHECK-NEXT:    [[TMP2:%.*]] = ashr exact i32 [[TMP1]], 2
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP1]], 0
-; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[TMP3]], i32 0, i32 [[TMP2]]
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[NARROW]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = ashr i64 [[Y:%.*]], [[TMP4]]
-; CHECK-NEXT:    ret i64 [[TMP5]]
+; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[X:%.*]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], -536870897
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = ashr i64 [[Y:%.*]], [[TMP3]]
+; CHECK-NEXT:    ret i64 [[TMP4]]
 ;
   %1 = and i32 %x, -2147483588
   %2 = ashr i32 %1, 2
   %3 = zext i32 %2 to i64
   %4 = icmp eq i32 %1, 0
   %5 = ashr i64 %y, %3
   %6 = select i1 %4, i64 %y, i64 %5
   ret i64 %6