Address review comments

Author: david-arm

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -469,8 +469,9 @@ Value *VPInstruction::generate(VPTransformState &State) {
     return Builder.CreateNot(A, Name);
   }
   case Instruction::ExtractElement: {
+    assert(State.VF.isVector() && "Only extract elements from vectors");
     Value *Vec = State.get(getOperand(0));
-    Value *Idx = State.get(getOperand(1), true);
+    Value *Idx = State.get(getOperand(1), /*IsScalar=*/true);
     return Builder.CreateExtractElement(Vec, Idx, Name);
   }
   case Instruction::ICmp: {
@@ -887,6 +888,8 @@ bool VPInstruction::onlyFirstLaneUsed(const VPValue *Op) const {
     return false;
   case Instruction::PHI:
     return true;
+  case Instruction::ExtractElement:
+    return Op == getOperand(1);
   case Instruction::ICmp:
   case Instruction::Select:
   case Instruction::Or:

llvm/test/Transforms/LoopVectorize/AArch64/early_exit_costs.ll

@@ -14,7 +14,7 @@ define i64 @same_exit_block_pre_inc_use1_sve() #1 {
 ; CHECK-NEXT: Cost of 4 for VF vscale x 16: EMIT vp<{{.*}}> = first-active-lane vp<{{.*}}>
 ; CHECK-NEXT: Cost of 2 for VF vscale x 16: EMIT vp<{{.*}}> = extractelement ir<{{.*}}>, vp<{{.*}}>
 ; CHECK-NEXT: Cost of 4 for VF vscale x 16: EMIT vp<{{.*}}>.1 = first-active-lane vp<{{.*}}>
-; CHECK-NEXT: Cost of 2 for VF vscale x 16: EMIT vp<{{.*}}>.1 = extractelement ir<{{.*}}>, vp<%first.active.lane>.1
+; CHECK-NEXT: Cost of 2 for VF vscale x 16: EMIT vp<{{.*}}>.1 = extractelement ir<{{.*}}>, vp<{{.*}}>.1
 ; CHECK: LV: Minimum required TC for runtime checks to be profitable:32
 entry:
   %p1 = alloca [1024 x i8]
@@ -53,7 +53,7 @@ define i64 @same_exit_block_pre_inc_use1_nosve() {
 ; CHECK-NEXT: Cost of 48 for VF 16: EMIT vp<{{.*}}> = first-active-lane vp<{{.*}}>
 ; CHECK-NEXT: Cost of 2 for VF 16: EMIT vp<{{.*}}> = extractelement ir<{{.*}}>, vp<{{.*}}>
 ; CHECK-NEXT: Cost of 48 for VF 16: EMIT vp<{{.*}}>.1 = first-active-lane vp<{{.*}}>
-; CHECK-NEXT: Cost of 2 for VF 16: EMIT vp<{{.*}}>.1 = extractelement ir<{{.*}}>, vp<%first.active.lane>.1
+; CHECK-NEXT: Cost of 2 for VF 16: EMIT vp<{{.*}}>.1 = extractelement ir<{{.*}}>, vp<{{.*}}>.1
 ; CHECK: LV: Minimum required TC for runtime checks to be profitable:176
 ; CHECK-NEXT: LV: Vectorization is not beneficial: expected trip count < minimum profitable VF (64 < 176)
 ; CHECK-NEXT: LV: Too many memory checks needed.