Step 3: Expand ExtractFromEnd to sub + extractelement

This transformation is merely to avoid potential issues.
The current folding by EVL is developed based on folding by mask. As a result,
ResumePhi should be dead code.

I just made this change so I can sleep better at night :P

Author: Mel-Chen

llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp

@@ -78,6 +78,7 @@ Type *VPTypeAnalysis::inferScalarTypeForRecipe(const VPInstruction *R) {
   case VPInstruction::CanonicalIVIncrementForPart:
   case VPInstruction::AnyOf:
     return SetResultTyFromOp();
+  case Instruction::ExtractElement:
   case VPInstruction::ExtractFirstActive:
   case VPInstruction::ExtractFromEnd: {
     Type *BaseTy = inferScalarType(R->getOperand(0));

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -478,6 +478,12 @@ Value *VPInstruction::generate(VPTransformState &State) {
     Value *Op2 = State.get(getOperand(2), OnlyFirstLaneUsed);
     return Builder.CreateSelect(Cond, Op1, Op2, 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), /*IsScalar*/ true);
+    return Builder.CreateExtractElement(Vec, Idx, Name);
+  }
   case VPInstruction::ActiveLaneMask: {
     // Get first lane of vector induction variable.
     Value *VIVElem0 = State.get(getOperand(0), VPLane(0));
@@ -752,7 +758,8 @@ InstructionCost VPInstruction::computeCost(ElementCount VF,
 }
 
 bool VPInstruction::isVectorToScalar() const {
-  return getOpcode() == VPInstruction::ExtractFromEnd ||
+  return getOpcode() == Instruction::ExtractElement ||
+         getOpcode() == VPInstruction::ExtractFromEnd ||
          getOpcode() == VPInstruction::ExtractFirstActive ||
          getOpcode() == VPInstruction::ComputeReductionResult ||
          getOpcode() == VPInstruction::AnyOf;
@@ -814,6 +821,7 @@ bool VPInstruction::opcodeMayReadOrWriteFromMemory() const {
   switch (getOpcode()) {
   case Instruction::ICmp:
   case Instruction::Select:
+  case Instruction::ExtractElement:
   case VPInstruction::AnyOf:
   case VPInstruction::CalculateTripCountMinusVF:
   case VPInstruction::CanonicalIVIncrementForPart:
@@ -851,6 +859,8 @@ bool VPInstruction::onlyFirstLaneUsed(const VPValue *Op) const {
   case VPInstruction::BranchOnCond:
   case VPInstruction::ResumePhi:
     return true;
+  case Instruction::ExtractElement:
+    return Op == getOperand(1);
   };
   llvm_unreachable("switch should return");
 }

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -1635,6 +1635,48 @@ void VPlanTransforms::addActiveLaneMask(
     HeaderMask->replaceAllUsesWith(LaneMask);
 }
 
+/// Adjust the way the resume value is obtained when using tail folding by EVL.
+/// Expanding ExtractFromEnd since the penultimate EVL could not equals to
+/// VFxUF. Expand
+///   %resume = ExtractFromEnd %vec, 1
+/// to
+///   %last.active.idx = sub %EVL, 1
+///   %resume = extractelement %vec, %last.active.idx
+static void adjustResumePhisForEVL(VPlan &Plan, VPValue &EVL) {
+  LLVMContext &Ctx = Plan.getCanonicalIV()->getScalarType()->getContext();
+  using namespace VPlanPatternMatch;
+  for (VPRecipeBase &R : *cast<VPBasicBlock>(Plan.getScalarPreheader())) {
+    VPValue *FromMiddleBlock;
+    if (!match(&R, m_VPInstruction<VPInstruction::ResumePhi>(
+                       m_VPValue(FromMiddleBlock), m_VPValue())))
+      continue;
+
+    VPValue *ExtractFrom;
+    if (match(FromMiddleBlock, m_VPInstruction<VPInstruction::ExtractFromEnd>(
+                                   m_VPValue(ExtractFrom), m_SpecificInt(1)))) {
+      // Skip if all elements are uniform.
+      if (vputils::isUniformAfterVectorization(ExtractFrom))
+        continue;
+      auto *ExtractR = cast<VPInstruction>(FromMiddleBlock);
+      VPBuilder Builder(ExtractR);
+      VPValue *OneVPV =
+          Plan.getOrAddLiveIn(ConstantInt::get(Type::getInt32Ty(Ctx), 1));
+      VPValue *LastActiveIdx =
+          Builder.createNaryOp(Instruction::Sub, {&EVL, OneVPV},
+                               ExtractR->getDebugLoc(), "last.active.idx");
+      VPValue *NewExtract = Builder.createNaryOp(
+          Instruction::ExtractElement, {ExtractFrom, LastActiveIdx},
+          ExtractR->getDebugLoc(), ExtractR->getName());
+      ExtractR->replaceAllUsesWith(NewExtract);
+      ExtractR->eraseFromParent();
+    }
+    assert((!dyn_cast<VPInstruction>(FromMiddleBlock) ||
+            cast<VPInstruction>(FromMiddleBlock)->getOpcode() !=
+                VPInstruction::ExtractFromEnd) &&
+           "Only extract the last lane for resumed values");
+  }
+}
+
 /// Try to convert \p CurRecipe to a corresponding EVL-based recipe. Returns
 /// nullptr if no EVL-based recipe could be created.
 /// \p HeaderMask  Header Mask.
@@ -1799,6 +1841,8 @@ static void transformRecipestoEVLRecipes(VPlan &Plan, VPValue &EVL) {
     }
   }
 
+  adjustResumePhisForEVL(Plan, EVL);
+
   for (VPRecipeBase *R : reverse(ToErase)) {
     SmallVector<VPValue *> PossiblyDead(R->operands());
     R->eraseFromParent();

llvm/lib/Transforms/Vectorize/VPlanVerifier.cpp

@@ -153,6 +153,9 @@ bool VPlanVerifier::verifyEVLRecipe(const VPInstruction &EVL) const {
         .Case<VPScalarCastRecipe>(
             [&](const VPScalarCastRecipe *S) { return VerifyEVLUse(*S, 0); })
         .Case<VPInstruction>([&](const VPInstruction *I) {
+          // Used by extracting the element at last active lane.
+          if (I->getOpcode() == Instruction::Sub)
+            return VerifyEVLUse(*I, 0);
           if (I->getOpcode() != Instruction::Add) {
             errs() << "EVL is used as an operand in non-VPInstruction::Add\n";
             return false;

llvm/test/Transforms/LoopVectorize/RISCV/vectorize-force-tail-with-evl-fixed-order-recurrence.ll

@@ -57,9 +57,7 @@ define void @first_order_recurrence(ptr noalias %A, ptr noalias %B, i64 %TC) {
 ; IF-EVL-NEXT:    [[TMP20:%.*]] = icmp eq i64 [[INDEX_NEXT]], [[N_VEC]]
 ; IF-EVL-NEXT:    br i1 [[TMP20]], label %[[MIDDLE_BLOCK:.*]], label %[[VECTOR_BODY]], !llvm.loop [[LOOP0:![0-9]+]]
 ; IF-EVL:       [[MIDDLE_BLOCK]]:
-; IF-EVL-NEXT:    [[TMP21:%.*]] = call i32 @llvm.vscale.i32()
-; IF-EVL-NEXT:    [[TMP22:%.*]] = mul i32 [[TMP21]], 4
-; IF-EVL-NEXT:    [[TMP23:%.*]] = sub i32 [[TMP22]], 1
+; IF-EVL-NEXT:    [[TMP23:%.*]] = sub i32 [[TMP12]], 1
 ; IF-EVL-NEXT:    [[VECTOR_RECUR_EXTRACT:%.*]] = extractelement <vscale x 4 x i32> [[VP_OP_LOAD]], i32 [[TMP23]]
 ; IF-EVL-NEXT:    br i1 true, label %[[FOR_END:.*]], label %[[SCALAR_PH]]
 ; IF-EVL:       [[SCALAR_PH]]:
@@ -210,24 +208,20 @@ define void @second_order_recurrence(ptr noalias %A, ptr noalias %B, i64 %TC) {
 ; IF-EVL-NEXT:    [[TMP24:%.*]] = icmp eq i64 [[INDEX_NEXT]], [[N_VEC]]
 ; IF-EVL-NEXT:    br i1 [[TMP24]], label %[[MIDDLE_BLOCK:.*]], label %[[VECTOR_BODY]], !llvm.loop [[LOOP4:![0-9]+]]
 ; IF-EVL:       [[MIDDLE_BLOCK]]:
-; IF-EVL-NEXT:    [[TMP25:%.*]] = call i32 @llvm.vscale.i32()
-; IF-EVL-NEXT:    [[TMP26:%.*]] = mul i32 [[TMP25]], 4
-; IF-EVL-NEXT:    [[TMP27:%.*]] = sub i32 [[TMP26]], 1
+; IF-EVL-NEXT:    [[TMP27:%.*]] = sub i32 [[TMP15]], 1
 ; IF-EVL-NEXT:    [[VECTOR_RECUR_EXTRACT:%.*]] = extractelement <vscale x 4 x i32> [[VP_OP_LOAD]], i32 [[TMP27]]
-; IF-EVL-NEXT:    [[TMP28:%.*]] = call i32 @llvm.vscale.i32()
-; IF-EVL-NEXT:    [[TMP29:%.*]] = mul i32 [[TMP28]], 4
-; IF-EVL-NEXT:    [[TMP30:%.*]] = sub i32 [[TMP29]], 1
+; IF-EVL-NEXT:    [[TMP30:%.*]] = sub i32 [[TMP15]], 1
 ; IF-EVL-NEXT:    [[VECTOR_RECUR_EXTRACT3:%.*]] = extractelement <vscale x 4 x i32> [[TMP19]], i32 [[TMP30]]
 ; IF-EVL-NEXT:    br i1 true, label %[[FOR_END:.*]], label %[[SCALAR_PH]]
 ; IF-EVL:       [[SCALAR_PH]]:
 ; IF-EVL-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i64 [ [[N_VEC]], %[[MIDDLE_BLOCK]] ], [ 0, %[[ENTRY]] ]
 ; IF-EVL-NEXT:    [[SCALAR_RECUR_INIT:%.*]] = phi i32 [ [[VECTOR_RECUR_EXTRACT]], %[[MIDDLE_BLOCK]] ], [ 33, %[[ENTRY]] ]
-; IF-EVL-NEXT:    [[SCALAR_RECUR_INIT4:%.*]] = phi i32 [ [[VECTOR_RECUR_EXTRACT3]], %[[MIDDLE_BLOCK]] ], [ 22, %[[ENTRY]] ]
+; IF-EVL-NEXT:    [[SCALAR_RECUR_INIT5:%.*]] = phi i32 [ [[VECTOR_RECUR_EXTRACT3]], %[[MIDDLE_BLOCK]] ], [ 22, %[[ENTRY]] ]
 ; IF-EVL-NEXT:    br label %[[FOR_BODY:.*]]
 ; IF-EVL:       [[FOR_BODY]]:
 ; IF-EVL-NEXT:    [[INDVARS:%.*]] = phi i64 [ [[BC_RESUME_VAL]], %[[SCALAR_PH]] ], [ [[INDVARS_NEXT:%.*]], %[[FOR_BODY]] ]
 ; IF-EVL-NEXT:    [[FOR1:%.*]] = phi i32 [ [[SCALAR_RECUR_INIT]], %[[SCALAR_PH]] ], [ [[TMP31:%.*]], %[[FOR_BODY]] ]
-; IF-EVL-NEXT:    [[FOR2:%.*]] = phi i32 [ [[SCALAR_RECUR_INIT4]], %[[SCALAR_PH]] ], [ [[FOR1]], %[[FOR_BODY]] ]
+; IF-EVL-NEXT:    [[FOR2:%.*]] = phi i32 [ [[SCALAR_RECUR_INIT5]], %[[SCALAR_PH]] ], [ [[FOR1]], %[[FOR_BODY]] ]
 ; IF-EVL-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds nuw i32, ptr [[A]], i64 [[INDVARS]]
 ; IF-EVL-NEXT:    [[TMP31]] = load i32, ptr [[ARRAYIDX]], align 4
 ; IF-EVL-NEXT:    [[ADD:%.*]] = add nsw i32 [[FOR1]], [[FOR2]]
@@ -389,30 +383,24 @@ define void @third_order_recurrence(ptr noalias %A, ptr noalias %B, i64 %TC) {
 ; IF-EVL-NEXT:    [[TMP28:%.*]] = icmp eq i64 [[INDEX_NEXT]], [[N_VEC]]
 ; IF-EVL-NEXT:    br i1 [[TMP28]], label %[[MIDDLE_BLOCK:.*]], label %[[VECTOR_BODY]], !llvm.loop [[LOOP6:![0-9]+]]
 ; IF-EVL:       [[MIDDLE_BLOCK]]:
-; IF-EVL-NEXT:    [[TMP29:%.*]] = call i32 @llvm.vscale.i32()
-; IF-EVL-NEXT:    [[TMP30:%.*]] = mul i32 [[TMP29]], 4
-; IF-EVL-NEXT:    [[TMP31:%.*]] = sub i32 [[TMP30]], 1
+; IF-EVL-NEXT:    [[TMP31:%.*]] = sub i32 [[TMP18]], 1
 ; IF-EVL-NEXT:    [[VECTOR_RECUR_EXTRACT:%.*]] = extractelement <vscale x 4 x i32> [[VP_OP_LOAD]], i32 [[TMP31]]
-; IF-EVL-NEXT:    [[TMP32:%.*]] = call i32 @llvm.vscale.i32()
-; IF-EVL-NEXT:    [[TMP33:%.*]] = mul i32 [[TMP32]], 4
-; IF-EVL-NEXT:    [[TMP34:%.*]] = sub i32 [[TMP33]], 1
+; IF-EVL-NEXT:    [[TMP34:%.*]] = sub i32 [[TMP18]], 1
 ; IF-EVL-NEXT:    [[VECTOR_RECUR_EXTRACT6:%.*]] = extractelement <vscale x 4 x i32> [[TMP22]], i32 [[TMP34]]
-; IF-EVL-NEXT:    [[TMP35:%.*]] = call i32 @llvm.vscale.i32()
-; IF-EVL-NEXT:    [[TMP36:%.*]] = mul i32 [[TMP35]], 4
-; IF-EVL-NEXT:    [[TMP37:%.*]] = sub i32 [[TMP36]], 1
+; IF-EVL-NEXT:    [[TMP37:%.*]] = sub i32 [[TMP18]], 1
 ; IF-EVL-NEXT:    [[VECTOR_RECUR_EXTRACT7:%.*]] = extractelement <vscale x 4 x i32> [[TMP23]], i32 [[TMP37]]
 ; IF-EVL-NEXT:    br i1 true, label %[[FOR_END:.*]], label %[[SCALAR_PH]]
 ; IF-EVL:       [[SCALAR_PH]]:
 ; IF-EVL-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i64 [ [[N_VEC]], %[[MIDDLE_BLOCK]] ], [ 0, %[[ENTRY]] ]
 ; IF-EVL-NEXT:    [[SCALAR_RECUR_INIT:%.*]] = phi i32 [ [[VECTOR_RECUR_EXTRACT]], %[[MIDDLE_BLOCK]] ], [ 33, %[[ENTRY]] ]
-; IF-EVL-NEXT:    [[SCALAR_RECUR_INIT8:%.*]] = phi i32 [ [[VECTOR_RECUR_EXTRACT6]], %[[MIDDLE_BLOCK]] ], [ 22, %[[ENTRY]] ]
-; IF-EVL-NEXT:    [[SCALAR_RECUR_INIT9:%.*]] = phi i32 [ [[VECTOR_RECUR_EXTRACT7]], %[[MIDDLE_BLOCK]] ], [ 11, %[[ENTRY]] ]
+; IF-EVL-NEXT:    [[SCALAR_RECUR_INIT10:%.*]] = phi i32 [ [[VECTOR_RECUR_EXTRACT6]], %[[MIDDLE_BLOCK]] ], [ 22, %[[ENTRY]] ]
+; IF-EVL-NEXT:    [[SCALAR_RECUR_INIT11:%.*]] = phi i32 [ [[VECTOR_RECUR_EXTRACT7]], %[[MIDDLE_BLOCK]] ], [ 11, %[[ENTRY]] ]
 ; IF-EVL-NEXT:    br label %[[FOR_BODY:.*]]
 ; IF-EVL:       [[FOR_BODY]]:
 ; IF-EVL-NEXT:    [[INDVARS:%.*]] = phi i64 [ [[BC_RESUME_VAL]], %[[SCALAR_PH]] ], [ [[INDVARS_NEXT:%.*]], %[[FOR_BODY]] ]
 ; IF-EVL-NEXT:    [[FOR1:%.*]] = phi i32 [ [[SCALAR_RECUR_INIT]], %[[SCALAR_PH]] ], [ [[TMP38:%.*]], %[[FOR_BODY]] ]
-; IF-EVL-NEXT:    [[FOR2:%.*]] = phi i32 [ [[SCALAR_RECUR_INIT8]], %[[SCALAR_PH]] ], [ [[FOR1]], %[[FOR_BODY]] ]
-; IF-EVL-NEXT:    [[FOR3:%.*]] = phi i32 [ [[SCALAR_RECUR_INIT9]], %[[SCALAR_PH]] ], [ [[FOR2]], %[[FOR_BODY]] ]
+; IF-EVL-NEXT:    [[FOR2:%.*]] = phi i32 [ [[SCALAR_RECUR_INIT10]], %[[SCALAR_PH]] ], [ [[FOR1]], %[[FOR_BODY]] ]
+; IF-EVL-NEXT:    [[FOR3:%.*]] = phi i32 [ [[SCALAR_RECUR_INIT11]], %[[SCALAR_PH]] ], [ [[FOR2]], %[[FOR_BODY]] ]
 ; IF-EVL-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds nuw i32, ptr [[A]], i64 [[INDVARS]]
 ; IF-EVL-NEXT:    [[TMP38]] = load i32, ptr [[ARRAYIDX]], align 4
 ; IF-EVL-NEXT:    [[ADD:%.*]] = add nsw i32 [[FOR2]], [[FOR3]]

llvm/test/Transforms/LoopVectorize/RISCV/vplan-vp-intrinsics-fixed-order-recurrence.ll

@@ -48,9 +48,15 @@ define void @first_order_recurrence(ptr noalias %A, ptr noalias %B, i64 %TC) {
 ; IF-EVL-NEXT: Successor(s): middle.block
 ; IF-EVL-EMPTY:
 ; IF-EVL: middle.block:
-; IF-EVL-NEXT:   EMIT vp<[[RESUME_EXTRACT:%.+]]> = extract-from-end ir<[[LD]]>, ir<1>
+; IF-EVL-NEXT:   EMIT vp<[[LAST_ACTIVE_IDX:%.+]]> = sub vp<[[EVL]]>, ir<1>
+; IF-EVL-NEXT:   EMIT vp<[[RESUME_EXTRACT:%.+]]> = extractelement ir<[[LD]]>, vp<[[LAST_ACTIVE_IDX]]>
 ; IF-EVL-NEXT:   EMIT branch-on-cond ir<true>
 ; IF-EVL-NEXT: Successor(s): ir-bb<for.end>, scalar.ph
+; IF-EVL-EMPTY:
+; IF-EVL: scalar.ph:
+; IF-EVL-NEXT:   EMIT vp<%bc.resume.val> = resume-phi vp<[[VTC]]>, ir<0>
+; IF-EVL-NEXT:   EMIT vp<%scalar.recur.init> = resume-phi vp<[[RESUME_EXTRACT]]>, ir<33>
+; IF-EVL-NEXT: Successor(s): ir-bb<for.body>
 
 entry:
   br label %for.body