[InstCombine] `select + mul` -> `select + shl` with power of twos.

[danlark]

Summary

Now select + mul transforms to select + shl with power of twos known in select.

Patch by Danila Kutenin. email:danilak at google.com, github:danlark1@ . I don't have commit rights.

[danlark]

Remove TODO

[danlark]

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

267–269

What happens with undef lanes?
(they should become 0, do not keep them as undef.
there was some helper for that, can't find it now "replace undef values with")

274–275

This isn't sound.
We want to not propagate NSW if *any* lane matches this predicate, not just if they all do..
So you want

if (I.hasNoSignedWrap() && !C1.isNotMinSignedValue())
  Shl->setHasNoSignedWrap();

[danlark]

Fix the issues and rebase

[danlark]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

267–269

They will remain undef as in division, why is 0 better?

274–275

Yeah, I agree. Maybe you meant !cast<Constant>(Op1)->isMinSignedValue()

[danlark]

Add simple min integer test

[danlark]

Add zext test

[danlark]

More robust zext tests

[danlark]

Even more robust tests

[danlark]

Rebase

[danlark]

Correct INT_MIN test

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

267–269

Keeping undef in this situation is (a very common) miscompilation.

define <2 x i8> @shift_if_different_lanes_undef_vector(<2 x i8> %px, i1 %cond) {
%0:
  %sel = select i1 %cond, <2 x i8> { 16, undef }, <2 x i8> { undef, 32 }
  %r = mul <2 x i8> %px, %sel
  ret <2 x i8> %r
}
=>
define <2 x i8> @shift_if_different_lanes_undef_vector(<2 x i8> %px, i1 %cond) {
%0:
  %sel = select i1 %cond, <2 x i8> { 4, undef }, <2 x i8> { undef, 5 }
  %r = shl <2 x i8> %px, %sel
  ret <2 x i8> %r
}
Transformation doesn't verify!
ERROR: Target is more poisonous than source

Example:
<2 x i8> %px = < undef, undef >
i1 %cond = undef

Source:
<2 x i8> %sel = < undef, undef >
<2 x i8> %r = < undef, undef >

Target:
<2 x i8> %sel = < #x04 (4), #x08 (8) >
<2 x i8> %r = < #x00 (0)        [based on undef value], poison >
Source value: < undef, undef >
Target value: < #x00 (0), poison >

define <2 x i8> @shift_if_different_lanes_undef_vector(<2 x i8> %px, i1 %cond) {
%0:
  %sel = select i1 %cond, <2 x i8> { 16, undef }, <2 x i8> { undef, 32 }
  %r = mul <2 x i8> %px, %sel
  ret <2 x i8> %r
}
=>
define <2 x i8> @shift_if_different_lanes_undef_vector(<2 x i8> %px, i1 %cond) {
%0:
  %sel = select i1 %cond, <2 x i8> { 4, 0 }, <2 x i8> { 0, 5 }
  %r = shl <2 x i8> %px, %sel
  ret <2 x i8> %r
}
Transformation seems to be correct!

274–275

I'm very sure i didn't - that has the same problem of only checking for splat INT_MIN while we want to check if any channel is INT_MIN.

[danlark]

undef should return zero in select+shl

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

274–275

I believe for now it is correct, C1 is a logarithm of Op1, if Op1 is INT_MIN, then it should not have nsw. So, if all are not int_min, it will be propagated, otherwise, it will not.

define i32 @shift_if_power2_nuw_nsw_min(i32 %x, i1 %cond) {
  %sel = select i1 %cond, i32 2, i32 -2147483648
  %r = mul nuw nsw i32 %sel, %x
  ret i32 %r
}

===>

define i32 @shift_if_power2_nuw_nsw_min(i32 %x, i1 %cond) {
  %sel = select i1 %cond, i32 1, i32 31
  %r =  shl nuw i32 i32 %x, %sel
  ret i32 %r
}

[danlark]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

267–269

Done.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

116

You really want to only use Constant::replaceUndefsWith() to fixup that single case,
not pessimize everything that uses this generic utility..

274–275

Can you point me to the (vector) test that shows it working as expected?

[danlark]

Update vector test

[danlark]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

274–275

Ok, I changed to isNot

define <2 x i32> @shift_if_power2_vector_nsw(<2 x i32> %px, i1 %cond) {
  %sel = select i1 %cond, <2 x i32> <i32 4, i32 -2147483648>, <2 x i32> <i32 1, i32 32>
  %r = mul nsw <2 x i32> %px, %sel
  ret <2 x i32> %r
}
=>
define <2 x i32> @shift_if_power2_vector_nsw(<2 x i32> %px, i1 %cond) {
  %sel = select i1 %cond, <2 x i32> <i32 2, i32 31>, <2 x i32> <i32 0, i32 5>
  %r = shl <2 x i32> %px, %sel
  ret <2 x i32> %r
}

[danlark]

Rebase and style

[danlark]

Now I think review is in a good shape

[lebedev.ri]

Thank you for working on this.

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

265

I think it would be better to change Op1 type to Constant, and make the caller do the cast.

268–269

Add a comment explaining that we really do need to sanitize undef to 0 here.

281

// X mul (zext (C1 << N)), where C1 is "1<<C2" --> X << (zext (add N, C2))

320–325

If i'm reading this correctly, these are miscompilations:
https://rise4fun.com/Alive/IKH

llvm/test/Transforms/InstCombine/mul.ll

669–682

This is a miscompile:

----------------------------------------
define i32 @shift_if_power2_double_select(i32 %x, i32 %y, i1 %cond1, i1 %cond2) {
%0:
  %shl.res = shl i32 8, %y
  %sel1 = select i1 %cond1, i32 %shl.res, i32 1024
  %sel2 = select i1 %cond2, i32 16, i32 %sel1
  %r = mul nuw i32 %x, %sel2
  ret i32 %r
}
=>
define i32 @shift_if_power2_double_select(i32 %X, i32 %Y, i1 %COND1, i1 %COND2) {
%0:
  %TMP1 = add i32 %Y, 3
  %DOTV = select i1 %COND1, i32 %TMP1, i32 10
  %R_V = select i1 %COND2, i32 4, i32 %DOTV
  %R = shl nuw i32 %X, %R_V
  ret i32 %R
}
Transformation doesn't verify!
ERROR: Target is more poisonous than source

Example:
i32 %x = #x0f360904 (255199492)
i32 %y = #x0000001d (29)
i1 %cond1 = #x1 (1)
i1 %cond2 = undef

Source:
i32 %shl.res = #x00000000 (0)
i32 %sel1 = #x00000000 (0)
i32 %sel2 = #x00000000 (0)      [based on undef value]
i32 %r = #x00000000 (0)

Target:
i32 %TMP1 = #x00000020 (32)
i32 %DOTV = #x00000020 (32)
i32 %R_V = #x00000020 (32)
i32 %R = poison
Source value: #x00000000 (0)
Target value: poison

Summary:
  0 correct transformations
  1 incorrect transformations
  0 errors

[danlark]

Fix overflow in shifting and further propagation

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

265

It would be harder to combine OperandFoldActions then in udiv and mul at the same time, though I changed to only one call

320–325

Yeah, thank you for the tool. if Shl NSW, this should not happen and this is correct I believe https://rise4fun.com/Alive/sog

llvm/test/Transforms/InstCombine/mul.ll

669–682

Same issue with nsw for shift, fixed

[danlark]

In D71568#1788655, @lebedev.ri wrote:

Thank you for working on this.

I want to say sorry if I am asking/doing stupid things, I am only getting used to things

[danlark]

Also add NUW

[danlark]

Don't add zero

[danlark]

Ping

[danlark]

Ping.

[danlark]

Ping

[danlark]

Another hopeless ping before NY

[xbolva00]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

298

Dtto

300

assert is better choice here?

324

emplace_back?

[xbolva00]

In D71568#1798937, @danlark wrote:

Another hopeless ping before NY

It is a holiday time so please be patient

[danlark]

Fix asserts

[danlark]

Remove nullptr assignment

[danlark]

Replace finally everything with asserts

[danlark]

Ping

[danlark]

@lebedev.ri Can you please take a look and possibly submit?

[lebedev.ri]

Sorry, i lost track of this patch.
I was hoping someone more familiar with that code would comment, too.

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

333–338

You can at least cut this in half by using m_ZExtOrSelf()

[danlark]

Applying m_ZExtOrSelf

[danlark]

@spatel You are doing a lot of stuff around this code, PTAL.

[spatel]

There's a lot going on here, and the patch doesn't apply cleanly to current source. Can you pre-commit any NFC changes/tests to make this patch smaller? For example, the udiv change/tests are not affected by the mul code?

[danlark]

Removing div tests, we don't need them

[spatel]

In D71568#1812355, @spatel wrote:

There's a lot going on here, and the patch doesn't apply cleanly to current source. Can you pre-commit any NFC changes/tests to make this patch smaller? For example, the udiv change/tests are not affected by the mul code?

Sorry - I didn't see earlier that this patch is part of a sequence.
I'm a bit skeptical about the need to extend the "udiv action" machine. Are the motivating cases really that complicated or could we get away with a simpler pattern match of mul(select...)?

[danlark]

Rebase and make tests committed before

[danlark]

In D71568#1812426, @spatel wrote:

In D71568#1812355, @spatel wrote:

There's a lot going on here, and the patch doesn't apply cleanly to current source. Can you pre-commit any NFC changes/tests to make this patch smaller? For example, the udiv change/tests are not affected by the mul code?

Sorry - I didn't see earlier that this patch is part of a sequence.
I'm a bit skeptical about the need to extend the "udiv action" machine. Are the motivating cases really that complicated or could we get away with a simpler pattern match of mul(select...)?

I am not sure but I just reused the code for both division and multiplication -- from my perspective the idea is the same.

Also, I rebased a bit and made the patch smaller as you asked.

[danlark]

In D71568#1812463, @danlark wrote:

In D71568#1812426, @spatel wrote:

In D71568#1812355, @spatel wrote:

There's a lot going on here, and the patch doesn't apply cleanly to current source. Can you pre-commit any NFC changes/tests to make this patch smaller? For example, the udiv change/tests are not affected by the mul code?

Sorry - I didn't see earlier that this patch is part of a sequence.
I'm a bit skeptical about the need to extend the "udiv action" machine. Are the motivating cases really that complicated or could we get away with a simpler pattern match of mul(select...)?

I am not sure but I just reused the code for both division and multiplication -- from my perspective the idea is the same.

Also, I rebased a bit and made the patch smaller as you asked.

So, I basically don't see any reason division machine is doing a lot of things but as abstractions and semantics are mostly the same, I decided to reuse it. And to generalize power 2 multiplication then.

That's true that most of patterns are just one mul(select) but I saw at least one where the depth level was two.

[spatel]

In D71568#1812498, @danlark wrote:

In D71568#1812463, @danlark wrote:

In D71568#1812426, @spatel wrote:

In D71568#1812355, @spatel wrote:

There's a lot going on here, and the patch doesn't apply cleanly to current source. Can you pre-commit any NFC changes/tests to make this patch smaller? For example, the udiv change/tests are not affected by the mul code?

Sorry - I didn't see earlier that this patch is part of a sequence.
I'm a bit skeptical about the need to extend the "udiv action" machine. Are the motivating cases really that complicated or could we get away with a simpler pattern match of mul(select...)?

I am not sure but I just reused the code for both division and multiplication -- from my perspective the idea is the same.

Also, I rebased a bit and made the patch smaller as you asked.

So, I basically don't see any reason division machine is doing a lot of things but as abstractions and semantics are mostly the same, I decided to reuse it. And to generalize power 2 multiplication then.

That's true that most of patterns are just one mul(select) but I saw at least one where the depth level was two.

Yes, I understand the similarity, I'm just not sold on the idea that we ever needed the div complexity in the first place. If @lebedev.ri is comfortable with all of the poison-related corner cases, then we can proceed. AFAIK, there are no active contributers on this code, so if you've looked it over - then you're the expert. :)

[danlark]

@lebedev.ri Friendly ping and see message from spatel@ above :)

