[InstCombine] Fix a couple crahses with extractelement on a scalable vector.

[david-arm]

LGTM

[sdesmalen]

llvm/lib/Analysis/ValueTracking.cpp

2573

nit: is this how whitespace changes are rendered after the latest Phabricator update, or are these tabs?

llvm/test/Transforms/InstCombine/vscale_extractelement.ll

168

nit: entry: can be removed.

[david-arm]

llvm/lib/Analysis/ValueTracking.cpp

2573

I've seen these rendered on mine and other patches too where there are no tabs. I think it's the latest Phabricator update showing indentation.

[sdesmalen]

llvm/lib/Analysis/ValueTracking.cpp

2573

Good to know, that's a nice improvement to Phabricator! Thanks for confirming.

llvm/lib/Analysis/ValueTracking.cpp

else if (const PHINode *PN = dyn_cast<PHINode>(V)) {
     if (AllNonZeroConstants)
       return true;
   }
   // ExtractElement
   else if (const auto *EEI = dyn_cast<ExtractElementInst>(V)) {
     const Value *Vec = EEI->getVectorOperand();
     const Value *Idx = EEI->getIndexOperand();
     auto *CIdx = dyn_cast<ConstantInt>(Idx);
-    unsigned NumElts = cast<FixedVectorType>(Vec->getType())->getNumElements();
+    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);
   computeKnownBits(V, DemandedElts, Known, Depth, Q);
   return Known.One != 0;
 }
 
 bool isKnownNonZero(const Value* V, unsigned Depth, const Query& Q) {
   // FIXME: We currently have no way to represent the DemandedElts of a scalable

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

if (!ExtElt || !ExtElt->hasOneUse())
     return nullptr;
 
   // The bitcast must be to a vectorizable type, otherwise we can't make a new
   // type to extract from.
   Type *DestType = BitCast.getType();
   if (!VectorType::isValidElementType(DestType))
     return nullptr;
 
-  unsigned NumElts = 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());
 }
 
 /// Change the type of a bitwise logic operation if we can eliminate a bitcast.
 static Instruction *foldBitCastBitwiseLogic(BitCastInst &BitCast,
                                             InstCombiner::BuilderTy &Builder) {

llvm/test/Transforms/InstCombine/vscale_extractelement.ll

;
   %vec.e2 = extractelement <vscale x 4 x i32> %vec, i32 6
   %vec.e3 = extractelement <vscale x 4 x i32> %vec, i32 7
   %1 = insertelement <vscale x 4 x i32> %vec, i32 %vec.e0, i32 0
   %2 = insertelement <vscale x 4 x i32> %1, i32 %vec.e1, i32 1
   %3 = insertelement <vscale x 4 x i32> %2, i32 %vec.e2, i32 2
   %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:  entry:
+; 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]]
+;
+entry:
+  %ext = extractelement <vscale x 2 x i32> %d, i32 0
+  %bb = icmp eq i32 %ext, 0
+  ret i1 %bb
+}