!fixup address lates comments, thanks!

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -168,6 +168,8 @@ class VPLane {
   static VPLane getFirstLane() { return VPLane(0, VPLane::Kind::First); }
 
   static VPLane getLaneFromEnd(const ElementCount &VF, unsigned Offset) {
+    assert(Offset <= VF.getKnownMinValue() &&
+           "trying to extract with invalid offset");
     unsigned LaneOffset = VF.getKnownMinValue() - Offset;
     Kind LaneKind;
     if (VF.isScalable())
@@ -1186,6 +1188,10 @@ class VPInstruction : public VPRecipeWithIRFlags {
     BranchOnCount,
     BranchOnCond,
     ComputeReductionResult,
+    // Takes the VPValue to extract from as first operand and the lane to
+    // extract from as second operand. The second operand must be a constant and
+    // <= VF when extracting from a vector or <= UF when extracting from a
+    // scalar.
     ExtractFromEnd,
     LogicalAnd, // Non-poison propagating logical And.
     // Add an offset in bytes (second operand) to a base pointer (first
@@ -1224,6 +1230,10 @@ class VPInstruction : public VPRecipeWithIRFlags {
   /// value for lane \p Lane.
   Value *generatePerLane(VPTransformState &State, const VPIteration &Lane);
 
+  /// Returns true if this VPInstruction converts a vector value to a scalar,
+  /// e.g. by performing a reduction or extracting a lane.
+  bool isVectorToScalar() const;
+
 #if !defined(NDEBUG)
   /// Return true if the VPInstruction is a floating point math operation, i.e.
   /// has fast-math flags.

llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp

@@ -45,6 +45,12 @@ Type *VPTypeAnalysis::inferScalarTypeForRecipe(const VPInstruction *R) {
     CachedTypes[OtherV] = ResTy;
     return ResTy;
   }
+  case VPInstruction::ExtractFromEnd: {
+    Type *BaseTy = inferScalarType(R->getOperand(0));
+    if (auto *VecTy = dyn_cast<VectorType>(BaseTy))
+      return VecTy->getElementType();
+    return BaseTy;
+  }
   case VPInstruction::Not: {
     Type *ResTy = inferScalarType(R->getOperand(0));
     assert(IntegerType::get(Ctx, 1) == ResTy &&

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -294,14 +294,13 @@ bool VPInstruction::doesGeneratePerAllLanes() const {
 bool VPInstruction::canGenerateScalarForFirstLane() const {
   if (Instruction::isBinaryOp(getOpcode()))
     return true;
-
+  if (isVectorToScalar())
+    return true;
   switch (Opcode) {
   case VPInstruction::BranchOnCond:
   case VPInstruction::BranchOnCount:
   case VPInstruction::CalculateTripCountMinusVF:
   case VPInstruction::CanonicalIVIncrementForPart:
-  case VPInstruction::ComputeReductionResult:
-  case VPInstruction::ExtractFromEnd:
   case VPInstruction::PtrAdd:
   case VPInstruction::ExplicitVectorLength:
     return true;
@@ -567,17 +566,15 @@ Value *VPInstruction::generatePerPart(VPTransformState &State, unsigned Part) {
     auto *CI = cast<ConstantInt>(getOperand(1)->getLiveInIRValue());
     unsigned Offset = CI->getZExtValue();
 
-    // Extract lane VF - Offset in the from the operand.
     Value *Res;
     if (State.VF.isVector()) {
+      // Extract lane VF - Offset from the operand.
       Res = State.get(
           getOperand(0),
           VPIteration(State.UF - 1, VPLane::getLaneFromEnd(State.VF, Offset)));
     } else {
       assert(State.UF > 1 && "VF and UF cannot both be 1");
-      // When loop is unrolled without vectorizing, retrieve the value just
-      // prior to the final unrolled value. This is analogous to the vectorized
-      // case above: extracting the second last element when VF > 1.
+      // When loop is unrolled without vectorizing, retrieve UF - Offset.
       Res = State.get(getOperand(0), State.UF - Offset);
     }
     Res->setName(Name);
@@ -600,6 +597,11 @@ Value *VPInstruction::generatePerPart(VPTransformState &State, unsigned Part) {
   }
 }
 
+bool VPInstruction::isVectorToScalar() const {
+  return getOpcode() == VPInstruction::ExtractFromEnd ||
+         getOpcode() == VPInstruction::ComputeReductionResult;
+}
+
 #if !defined(NDEBUG)
 bool VPInstruction::isFPMathOp() const {
   // Inspired by FPMathOperator::classof. Notable differences are that we don't
@@ -622,9 +624,7 @@ void VPInstruction::execute(VPTransformState &State) {
   State.setDebugLocFrom(getDebugLoc());
   bool GeneratesPerFirstLaneOnly =
       canGenerateScalarForFirstLane() &&
-      (vputils::onlyFirstLaneUsed(this) ||
-       getOpcode() == VPInstruction::ExtractFromEnd ||
-       getOpcode() == VPInstruction::ComputeReductionResult);
+      (vputils::onlyFirstLaneUsed(this) || isVectorToScalar());
   bool GeneratesPerAllLanes = doesGeneratePerAllLanes();
   for (unsigned Part = 0; Part < State.UF; ++Part) {
     if (GeneratesPerAllLanes) {