[danlark]

@spatel @lebedev.ri Ping

[spatel]

I'm struggling to keep all of the poison potential straight. Can we make this patch only include foldMulPow2Cst() (remove foldMulShl())?

Double-check to make sure I've translated this correctly:

define i32 @shift_if_power2_double_select_zext(i32 %x, i32 %y, i1 %cond1, i1 %cond2) {
  %shl = shl nsw i32 2147483648, %y
  %sel1 = select i1 %cond1, i32 %shl, i32 1024
  %sel2 = select i1 %cond2, i32 16, i32 %sel1
  %r = mul nsw i32 %x, %sel2
  ret i32 %r
}

$ opt -instcombine mul.ll -S

define i32 @shift_if_power2_double_select_zext(i32 %x, i32 %y, i1 %cond1, i1 %cond2) {
  %1 = add i32 %y, 31
  %.v = select i1 %cond1, i32 %1, i32 10
  %r.v = select i1 %cond2, i32 4, i32 %.v
  %r = shl nsw i32 %x, %r.v
  ret i32 %r
}

Miscompile?
https://rise4fun.com/Alive/YcG

[lebedev.ri]

This review seems to be stuck/dead, consider abandoning if no longer relevant.

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

if (!Ty->isVectorTy())
     return nullptr;
 
   SmallVector<Constant *, 4> Elts;
   for (unsigned I = 0, E = Ty->getVectorNumElements(); I != E; ++I) {
     Constant *Elt = C->getAggregateElement(I);
     if (!Elt)
       return nullptr;
     if (isa<UndefValue>(Elt)) {
       Elts.push_back(UndefValue::get(Ty->getScalarType()));
       continue;
     }
     if (!match(Elt, m_APInt(IVal)) || !IVal->isPowerOf2())
       return nullptr;
     Elts.push_back(ConstantInt::get(Ty->getScalarType(), IVal->logBase2()));
   }
 
   return ConstantVector::get(Elts);
