[llvm] 96ef699 - [InstCombine] Fix a couple crashes with extractelement on a scalable vector.

[Eli Friedman via llvm-commits]

Author: Eli Friedman
Date: 2020-09-02T18:02:07-07:00
New Revision: 96ef6998dfcc4d7f2c7fcb87eb9369e94a40e127

URL: https://github.com/llvm/llvm-project/commit/96ef6998dfcc4d7f2c7fcb87eb9369e94a40e127
DIFF: https://github.com/llvm/llvm-project/commit/96ef6998dfcc4d7f2c7fcb87eb9369e94a40e127.diff

LOG: [InstCombine] Fix a couple crashes with extractelement on a scalable vector.

Differential Revision: https://reviews.llvm.org/D86989

Added: 
    

Modified: 
    llvm/lib/Analysis/ValueTracking.cpp
    llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
    llvm/test/Transforms/InstCombine/vscale_extractelement.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index b592412ed0b6..3f7cf296ac1c 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -2568,11 +2568,13 @@ bool isKnownNonZero(const Value *V, const APInt &DemandedElts, unsigned Depth,
     const Value *Vec = EEI->getVectorOperand();
     const Value *Idx = EEI->getIndexOperand();
     auto *CIdx = dyn_cast<ConstantInt>(Idx);
-    unsigned NumElts = cast<FixedVectorType>(Vec->getType())->getNumElements();
-    APInt DemandedVecElts = APInt::getAllOnesValue(NumElts);
-    if (CIdx && CIdx->getValue().ult(NumElts))
-      DemandedVecElts = APInt::getOneBitSet(NumElts, CIdx->getZExtValue());
-    return isKnownNonZero(Vec, DemandedVecElts, Depth, Q);
+    if (auto *VecTy = dyn_cast<FixedVectorType>(Vec->getType())) {
+      unsigned NumElts = VecTy->getNumElements();
+      APInt DemandedVecElts = APInt::getAllOnesValue(NumElts);
+      if (CIdx && CIdx->getValue().ult(NumElts))
+        DemandedVecElts = APInt::getOneBitSet(NumElts, CIdx->getZExtValue());
+      return isKnownNonZero(Vec, DemandedVecElts, Depth, Q);
+    }
   }
 
   KnownBits Known(BitWidth);

diff  --git a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
index 7cbbcddc14ba..35db21245754 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
@@ -2224,9 +2224,7 @@ static Instruction *canonicalizeBitCastExtElt(BitCastInst &BitCast,
   if (!VectorType::isValidElementType(DestType))
     return nullptr;
 
-  unsigned NumElts =
-      cast<FixedVectorType>(ExtElt->getVectorOperandType())->getNumElements();
-  auto *NewVecType = FixedVectorType::get(DestType, NumElts);
+  auto *NewVecType = VectorType::get(DestType, ExtElt->getVectorOperandType());
   auto *NewBC = IC.Builder.CreateBitCast(ExtElt->getVectorOperand(),
                                          NewVecType, "bc");
   return ExtractElementInst::Create(NewBC, ExtElt->getIndexOperand());

diff  --git a/llvm/test/Transforms/InstCombine/vscale_extractelement.ll b/llvm/test/Transforms/InstCombine/vscale_extractelement.ll
index 8dcb16740869..0b148172a42a 100644
--- a/llvm/test/Transforms/InstCombine/vscale_extractelement.ll
+++ b/llvm/test/Transforms/InstCombine/vscale_extractelement.ll
@@ -146,3 +146,25 @@ define <vscale x 4 x i32> @extractelement_insertelement_
diff _positions(<vscale x
   %4 = insertelement <vscale x 4 x i32> %3, i32 %vec.e3, i32 3
   ret <vscale x 4 x i32> %4
 }
+
+define i32 @bitcast_of_extractelement( <vscale x 2 x float> %d) {
+; CHECK-LABEL: @bitcast_of_extractelement(
+; CHECK-NEXT:    [[BC:%.*]] = bitcast <vscale x 2 x float> [[D:%.*]] to <vscale x 2 x i32>
+; CHECK-NEXT:    [[CAST:%.*]] = extractelement <vscale x 2 x i32> [[BC]], i32 0
+; CHECK-NEXT:    ret i32 [[CAST]]
+;
+  %ext = extractelement <vscale x 2 x float> %d, i32 0
+  %cast = bitcast float %ext to i32
+  ret i32 %cast
+}
+
+define i1 @extractelement_is_zero(<vscale x 2 x i32> %d, i1 %b, i32 %z) {
+; CHECK-LABEL: @extractelement_is_zero(
+; CHECK-NEXT:    [[EXT:%.*]] = extractelement <vscale x 2 x i32> [[D:%.*]], i32 0
+; CHECK-NEXT:    [[BB:%.*]] = icmp eq i32 [[EXT]], 0
+; CHECK-NEXT:    ret i1 [[BB]]
+;
+  %ext = extractelement <vscale x 2 x i32> %d, i32 0
+  %bb = icmp eq i32 %ext, 0
+  ret i1 %bb
+}