[llvm] [LV] Vectorize conditional scalar assignments (PR #158088)

[Florian Hahn via llvm-commits]

================
@@ -4722,3 +4722,75 @@ void VPlanTransforms::addExitUsersForFirstOrderRecurrences(VPlan &Plan,
     }
   }
 }
+
+void VPlanTransforms::convertFindLastRecurrences(
+    VPlan &Plan, VPRecipeBuilder &RecipeBuilder) {
+
+  // May need to do something better than this?
+  if (Plan.hasScalarVFOnly())
+    return;
+
+  // We want to create the following nodes:
+  // vec.body:
+  //   mask.phi = phi <VF x i1> [ all.false, vec.ph ], [ new.mask, vec.body ]
+  //   ...data.phi already exists, but needs updating...
+  //   data.phi = phi <VF x Ty> [ default.val, vec.ph ], [ new.data, vec.body ]
+  //
+  //   ...'data' and 'compare' created by existing nodes...
+  //
+  //   any_active = i1 any_of_reduction(compare)
+  //   new.mask = select any_active, compare, mask.phi
+  //   new.data = select any_active, data, data.phi
+  //
+  // middle.block:
+  //   result = extract-last-active new.data, new.mask, default.val
+
+  for (auto &Phi : Plan.getVectorLoopRegion()->getEntryBasicBlock()->phis()) {
+    auto *PhiR = dyn_cast<VPReductionPHIRecipe>(&Phi);
+    if (!PhiR || !RecurrenceDescriptor::isFindLastRecurrenceKind(
+                     PhiR->getRecurrenceKind()))
+      continue;
+
+    // Find the condition for the select
+    auto *SR = dyn_cast<VPWidenSelectRecipe>(&PhiR->getBackedgeRecipe());
+    if (!SR)
+      continue;
+    VPValue *Cond = SR->getCond();
+
+    // Add mask phi
+    VPBuilder Builder = VPBuilder::getToInsertAfter(PhiR);
+    VPValue *False = Plan.getOrAddLiveIn(
+        ConstantInt::getFalse(PhiR->getUnderlyingValue()->getContext()));
+    auto *MaskPHI = new VPWidenPHIRecipe(nullptr, False, DebugLoc());
+    Builder.insert(MaskPHI);
+
+    // Add select for mask
+    Builder.setInsertPoint(SR);
+    VPValue *AnyOf = Builder.createNaryOp(VPInstruction::AnyOf, {Cond});
+    VPValue *MaskSelect = Builder.createSelect(AnyOf, Cond, MaskPHI);
+    MaskPHI->addOperand(MaskSelect);
+
+    // Replace select for data
+    VPValue *DataSelect = Builder.createSelect(
+        AnyOf, SR->getOperand(1), SR->getOperand(2), SR->getDebugLoc());
+    SR->replaceAllUsesWith(DataSelect);
+    SR->eraseFromParent();
+
+    // Find final reduction computation and replace it with an
+    // extract.last.active intrinsic.
+    VPUser **ComputeRdx = find_if(DataSelect->users(), [](VPUser *U) {
+      VPInstruction *I = dyn_cast<VPInstruction>(U);
+      return I && I->getOpcode() == VPInstruction::ComputeReductionResult;
+    });
----------------
fhahn wrote:

In VPlanConstruction.cpp there should be a `findComputeReductionResult` helper, could you use that one here as well after moving?

https://github.com/llvm/llvm-project/pull/158088