if (size_t LHSIdx =
       if (visitUDivOperand(Op0, SI->getOperand(2), Actions, Depth)) {
         Actions.push_back(OperandFoldAction(nullptr, Op1, LHSIdx - 1));
         return Actions.size();
       }
 
   return 0;
 }
 
+// X mul 2^C -> X << C
+static Instruction *foldMulPow2Cst(Value *Op0, Value *Op1,
+                                   const BinaryOperator &I, InstCombiner &IC) {
+  Constant *C1 = cast<Constant>(Op1);
+  Constant *LogC1 = getLogBase2(Op0->getType(), C1);
+  // An undef value becomes zero because max bit shift can produce poisoned
+  // value.
+  LogC1 = Constant::replaceUndefsWith(
+      LogC1, ConstantInt::get(LogC1->getType()->getScalarType(), 0));
+  assert(LogC1 && "Failed to constant fold mul -> logbase2");
+  BinaryOperator *Shl = BinaryOperator::CreateShl(Op0, LogC1);
+  if (I.hasNoUnsignedWrap())
+    Shl->setHasNoUnsignedWrap();
+  if (I.hasNoSignedWrap() && C1->isNotMinSignedValue())
+    Shl->setHasNoSignedWrap();
+  return Shl;
+}
+
+// The caller must guarantee nsw/nuw properties.
+// X mul (C1 << N), where C1 is "1<<C2" and (C1 <<nsw/nuw N) or C2 is 0
+// --> X << (add N, C2)
+
+// X mul (zext (C1 << N)), where C1 is "1<<C2" and (C1 <<nsw/nuw N) or
+// C2 is 0 -->  X << (zext (add N, C2))
+static Instruction *foldMulShl(Value *Op0, Value *Op1, const BinaryOperator &I,
+                               InstCombiner &IC) {
+  Value *ShiftLeft;
+  if (!match(Op1, m_ZExt(m_Value(ShiftLeft))))
+    ShiftLeft = Op1;
+
+  Constant *CI;
+  Value *N;
+  if (!match(ShiftLeft, m_Shl(m_Constant(CI), m_Value(N))))
+    assert(false && "match should never fail in foldMulShl");
+  Constant *Log2Base = getLogBase2(N->getType(), CI);
+  assert(Log2Base && "getLogBase2 should never fail");
+  if (!Log2Base->isZeroValue())
+    N = IC.Builder.CreateAdd(N, Log2Base);
+  if (Op1 != ShiftLeft)
+    N = IC.Builder.CreateZExt(N, Op1->getType());
+  BinaryOperator *Shl = BinaryOperator::CreateShl(Op0, N);
+  if (I.hasNoUnsignedWrap())
+    Shl->setHasNoUnsignedWrap();
+  if (I.hasNoSignedWrap() && cast<ShlOperator>(ShiftLeft)->hasNoSignedWrap())
+    Shl->setHasNoSignedWrap();
+  return Shl;
+}
+
+// Recursively visits the possible right hand operands of a mul
+// instruction, seeing through select instructions, to determine if we can
+// replace the mul with something simpler. If we find that an operand is not
+// able to simplify the mul, we abort the entire transformation.
+static size_t visitMulOperand(Value *Op0, Value *Op1,
+                              SmallVectorImpl<OperandFoldAction> &Actions,
+                              unsigned Depth = 0) {
+  // Check to see if this is a multiplication with an exact power of 2, if so,
+  // convert to a left shift.
+  if (match(Op1, m_Power2())) {
+    Actions.emplace_back(foldMulPow2Cst, Op1);
+    return Actions.size();
+  }
+
+  // X mul (C1 << N), where C1 is "1<<C2" and (C1 <<nsw/nuw N) or C2 is 0
+  // --> X << (add N, C2)
+
+  // X mul (zext (C1 << N)), where C1 is "1<<C2" and (C1 <<nsw/nuw N) or
+  // C2 is 0 -->  X << (zext (add N, C2))
+  if (match(Op1, m_Shl(m_One(), m_Value())) ||
+      match(Op1, m_NSWShl(m_Power2(), m_Value())) ||
+      match(Op1, m_NUWShl(m_Power2(), m_Value())) ||
+      match(Op1, m_ZExt(m_Shl(m_One(), m_Value()))) ||
+      match(Op1, m_ZExt(m_NSWShl(m_Power2(), m_Value()))) ||
+      match(Op1, m_ZExt(m_NUWShl(m_Power2(), m_Value())))) {
+    Actions.emplace_back(foldMulShl, Op1);
+    return Actions.size();
+  }
+
+  // The remaining tests are all recursive, so bail out if we hit the limit.
+  if (Depth++ == MaxDepth)
+    return 0;
+
+  if (SelectInst *SI = dyn_cast<SelectInst>(Op1))
+    if (size_t LHSIdx = visitMulOperand(Op0, SI->getOperand(1), Actions, Depth))
+      if (visitMulOperand(Op0, SI->getOperand(2), Actions, Depth)) {
+        Actions.emplace_back(nullptr, Op1, LHSIdx - 1);
+        return Actions.size();
+      }
+
+  return 0;
+}
+
 // TODO: This is a specific form of a much more general pattern.
 //       We could detect a select with any binop identity constant, or we
 //       could use SimplifyBinOp to see if either arm of the select reduces.
 //       But that needs to be done carefully and/or while removing potential
 //       reverse canonicalizations as in InstCombiner::foldSelectIntoOp().
 static Value *foldMulSelectToNegate(BinaryOperator &I,
                                     InstCombiner::BuilderTy &Builder) {
   Value *Cond, *OtherOp;
Instruction *InstCombiner::visitMul(BinaryOperator &I) {
   }
 
   if (Instruction *FoldedMul = foldBinOpIntoSelectOrPhi(I))
     return FoldedMul;
 
   if (Value *FoldedMul = foldMulSelectToNegate(I, Builder))
     return replaceInstUsesWith(I, FoldedMul);
 
+  // (LHS mul (select (select (...)))) -> (LHS << (select (select (...))))
+  SmallVector<OperandFoldAction, 6> MulActions;
+  if (visitMulOperand(Op0, Op1, MulActions))
+    if (Instruction *Inst = combineActions(Op0, I, MulActions, *this))
+      return Inst;
+  MulActions.clear();
+  // ((select (select (...))) mul LHS) -> (LHS << (select (select (...))))
+  if (visitMulOperand(Op1, Op0, MulActions))
+    if (Instruction *Inst = combineActions(Op1, I, MulActions, *this))
+      return Inst;
+
   // Simplify mul instructions with a constant RHS.
   if (isa<Constant>(Op1)) {
     // Canonicalize (X+C1)*CI -> X*CI+C1*CI.
     Value *X;
     Constant *C1;
     if (match(Op0, m_OneUse(m_Add(m_Value(X), m_Constant(C1))))) {
       Value *Mul = Builder.CreateMul(C1, Op1);
       // Only go forward with the transform if C1*CI simplifies to a tidier
if (Div && Div->hasOneUse() &&
       return BinaryOperator::CreateSub(Rem, X);
     }
   }
 
   /// i1 mul -> i1 and.
   if (I.getType()->isIntOrIntVectorTy(1))
     return BinaryOperator::CreateAnd(Op0, Op1);
 
-  // X*(1 << Y) --> X << Y
-  // (1 << Y)*X --> X << Y
-  {
-    Value *Y;
-    BinaryOperator *BO = nullptr;
-    bool ShlNSW = false;
-    if (match(Op0, m_Shl(m_One(), m_Value(Y)))) {
-      BO = BinaryOperator::CreateShl(Op1, Y);
-      ShlNSW = cast<ShlOperator>(Op0)->hasNoSignedWrap();
-    } else if (match(Op1, m_Shl(m_One(), m_Value(Y)))) {
-      BO = BinaryOperator::CreateShl(Op0, Y);
-      ShlNSW = cast<ShlOperator>(Op1)->hasNoSignedWrap();
-    }
-    if (BO) {
-      if (I.hasNoUnsignedWrap())
-        BO->setHasNoUnsignedWrap();
-      if (I.hasNoSignedWrap() && ShlNSW)
-        BO->setHasNoSignedWrap();
-      return BO;
-    }
-  }
-
   // (bool X) * Y --> X ? Y : 0
   // Y * (bool X) --> X ? Y : 0
   if (match(Op0, m_ZExt(m_Value(X))) && X->getType()->isIntOrIntVectorTy(1))
     return SelectInst::Create(X, Op1, ConstantInt::get(I.getType(), 0));
   if (match(Op1, m_ZExt(m_Value(X))) && X->getType()->isIntOrIntVectorTy(1))
     return SelectInst::Create(X, Op0, ConstantInt::get(I.getType(), 0));
 
   // (lshr X, 31) * Y --> (ashr X, 31) & Y

llvm/test/Transforms/InstCombine/mul.ll

 ;
   %div = udiv exact i32 %x, %y
   %sel = select i1 %c, i32 %div, i32 1
   %mul = mul i32 %sel, %y
   ret i32 %mul
 }
 
 
-; TODO. 'select + mul' -> 'select + shl' for power of twos
+; 'select + mul' -> 'select + shl' for power of twos
 
 define i32 @shift_if_power2(i32 %x, i1 %cond) {
 ; CHECK-LABEL: @shift_if_power2(
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 16, i32 4
-; CHECK-NEXT:    [[R:%.*]] = mul i32 [[SEL]], [[X:%.*]]
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND:%.*]], i32 4, i32 2
+; CHECK-NEXT:    [[R:%.*]] = shl i32 [[X:%.*]], [[R_V]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %sel = select i1 %cond, i32 16, i32 4
   %r = mul i32 %sel, %x
   ret i32 %r
 }
 
 define i32 @shift_if_power2_nuw(i32 %x, i1 %cond) {
 ; CHECK-LABEL: @shift_if_power2_nuw(
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 16, i32 4
-; CHECK-NEXT:    [[R:%.*]] = mul nuw i32 [[SEL]], [[X:%.*]]
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND:%.*]], i32 4, i32 2
+; CHECK-NEXT:    [[R:%.*]] = shl nuw i32 [[X:%.*]], [[R_V]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %sel = select i1 %cond, i32 16, i32 4
   %r = mul nuw i32 %sel, %x
   ret i32 %r
 }
 
 define i32 @shift_if_power2_nsw(i32 %x, i1 %cond) {
 ; CHECK-LABEL: @shift_if_power2_nsw(
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 16, i32 4
-; CHECK-NEXT:    [[R:%.*]] = mul nsw i32 [[SEL]], [[X:%.*]]
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND:%.*]], i32 4, i32 2
+; CHECK-NEXT:    [[R:%.*]] = shl nsw i32 [[X:%.*]], [[R_V]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %sel = select i1 %cond, i32 16, i32 4
   %r = mul nsw i32 %sel, %x
   ret i32 %r
 }
 
 define i32 @shift_if_power2_nuw_nsw(i32 %x, i1 %cond) {
 ; CHECK-LABEL: @shift_if_power2_nuw_nsw(
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 16, i32 4
-; CHECK-NEXT:    [[R:%.*]] = mul nuw nsw i32 [[SEL]], [[X:%.*]]
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND:%.*]], i32 4, i32 2
+; CHECK-NEXT:    [[R:%.*]] = shl nuw nsw i32 [[X:%.*]], [[R_V]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %sel = select i1 %cond, i32 16, i32 4
   %r = mul nuw nsw i32 %sel, %x
   ret i32 %r
 }
 
 define i32 @shift_if_power2_nuw_nsw_min(i32 %x, i1 %cond) {
 ; CHECK-LABEL: @shift_if_power2_nuw_nsw_min(
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 2, i32 -2147483648
-; CHECK-NEXT:    [[R:%.*]] = mul nuw nsw i32 [[SEL]], [[X:%.*]]
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND:%.*]], i32 1, i32 31
+; CHECK-NEXT:    [[R:%.*]] = shl nuw i32 [[X:%.*]], [[R_V]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %sel = select i1 %cond, i32 2, i32 -2147483648
   %r = mul nuw nsw i32 %sel, %x
   ret i32 %r
 }
 
 define i32 @shift_if_power2_double_select(i32 %x, i32 %y, i1 %cond1, i1 %cond2) {
 ; CHECK-LABEL: @shift_if_power2_double_select(
 ; CHECK-NEXT:    [[SHL_RES:%.*]] = shl i32 8, [[Y:%.*]]
 ; CHECK-NEXT:    [[SEL1:%.*]] = select i1 [[COND1:%.*]], i32 [[SHL_RES]], i32 1024
 ; CHECK-NEXT:    [[SEL2:%.*]] = select i1 [[COND2:%.*]], i32 16, i32 [[SEL1]]
 ; CHECK-NEXT:    [[R:%.*]] = mul nuw i32 [[SEL2]], [[X:%.*]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %shl.res = shl i32 8, %y
   %sel1 = select i1 %cond1, i32 %shl.res, i32 1024
   %sel2 = select i1 %cond2, i32 16, i32 %sel1
   %r = mul nuw i32 %x, %sel2
   ret i32 %r
 }
 
 define i32 @shift_if_power2_double_select_zext(i32 %x, i16 %y, i1 %cond1, i1 %cond2) {
 ; CHECK-LABEL: @shift_if_power2_double_select_zext(
-; CHECK-NEXT:    [[SHL_RES:%.*]] = shl nsw i16 8, [[Y:%.*]]
-; CHECK-NEXT:    [[SHL:%.*]] = zext i16 [[SHL_RES]] to i32
-; CHECK-NEXT:    [[SEL1:%.*]] = select i1 [[COND1:%.*]], i32 [[SHL]], i32 1024
-; CHECK-NEXT:    [[SEL2:%.*]] = select i1 [[COND2:%.*]], i32 16, i32 [[SEL1]]
-; CHECK-NEXT:    [[R:%.*]] = mul nuw i32 [[SEL2]], [[X:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i16 [[Y:%.*]], 3
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i16 [[TMP1]] to i32
+; CHECK-NEXT:    [[DOTV:%.*]] = select i1 [[COND1:%.*]], i32 [[TMP2]], i32 10
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND2:%.*]], i32 4, i32 [[DOTV]]
+; CHECK-NEXT:    [[R:%.*]] = shl nuw i32 [[X:%.*]], [[R_V]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %shl.res = shl nsw i16 8, %y
   %shl = zext i16 %shl.res to i32
   %sel1 = select i1 %cond1, i32 %shl, i32 1024
   %sel2 = select i1 %cond2, i32 16, i32 %sel1
   %r = mul nuw i32 %x, %sel2
   ret i32 %r
;
   %shl.res = shl i16 8, %y
   %shl = zext i16 %shl.res to i32
   %sel1 = select i1 %cond1, i32 %shl, i32 1024
   %sel2 = select i1 %cond2, i32 16, i32 %sel1
   %r = mul nsw i32 %x, %sel2
   ret i32 %r
 }
 
+define i32 @shift_if_power2_double_select_zext_shift_1(i32 %x, i16 %y, i1 %cond1, i1 %cond2) {
+; CHECK-LABEL: @shift_if_power2_double_select_zext_shift_1(
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i16 [[Y:%.*]] to i32
+; CHECK-NEXT:    [[DOTV:%.*]] = select i1 [[COND1:%.*]], i32 [[TMP1]], i32 10
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND2:%.*]], i32 4, i32 [[DOTV]]
+; CHECK-NEXT:    [[R:%.*]] = shl i32 [[X:%.*]], [[R_V]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %shl.res = shl i16 1, %y
+  %shl = zext i16 %shl.res to i32
+  %sel1 = select i1 %cond1, i32 %shl, i32 1024
+  %sel2 = select i1 %cond2, i32 16, i32 %sel1
+  %r = mul nsw i32 %x, %sel2
+  ret i32 %r
+}
+
 define i32 @shift_if_power2_double_select_zext_shift_nsw(i32 %x, i16 %y, i1 %cond1, i1 %cond2) {
 ; CHECK-LABEL: @shift_if_power2_double_select_zext_shift_nsw(
-; CHECK-NEXT:    [[SHL_RES:%.*]] = shl nsw i16 8, [[Y:%.*]]
-; CHECK-NEXT:    [[SHL:%.*]] = zext i16 [[SHL_RES]] to i32
-; CHECK-NEXT:    [[SEL1:%.*]] = select i1 [[COND1:%.*]], i32 [[SHL]], i32 1024
-; CHECK-NEXT:    [[SEL2:%.*]] = select i1 [[COND2:%.*]], i32 16, i32 [[SEL1]]
-; CHECK-NEXT:    [[R:%.*]] = mul nuw nsw i32 [[SEL2]], [[X:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i16 [[Y:%.*]], 3
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i16 [[TMP1]] to i32
+; CHECK-NEXT:    [[DOTV:%.*]] = select i1 [[COND1:%.*]], i32 [[TMP2]], i32 10
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND2:%.*]], i32 4, i32 [[DOTV]]
+; CHECK-NEXT:    [[R:%.*]] = shl nuw nsw i32 [[X:%.*]], [[R_V]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %shl.res = shl nsw i16 8, %y
   %shl = zext i16 %shl.res to i32
   %sel1 = select i1 %cond1, i32 %shl, i32 1024
   %sel2 = select i1 %cond2, i32 16, i32 %sel1
   %r = mul nuw nsw i32 %x, %sel2
   ret i32 %r
 }
 
+define i32 @shift_if_power2_double_select_zext_shift_nuw(i32 %x, i16 %y, i1 %cond1, i1 %cond2) {
+; CHECK-LABEL: @shift_if_power2_double_select_zext_shift_nuw(
+; CHECK-NEXT:    [[TMP1:%.*]] = add i16 [[Y:%.*]], 3
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i16 [[TMP1]] to i32
+; CHECK-NEXT:    [[DOTV:%.*]] = select i1 [[COND1:%.*]], i32 [[TMP2]], i32 10
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND2:%.*]], i32 4, i32 [[DOTV]]
+; CHECK-NEXT:    [[R:%.*]] = shl nuw i32 [[X:%.*]], [[R_V]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %shl.res = shl nuw i16 8, %y
+  %shl = zext i16 %shl.res to i32
+  %sel1 = select i1 %cond1, i32 %shl, i32 1024
+  %sel2 = select i1 %cond2, i32 16, i32 %sel1
+  %r = mul nuw nsw i32 %x, %sel2
+  ret i32 %r
+}
+
 define i32 @shift_if_power2_zero(i32 %x, i1 %cond) {
 ; CHECK-LABEL: @shift_if_power2_zero(
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 1, i32 4
-; CHECK-NEXT:    [[R:%.*]] = mul i32 [[SEL]], [[X:%.*]]
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND:%.*]], i32 0, i32 2
+; CHECK-NEXT:    [[R:%.*]] = shl i32 [[X:%.*]], [[R_V]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %sel = select i1 %cond, i32 1, i32 4
   %r = mul i32 %sel, %x
   ret i32 %r
 }
 
 define <2 x i8> @shift_if_power2_vector(<2 x i8> %px, i1 %cond) {
 ; CHECK-LABEL: @shift_if_power2_vector(
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], <2 x i8> <i8 16, i8 4>, <2 x i8> <i8 1, i8 32>
-; CHECK-NEXT:    [[R:%.*]] = mul <2 x i8> [[SEL]], [[PX:%.*]]
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND:%.*]], <2 x i8> <i8 4, i8 2>, <2 x i8> <i8 0, i8 5>
+; CHECK-NEXT:    [[R:%.*]] = shl <2 x i8> [[PX:%.*]], [[R_V]]
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %sel = select i1 %cond, <2 x i8> <i8 16, i8 4>, <2 x i8> <i8 1, i8 32>
   %r = mul <2 x i8> %px, %sel
   ret <2 x i8> %r
 }
 
+define <2 x i32> @shift_if_power2_vector_nsw(<2 x i32> %px, i1 %cond) {
+; CHECK-LABEL: @shift_if_power2_vector_nsw(
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND:%.*]], <2 x i32> <i32 2, i32 31>, <2 x i32> <i32 0, i32 5>
+; CHECK-NEXT:    [[R:%.*]] = shl <2 x i32> [[PX:%.*]], [[R_V]]
+; CHECK-NEXT:    ret <2 x i32> [[R]]
+;
+  %sel = select i1 %cond, <2 x i32> <i32 4, i32 -2147483648>, <2 x i32> <i32 1, i32 32>
+  %r = mul nsw <2 x i32> %px, %sel
+  ret <2 x i32> %r
+}
+
 define i32 @shift_if_extra_use(i32 %x, i1 %cond) {
 ; CHECK-LABEL: @shift_if_extra_use(
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 4, i32 128
 ; CHECK-NEXT:    call void @use32(i32 [[SEL]])
-; CHECK-NEXT:    [[R:%.*]] = mul i32 [[SEL]], [[X:%.*]]
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND]], i32 2, i32 7
+; CHECK-NEXT:    [[R:%.*]] = shl i32 [[X:%.*]], [[R_V]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %sel = select i1 %cond, i32 4, i32 128
   call void @use32(i32 %sel)
   %r = mul i32 %sel, %x
   ret i32 %r
 }
 
;
   %r = mul <2 x i8> %px, %sel
   ret <2 x i8> %r
 }
 
 ; Undef select + mul tests
 
 define <2 x i8> @shift_if_undef_vector(<2 x i8> %px, i1 %cond) {
 ; CHECK-LABEL: @shift_if_undef_vector(
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], <2 x i8> <i8 16, i8 4>, <2 x i8> <i8 undef, i8 32>
-; CHECK-NEXT:    [[R:%.*]] = mul <2 x i8> [[SEL]], [[PX:%.*]]
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND:%.*]], <2 x i8> <i8 4, i8 2>, <2 x i8> <i8 0, i8 5>
+; CHECK-NEXT:    [[R:%.*]] = shl <2 x i8> [[PX:%.*]], [[R_V]]
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %sel = select i1 %cond, <2 x i8> <i8 16, i8 4>, <2 x i8> <i8 undef, i8 32>
   %r = mul <2 x i8> %px, %sel
   ret <2 x i8> %r
 }
 
 define <2 x i8> @shift_if_different_lanes_undef_vector(<2 x i8> %px, i1 %cond) {
 ; CHECK-LABEL: @shift_if_different_lanes_undef_vector(
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], <2 x i8> <i8 16, i8 undef>, <2 x i8> <i8 undef, i8 32>
-; CHECK-NEXT:    [[R:%.*]] = mul <2 x i8> [[SEL]], [[PX:%.*]]
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND:%.*]], <2 x i8> <i8 4, i8 0>, <2 x i8> <i8 0, i8 5>
+; CHECK-NEXT:    [[R:%.*]] = shl <2 x i8> [[PX:%.*]], [[R_V]]
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %sel = select i1 %cond, <2 x i8> <i8 16, i8 undef>, <2 x i8> <i8 undef, i8 32>
   %r = mul <2 x i8> %px, %sel
   ret <2 x i8> %r
 }
 
 define <2 x i8> @shift_if_same_lane_undef_vector(<2 x i8> %px, i1 %cond) {
 ; CHECK-LABEL: @shift_if_same_lane_undef_vector(
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], <2 x i8> <i8 undef, i8 4>, <2 x i8> <i8 undef, i8 32>
-; CHECK-NEXT:    [[R:%.*]] = mul <2 x i8> [[SEL]], [[PX:%.*]]
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND:%.*]], <2 x i8> <i8 0, i8 2>, <2 x i8> <i8 0, i8 5>
+; CHECK-NEXT:    [[R:%.*]] = shl <2 x i8> [[PX:%.*]], [[R_V]]
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %sel = select i1 %cond, <2 x i8> <i8 undef, i8 4>, <2 x i8> <i8 undef, i8 32>
   %r = mul <2 x i8> %px, %sel
   ret <2 x i8> %r
 }

