!fixup address comments, thanks!

Author: fhahn

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

@@ -1782,15 +1782,18 @@ bool VectorCombine::scalarizeExtExtract(Instruction &I) {
   // Try to convert a vector zext feeding only extracts to a set of scalar (Src
   // << ExtIdx *Size) & (Size -1), if profitable.
   auto *Ext = cast<ZExtInst>(&I);
-  auto *SrcTy = cast<FixedVectorType>(Ext->getOperand(0)->getType());
+  auto *SrcTy = dyn_cast<FixedVectorType>(Ext->getOperand(0)->getType());
+  if (!SrcTy)
+    return false;
   auto *DstTy = cast<FixedVectorType>(Ext->getType());
 
-  if (DL->getTypeSizeInBits(SrcTy) !=
-      DL->getTypeSizeInBits(DstTy->getElementType()))
+  Type *ScalarDstTy = DstTy->getElementType();
+  if (DL->getTypeSizeInBits(SrcTy) != DL->getTypeSizeInBits(ScalarDstTy))
     return false;
 
-  InstructionCost VectorCost = TTI.getCastInstrCost(
-      Instruction::ZExt, DstTy, SrcTy, TTI::CastContextHint::None, CostKind);
+  InstructionCost VectorCost =
+      TTI.getCastInstrCost(Instruction::ZExt, DstTy, SrcTy,
+                           TTI::CastContextHint::None, CostKind, Ext);
   unsigned ExtCnt = 0;
   bool ExtLane0 = false;
   for (User *U : Ext->users()) {
@@ -1805,15 +1808,13 @@ bool VectorCombine::scalarizeExtExtract(Instruction &I) {
                                          CostKind, Idx->getZExtValue(), U);
   }
 
-  Type *ScalarDstTy = DstTy->getElementType();
   InstructionCost ScalarCost =
       ExtCnt * TTI.getArithmeticInstrCost(
                    Instruction::And, ScalarDstTy, CostKind,
                    {TTI::OK_AnyValue, TTI::OP_None},
                    {TTI::OK_NonUniformConstantValue, TTI::OP_None}) +
       (ExtCnt - ExtLane0) *
           TTI.getArithmeticInstrCost(
-
               Instruction::LShr, ScalarDstTy, CostKind,
               {TTI::OK_AnyValue, TTI::OP_None},
               {TTI::OK_NonUniformConstantValue, TTI::OP_None});

llvm/test/Transforms/VectorCombine/AArch64/ext-extract.ll

@@ -329,3 +329,32 @@ entry:
   call void @use.i64(i64 %ext.1)
   ret void
 }
+
+define void @zext_nv4i8_all_lanes_used(<vscale x 4 x i8> %src) {
+; CHECK-LABEL: define void @zext_nv4i8_all_lanes_used(
+; CHECK-SAME: <vscale x 4 x i8> [[SRC:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[EXT9:%.*]] = zext nneg <vscale x 4 x i8> [[SRC]] to <vscale x 4 x i32>
+; CHECK-NEXT:    [[EXT_0:%.*]] = extractelement <vscale x 4 x i32> [[EXT9]], i64 0
+; CHECK-NEXT:    [[EXT_1:%.*]] = extractelement <vscale x 4 x i32> [[EXT9]], i64 1
+; CHECK-NEXT:    [[EXT_2:%.*]] = extractelement <vscale x 4 x i32> [[EXT9]], i64 2
+; CHECK-NEXT:    [[EXT_3:%.*]] = extractelement <vscale x 4 x i32> [[EXT9]], i64 3
+; CHECK-NEXT:    call void @use.i32(i32 [[EXT_0]])
+; CHECK-NEXT:    call void @use.i32(i32 [[EXT_1]])
+; CHECK-NEXT:    call void @use.i32(i32 [[EXT_2]])
+; CHECK-NEXT:    call void @use.i32(i32 [[EXT_3]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %ext9 = zext nneg <vscale x 4 x i8> %src to <vscale x 4 x i32>
+  %ext.0 = extractelement <vscale x 4 x i32> %ext9, i64 0
+  %ext.1 = extractelement <vscale x 4 x i32> %ext9, i64 1
+  %ext.2 = extractelement <vscale x 4 x i32> %ext9, i64 2
+  %ext.3 = extractelement <vscale x 4 x i32> %ext9, i64 3
+
+  call void @use.i32(i32 %ext.0)
+  call void @use.i32(i32 %ext.1)
+  call void @use.i32(i32 %ext.2)
+  call void @use.i32(i32 %ext.3)
+  ret void
+}