[InstCombine] Fix issue with X udiv 2^C -> X >> C for non-splat constant vectors

foldUDivPow2Cst was assuming that the input was a scalar or a splat constant

llvm-svn: 324608

2 files changed, 25 insertions, 3 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
index e751c64caae..9efc797d464 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
@@ -1055,9 +1055,10 @@ struct UDivFoldAction {
 // X udiv 2^C -> X >> C
 static Instruction *foldUDivPow2Cst(Value *Op0, Value *Op1,
                                     const BinaryOperator &I, InstCombiner &IC) {
-  const APInt &C = cast<Constant>(Op1)->getUniqueInteger();
-  BinaryOperator *LShr = BinaryOperator::CreateLShr(
-      Op0, ConstantInt::get(Op0->getType(), C.logBase2()));
+  Constant *C1 = getLogBase2(Op0->getType(), cast<Constant>(Op1));
+  if (!C1)
+    llvm_unreachable("Failed to constant fold udiv -> logbase2");
+  BinaryOperator *LShr = BinaryOperator::CreateLShr(Op0, C1);
   if (I.isExact())
     LShr->setIsExact();
   return LShr;
diff --git a/llvm/test/Transforms/InstCombine/vector-udiv.ll b/llvm/test/Transforms/InstCombine/vector-udiv.ll
new file mode 100644
index 00000000000..5052661ee4e
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/vector-udiv.ll
@@ -0,0 +1,21 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -instcombine -S | FileCheck %s
+
+define <4 x i32> @test_v4i32_splatconst_pow2(<4 x i32> %a0) {
+; CHECK-LABEL: @test_v4i32_splatconst_pow2(
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr <4 x i32> [[A0:%.*]], <i32 1, i32 1, i32 1, i32 1>
+; CHECK-NEXT:    ret <4 x i32> [[TMP1]]
+;
+  %1 = udiv <4 x i32> %a0, <i32 2, i32 2, i32 2, i32 2>
+  ret <4 x i32> %1
+}
+
+define <4 x i32> @test_v4i32_const_pow2(<4 x i32> %a0) {
+; CHECK-LABEL: @test_v4i32_const_pow2(
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr <4 x i32> [[A0:%.*]], <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT:    ret <4 x i32> [[TMP1]]
+;
+  %1 = udiv <4 x i32> %a0, <i32 1, i32 2, i32 4, i32 8>
+  ret <4 x i32> %1
+}
+