llvm/test/Transforms/InstCombine/udiv_select_to_select_shift.ll

;
   %divisor1 = select i1 %Cond, i64 16, i64 8
   %quotient1 = udiv i64 %X, %divisor1
   %divisor2 = select i1 %Cond, i64 8, i64 0
   %quotient2 = udiv i64 %X, %divisor2
   %sum = add i64 %quotient1, %quotient2
   ret i64 %sum
 }
 
+define <2 x i8> @test2(<2 x i8> %px, i1 %cond) {
+; CHECK-LABEL: @test2(
+; CHECK-NEXT:    [[R_V:%.*]] = select i1 [[COND:%.*]], <2 x i8> <i8 4, i8 0>, <2 x i8> <i8 undef, i8 5>
+; CHECK-NEXT:    [[R:%.*]] = lshr <2 x i8> [[PX:%.*]], [[R_V]]
+; CHECK-NEXT:    ret <2 x i8> [[R]]
+;
+  %sel = select i1 %cond, <2 x i8> <i8 16, i8 1>, <2 x i8> <i8 undef, i8 32>
+  %r = udiv <2 x i8> %px, %sel
+  ret <2 x i8> %r
+}
+
+define <2 x i8> @test3(<2 x i8> %px, i1 %cond) {
+; CHECK-LABEL: @test3(
+; CHECK-NEXT:    ret <2 x i8> undef
+;
+  %sel = select i1 %cond, <2 x i8> <i8 16, i8 undef>, <2 x i8> <i8 undef, i8 32>
+  %r = udiv <2 x i8> %px, %sel
+  ret <2 x i8> %r
+}
+
+define <2 x i8> @test4(<2 x i8> %px, i1 %cond) {
+; CHECK-LABEL: @test4(
+; CHECK-NEXT:    ret <2 x i8> undef
+;
+  %sel = select i1 %cond, <2 x i8> <i8 undef, i8 1>, <2 x i8> <i8 undef, i8 32>
+  %r = udiv <2 x i8> %px, %sel
+  ret <2 x i8> %r
+}
+
 ; https://bugs.llvm.org/show_bug.cgi?id=34856
 ; This would assert/crash because we didn't propagate the condition with the correct vector type.
 
 define <2 x i32> @PR34856(<2 x i32> %t0, <2 x i32> %t1) {
 ; CHECK-LABEL: @PR34856(
 ; CHECK-NEXT:    [[TMP1:%.*]] = icmp ugt <2 x i32> [[T1:%.*]], <i32 -8, i32 -8>
 ; CHECK-NEXT:    [[DIV1:%.*]] = zext <2 x i1> [[TMP1]] to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[DIV1]]