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

[Benjamin Maxwell via llvm-commits]

================
@@ -4530,3 +4531,84 @@ void VPlanTransforms::addExitUsersForFirstOrderRecurrences(VPlan &Plan,
     }
   }
 }
+
+void VPlanTransforms::convertFindLastRecurrences(
+    VPlan &Plan, VPRecipeBuilder &RecipeBuilder,
+    LoopVectorizationLegality *Legal) {
+  assert(Legal && "Need valid LoopVecLegality");
+
+  // 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:
+  //   ...the extract already exists, but needs updating...
+  //   result = extract-last-active new.data, new.mask, default.val
+
+  for (const auto &[Phi, RdxDesc] : Legal->getReductionVars()) {
+    if (RecurrenceDescriptor::isFindLastRecurrenceKind(
+            RdxDesc.getRecurrenceKind())) {
+      VPRecipeBase *PhiR = RecipeBuilder.getRecipe(Phi);
+      VPBuilder Builder = VPBuilder::getToInsertAfter(PhiR);
+
+      // Add mask phi...
+      VPValue *False =
+          Plan.getOrAddLiveIn(ConstantInt::getFalse(Phi->getContext()));
+      // FIXME: Either come up with a new phi recipe or make an existing one
+      //        more generic. There's only supposed to be one ALM PHI.
+      VPActiveLaneMaskPHIRecipe *MaskPHI =
+          new VPActiveLaneMaskPHIRecipe(False, DebugLoc());
+      Builder.insert(MaskPHI);
+
+      SelectInst *Select = cast<SelectInst>(RdxDesc.getLoopExitInstr());
+      auto *SR = cast<VPWidenSelectRecipe>(RecipeBuilder.getRecipe(Select));
+
+      // Add select for mask...
+      VPValue *Cond = SR->getCond();
+      Builder.setInsertPoint(SR);
+      VPInstruction *AnyOf = Builder.createNaryOp(VPInstruction::AnyOf, {Cond});
+      auto *MaskSelect = new VPWidenSelectVectorRecipe({AnyOf, Cond, MaskPHI});
+      Builder.insert(MaskSelect);
+      MaskPHI->addOperand(MaskSelect);
+
+      // Create new select for data...
+      auto *DataSelect = new VPWidenSelectVectorRecipe(
+          {AnyOf, SR->getOperand(1), SR->getOperand(2)});
+      Builder.insert(DataSelect);
+      SR->replaceAllUsesWith(DataSelect);
+      SR->eraseFromParent();
+
+      // FIXME: I initially tried to replace the reduction calculation in
+      //        the middle block here. However, it appears that the Def-Use
+      //        graph is somewhat broken. Calling replaceAllUsesWith on the
+      //        extract-from-end replaced the one use it knew about in
+      //        ir-bb<loop>, but not a second reference in ir-bb<exit.loopexit>
+      //        Both references were pending extra operands to phi nodes.
+
+      // I could either just replace the compute-reduction-result node alone,
+      // or (as I have done) plant the extract-last-active instead. But we
+      // now need to add the other operands to extract-last-active.
+      VPValue *Default = RecipeBuilder.getVPValueOrAddLiveIn(
+          RdxDesc.getRecurrenceStartValue());
+      for (VPUser *U : DataSelect->users()) {
+        VPInstruction *I = dyn_cast<VPInstruction>(U);
+        if (I && I->getOpcode() == VPInstruction::ExtractLastActive) {
+          I->addOperand(MaskSelect);
+          I->addOperand(Default);
+        }
+      }
----------------
MacDue wrote:

I was just trying to understand this `FIXME`. Is this what you meant here (replacing the `ComputeReductionResult` rather than modify a previously inserted `ExtractLastActive`)? I tried these changes (along with removing the change that adds the `ExtractLastActive`) and I didn't see any test failures. 
```suggestion
      VPInstruction *RdxResult = nullptr;
      for (VPUser *U : DataSelect->users()) {
        VPInstruction *I = dyn_cast<VPInstruction>(U);
        if (I && I->getOpcode() == VPInstruction::ComputeReductionResult) {
          RdxResult = I;
          break;
        }
      }

      assert(RdxResult);
      Builder.setInsertPoint(RdxResult);
      VPValue *Default = RecipeBuilder.getVPValueOrAddLiveIn(
          RdxDesc.getRecurrenceStartValue());
      auto *ExtractLastActive = Builder.createNaryOp(
          VPInstruction::ExtractLastActive, {DataSelect, MaskSelect, Default},
          RdxResult->getDebugLoc());
      RdxResult->replaceAllUsesWith(ExtractLastActive);
      RdxResult->eraseFromParent();
```

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