Consider vscale_range

Author: Mel-Chen

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -8566,8 +8566,13 @@ static void addExitUsersForFirstOrderRecurrences(
   VPBuilder ScalarPHBuilder(ScalarPHVPBB);
   VPBuilder MiddleBuilder(MiddleVPBB, MiddleVPBB->getFirstNonPhi());
 
-  auto IsScalableOne = [](ElementCount VF) -> bool {
-    return VF.isScalable() && VF.getKnownMinValue() == 1;
+  auto MayOneElement = [&Plan](ElementCount VF) -> bool {
+    if (!VF.isScalable() || VF.getKnownMinValue() != 1)
+      return false;
+
+    const Function *Fn = Plan.getScalarHeader()->getIRBasicBlock()->getParent();
+    const auto &Attr = Fn->getFnAttribute(Attribute::VScaleRange);
+    return !Attr.isValid() || Attr.getVScaleRangeMin() == 1;
   };
 
   for (auto &HeaderPhi : VectorRegion->getEntryBasicBlock()->phis()) {
@@ -8654,8 +8659,8 @@ static void addExitUsersForFirstOrderRecurrences(
       // For VF vscale x 1, if vscale = 1, we are unable to extract the
       // penultimate value of the recurrence. Instead, we can extract the last
       // element directly from VPInstruction::FirstOrderRecurrenceSplice.
-      // TODO: Consider vscale_range info and UF.
-      if (LoopVectorizationPlanner::getDecisionAndClampRange(IsScalableOne,
+      // TODO: Consider UF > 1.
+      if (LoopVectorizationPlanner::getDecisionAndClampRange(MayOneElement,
                                                              Range))
         return;
       VPValue *PenultimateElement = MiddleBuilder.createNaryOp(

llvm/lib/Transforms/Vectorize/VPlanHelpers.h

@@ -149,8 +149,7 @@ class VPLane {
   static VPLane getFirstLane() { return VPLane(0, VPLane::Kind::First); }
 
   static VPLane getLaneFromEnd(const ElementCount &VF, unsigned Offset) {
-    assert(Offset > 0 && Offset <= VF.getKnownMinValue() &&
-           "trying to extract with invalid offset");
+    assert(Offset > 0 && "trying to extract with invalid offset");
     unsigned LaneOffset = VF.getKnownMinValue() - Offset;
     Kind LaneKind;
     if (VF.isScalable())

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -715,8 +715,6 @@ Value *VPInstruction::generate(VPTransformState &State) {
     unsigned Offset = getOpcode() == VPInstruction::ExtractLastElement ? 1 : 2;
     Value *Res;
     if (State.VF.isVector()) {
-      assert(Offset <= State.VF.getKnownMinValue() &&
-             "invalid offset to extract from");
       // Extract lane VF - Offset from the operand.
       Res = State.get(getOperand(0), VPLane::getLaneFromEnd(State.VF, Offset));
     } else {