[LV] Build VPlan for CSA

Author: michaelmaitland

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -180,6 +180,8 @@ const char LLVMLoopVectorizeFollowupEpilogue[] =
 STATISTIC(LoopsVectorized, "Number of loops vectorized");
 STATISTIC(LoopsAnalyzed, "Number of loops analyzed for vectorization");
 STATISTIC(LoopsEpilogueVectorized, "Number of epilogues vectorized");
+STATISTIC(CSAsVectorized,
+          "Number of conditional scalar assignments vectorized");
 
 static cl::opt<bool> EnableEpilogueVectorization(
     "enable-epilogue-vectorization", cl::init(true), cl::Hidden,
@@ -500,6 +502,10 @@ class InnerLoopVectorizer {
   virtual std::pair<BasicBlock *, Value *>
   createVectorizedLoopSkeleton(const SCEV2ValueTy &ExpandedSCEVs);
 
+  /// For all vectorized CSAs, replace uses of live-out scalar from the orignal
+  /// loop with the extracted scalar from the vector loop for.
+  void fixCSALiveOuts(VPTransformState &State, VPlan &Plan);
+
   /// Fix the vectorized code, taking care of header phi's, live-outs, and more.
   void fixVectorizedLoop(VPTransformState &State, VPlan &Plan);
 
@@ -2932,6 +2938,25 @@ LoopVectorizationCostModel::getVectorIntrinsicCost(CallInst *CI,
                                    TargetTransformInfo::TCK_RecipThroughput);
 }
 
+void InnerLoopVectorizer::fixCSALiveOuts(VPTransformState &State, VPlan &Plan) {
+  for (const auto &CSA : Plan.getCSAStates()) {
+    VPCSADataUpdateRecipe *VPDataUpdate = CSA.second->getDataUpdate();
+    assert(VPDataUpdate &&
+           "VPDataUpdate must have been introduced prior to fixing live outs");
+    Value *V = VPDataUpdate->getUnderlyingValue();
+    Value *ExtractedScalar = State.get(CSA.second->getExtractScalarRecipe(), 0,
+                                       /*NeedsScalar=*/true);
+    // Fix LCSSAPhis
+    llvm::SmallPtrSet<PHINode *, 2> ToFix;
+    for (User *U : V->users())
+      if (auto *Phi = dyn_cast<PHINode>(U);
+          Phi && Phi->getParent() == LoopExitBlock)
+        ToFix.insert(Phi);
+    for (PHINode *Phi : ToFix)
+      Phi->addIncoming(ExtractedScalar, LoopMiddleBlock);
+  }
+}
+
 void InnerLoopVectorizer::fixVectorizedLoop(VPTransformState &State,
                                             VPlan &Plan) {
   // Fix widened non-induction PHIs by setting up the PHI operands.
@@ -2972,6 +2997,8 @@ void InnerLoopVectorizer::fixVectorizedLoop(VPTransformState &State,
                    getOrCreateVectorTripCount(VectorLoop->getLoopPreheader()),
                    IVEndValues[Entry.first], LoopMiddleBlock,
                    VectorLoop->getHeader(), Plan, State);
+
+    fixCSALiveOuts(State, Plan);
   }
 
   // Fix live-out phis not already fixed earlier.
@@ -4482,6 +4509,9 @@ static bool willGenerateVectors(VPlan &Plan, ElementCount VF,
       case VPDef::VPEVLBasedIVPHISC:
       case VPDef::VPPredInstPHISC:
       case VPDef::VPBranchOnMaskSC:
+      case VPRecipeBase::VPCSADataUpdateSC:
+      case VPRecipeBase::VPCSAExtractScalarSC:
+      case VPRecipeBase::VPCSAHeaderPHISC:
         continue;
       case VPDef::VPReductionSC:
       case VPDef::VPActiveLaneMaskPHISC:
@@ -8508,9 +8538,6 @@ VPRecipeBuilder::tryToCreateWidenRecipe(Instruction *Instr,
       return Recipe;
 
     VPHeaderPHIRecipe *PhiRecipe = nullptr;
-    assert((Legal->isReductionVariable(Phi) ||
-            Legal->isFixedOrderRecurrence(Phi)) &&
-           "can only widen reductions and fixed-order recurrences here");
     VPValue *StartV = Operands[0];
     if (Legal->isReductionVariable(Phi)) {
       const RecurrenceDescriptor &RdxDesc =
@@ -8520,12 +8547,23 @@ VPRecipeBuilder::tryToCreateWidenRecipe(Instruction *Instr,
       PhiRecipe = new VPReductionPHIRecipe(Phi, RdxDesc, *StartV,
                                            CM.isInLoopReduction(Phi),
                                            CM.useOrderedReductions(RdxDesc));
-    } else {
+    } else if (Legal->isFixedOrderRecurrence(Phi)) {
       // TODO: Currently fixed-order recurrences are modeled as chains of
       // first-order recurrences. If there are no users of the intermediate
       // recurrences in the chain, the fixed order recurrence should be modeled
       // directly, enabling more efficient codegen.
       PhiRecipe = new VPFirstOrderRecurrencePHIRecipe(Phi, *StartV);
+    } else if (Legal->isCSAPhi(Phi)) {
+      VPCSAState *State = Plan.getCSAStates().find(Phi)->second;
+      VPValue *InitData = State->getVPInitData();
+      // When the VF=getFixed(1), InitData is just InitScalar.
+      if (!InitData)
+        InitData = State->getVPInitScalar();
+      PhiRecipe = new VPCSAHeaderPHIRecipe(Phi, InitData);
+      State->setPhiRecipe(cast<VPCSAHeaderPHIRecipe>(PhiRecipe));
+    } else {
+      llvm_unreachable(
+          "can only widen reductions, fixed-order recurrences, and CSAs here");
     }
 
     PhisToFix.push_back(PhiRecipe);
@@ -8555,6 +8593,19 @@ VPRecipeBuilder::tryToCreateWidenRecipe(Instruction *Instr,
                                 make_range(Operands.begin(), Operands.end()));
 
   if (auto *SI = dyn_cast<SelectInst>(Instr)) {
+    auto *CSADescIt = find_if(Legal->getCSAs(), [&](auto CSA) {
+      return CSADescriptor::isCSASelect(CSA.second, SI);
+    });
+    if (CSADescIt != Legal->getCSAs().end()) {
+      PHINode *CSAPhi = CSADescIt->first;
+      VPCSAState *State = Plan.getCSAStates().find(CSAPhi)->second;
+      VPValue *VPDataPhi = State->getPhiRecipe();
+      auto *R = new VPCSADataUpdateRecipe(
+          SI, {VPDataPhi, Operands[0], Operands[1], Operands[2]});
+      State->setDataUpdate(R);
+      return R;
+    }
+
     return new VPWidenSelectRecipe(
         *SI, make_range(Operands.begin(), Operands.end()));
   }
@@ -8567,6 +8618,107 @@ VPRecipeBuilder::tryToCreateWidenRecipe(Instruction *Instr,
   return tryToWiden(Instr, Operands, VPBB);
 }
 
+/// Add CSA Recipes that can occur before each instruction in the input IR
+/// is processed and introduced into VPlan.
+static void
+addCSAPreprocessRecipes(const LoopVectorizationLegality::CSAList &CSAs,
+                        Loop *OrigLoop, VPBasicBlock *PreheaderVPBB,
+                        VPBasicBlock *HeaderVPBB, DebugLoc DL, VFRange &Range,
+                        VPlan &Plan) {
+
+  // Don't build full CSA for VF=ElementCount::getFixed(1)
+  bool IsScalarVF = LoopVectorizationPlanner::getDecisionAndClampRange(
+      [&](ElementCount VF) { return VF.isScalar(); }, Range);
+
+  for (const auto &CSA : CSAs) {
+    VPValue *VPInitScalar = Plan.getOrAddLiveIn(
+        CSA.first->getIncomingValueForBlock(OrigLoop->getLoopPreheader()));
+
+    // Scalar VF builds the scalar version of the loop. In that case,
+    // no maintenence of mask nor extraction in middle block is needed.
+    if (IsScalarVF) {
+      VPCSAState *S = new VPCSAState(VPInitScalar);
+      Plan.addCSAState(CSA.first, S);
+      continue;
+    }
+
+    auto *VPInitMask =
+        new VPInstruction(VPInstruction::CSAInitMask, {}, DL, "csa.init.mask");
+    auto *VPInitData = new VPInstruction(VPInstruction::CSAInitData,
+                                         {VPInitScalar}, DL, "csa.init.data");
+    PreheaderVPBB->appendRecipe(VPInitMask);
+    PreheaderVPBB->appendRecipe(VPInitData);
+
+    auto *VPMaskPhi = new VPInstruction(VPInstruction::CSAMaskPhi, {VPInitMask},
+                                        DL, "csa.mask.phi");
+    HeaderVPBB->appendRecipe(VPMaskPhi);
+
+    auto *S = new VPCSAState(VPInitScalar, VPInitData, VPMaskPhi);
+    Plan.addCSAState(CSA.first, S);
+  }
+}
+
+/// Add CSA Recipes that must occur after each instruction in the input IR
+/// is processed and introduced into VPlan.
+static void
+addCSAPostprocessRecipes(VPRecipeBuilder &RecipeBuilder,
+                         const LoopVectorizationLegality::CSAList &CSAs,
+                         VPBasicBlock *MiddleVPBB, DebugLoc DL, VFRange &Range,
+                         VPlan &Plan) {
+  // Don't build CSA for VF=ElementCount::getFixed(1)
+  if (LoopVectorizationPlanner::getDecisionAndClampRange(
+          [&](ElementCount VF) { return VF.isScalar(); }, Range))
+    return;
+
+  for (const auto &CSA : CSAs) {
+    VPCSAState *CSAState = Plan.getCSAStates().find(CSA.first)->second;
+    VPCSADataUpdateRecipe *VPDataUpdate = CSAState->getDataUpdate();
+
+    assert(VPDataUpdate &&
+           "VPDataUpdate must have been introduced prior to postprocess");
+    assert(CSA.second.getCond() &&
+           "CSADescriptor must know how to describe the condition");
+    auto GetVPValue = [&](Value *I) {
+      return RecipeBuilder.getRecipe(cast<Instruction>(I))->getVPSingleValue();
+    };
+    VPValue *WidenedCond = GetVPValue(CSA.second.getCond());
+    VPValue *VPInitScalar = CSAState->getVPInitScalar();
+
+    // The CSA optimization wants to use a condition such that when it is
+    // true, a new value is assigned. However, it is possible that a true lane
+    // in WidenedCond corresponds to selection of the initial value instead.
+    // In that case, we must use the negation of WidenedCond.
+    // i.e. select cond new_val old_val versus select cond.not old_val new_val
+    VPValue *CondToUse = WidenedCond;
+    if (cast<SelectInst>(CSA.second.getAssignment())->getTrueValue() ==
+        CSA.first) {
+      auto *VPNotCond = new VPInstruction(VPInstruction::Not, WidenedCond, DL);
+      VPNotCond->insertBefore(
+          GetVPValue(CSA.second.getAssignment())->getDefiningRecipe());
+      CondToUse = VPNotCond;
+    }
+
+    auto *VPAnyActive = new VPInstruction(
+        VPInstruction::CSAAnyActive, {CondToUse}, DL, "csa.cond.anyactive");
+    VPAnyActive->insertBefore(
+        GetVPValue(CSA.second.getAssignment())->getDefiningRecipe());
+
+    auto *VPMaskSel = new VPInstruction(
+        VPInstruction::CSAMaskSel,
+        {CondToUse, CSAState->getVPMaskPhi(), VPAnyActive}, DL, "csa.mask.sel");
+    VPMaskSel->insertAfter(VPAnyActive);
+    VPDataUpdate->setVPNewMaskAndVPAnyActive(VPMaskSel, VPAnyActive);
+    VPCSAExtractScalarRecipe *ExtractScalarRecipe =
+        new VPCSAExtractScalarRecipe({VPInitScalar, VPMaskSel, VPDataUpdate});
+
+    MiddleVPBB->insert(ExtractScalarRecipe, MiddleVPBB->getFirstNonPhi());
+
+    // Update CSAState with new recipes
+    CSAState->setExtractScalarRecipe(ExtractScalarRecipe);
+    CSAState->setVPAnyActive(VPAnyActive);
+  }
+}
+
 void LoopVectorizationPlanner::buildVPlansWithVPRecipes(ElementCount MinVF,
                                                         ElementCount MaxVF) {
   assert(OrigLoop->isInnermost() && "Inner loop expected.");
@@ -8623,7 +8775,8 @@ static void addCanonicalIVRecipes(VPlan &Plan, Type *IdxTy, bool HasNUW,
 // VPWidenPointerInductionRecipe and induction increments.
 static MapVector<PHINode *, VPValue *> collectUsersInExitBlock(
     Loop *OrigLoop, VPRecipeBuilder &Builder, VPlan &Plan,
-    const MapVector<PHINode *, InductionDescriptor> &Inductions) {
+    const MapVector<PHINode *, InductionDescriptor> &Inductions,
+    const MapVector<PHINode *, CSADescriptor> &CSAs) {
   auto MiddleVPBB =
       cast<VPBasicBlock>(Plan.getVectorLoopRegion()->getSingleSuccessor());
   // No edge from the middle block to the unique exit block has been inserted
@@ -8652,6 +8805,17 @@ static MapVector<PHINode *, VPValue *> collectUsersInExitBlock(
            return P && Inductions.contains(P);
          })))
       continue;
+    // Exit values for CSAs are computed and updated outside of VPlan and
+    // independent of induction recipes.
+    // TODO: Compute induction exit values in VPlan, use VPLiveOuts to update
+    // live-outs.
+    if (isa<VPCSADataUpdateRecipe>(V) &&
+        (isa<Instruction>(IncomingValue) &&
+         any_of(IncomingValue->users(), [&CSAs](User *U) {
+           auto *P = dyn_cast<PHINode>(U);
+           return P && CSAs.contains(P);
+         })))
+      continue;
     ExitingValuesToFix.insert({&ExitPhi, V});
   }
   return ExitingValuesToFix;
@@ -8893,6 +9057,10 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
   bool HasNUW = Style == TailFoldingStyle::None;
   addCanonicalIVRecipes(*Plan, Legal->getWidestInductionType(), HasNUW, DL);
 
+  addCSAPreprocessRecipes(Legal->getCSAs(), OrigLoop, Plan->getPreheader(),
+                          Plan->getVectorLoopRegion()->getEntryBasicBlock(), DL,
+                          Range, *Plan);
+
   VPRecipeBuilder RecipeBuilder(*Plan, OrigLoop, TLI, Legal, CM, PSE, Builder);
 
   // ---------------------------------------------------------------------------
@@ -8999,6 +9167,11 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
     VPBB = cast<VPBasicBlock>(VPBB->getSingleSuccessor());
   }
 
+  VPBasicBlock *MiddleVPBB =
+      cast<VPBasicBlock>(Plan->getVectorLoopRegion()->getSingleSuccessor());
+  addCSAPostprocessRecipes(RecipeBuilder, Legal->getCSAs(), MiddleVPBB, DL,
+                           Range, *Plan);
+
   // After here, VPBB should not be used.
   VPBB = nullptr;
 
@@ -9008,8 +9181,9 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
          "VPBasicBlock");
   RecipeBuilder.fixHeaderPhis();
 
-  MapVector<PHINode *, VPValue *> ExitingValuesToFix = collectUsersInExitBlock(
-      OrigLoop, RecipeBuilder, *Plan, Legal->getInductionVars());
+  MapVector<PHINode *, VPValue *> ExitingValuesToFix =
+      collectUsersInExitBlock(OrigLoop, RecipeBuilder, *Plan,
+                              Legal->getInductionVars(), Legal->getCSAs());
 
   addLiveOutsForFirstOrderRecurrences(*Plan, ExitingValuesToFix);
   addUsersInExitBlock(*Plan, ExitingValuesToFix);
@@ -10106,6 +10280,7 @@ bool LoopVectorizePass::processLoop(Loop *L) {
         const auto &[ExpandedSCEVs, ReductionResumeValues] = LVP.executePlan(
             EPI.MainLoopVF, EPI.MainLoopUF, *BestMainPlan, MainILV, DT, true);
         ++LoopsVectorized;
+        CSAsVectorized += LVL.getCSAs().size();
 
         // Second pass vectorizes the epilogue and adjusts the control flow
         // edges from the first pass.
@@ -10198,6 +10373,7 @@ bool LoopVectorizePass::processLoop(Loop *L) {
                                PSI, Checks);
         LVP.executePlan(VF.Width, IC, BestPlan, LB, DT, false);
         ++LoopsVectorized;
+        CSAsVectorized += LVL.getCSAs().size();
 
         // Add metadata to disable runtime unrolling a scalar loop when there
         // are no runtime checks about strides and memory. A scalar loop that is

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -216,7 +216,7 @@ void VPBlockBase::deleteCFG(VPBlockBase *Entry) {
 
 VPBasicBlock::iterator VPBasicBlock::getFirstNonPhi() {
   iterator It = begin();
-  while (It != end() && It->isPhi())
+  while (It != end() && vputils::isPhi(*It))
     It++;
   return It;
 }
@@ -858,6 +858,9 @@ VPlan::~VPlan() {
     delete VPV;
   if (BackedgeTakenCount)
     delete BackedgeTakenCount;
+
+  for (std::pair<PHINode *, VPCSAState *> &S : CSAStates)
+    delete S.second;
 }
 
 VPlanPtr VPlan::createInitialVPlan(const SCEV *TripCount, ScalarEvolution &SE,
@@ -1032,7 +1035,7 @@ void VPlan::execute(VPTransformState *State) {
   VPBasicBlock *Header = getVectorLoopRegion()->getEntryBasicBlock();
   for (VPRecipeBase &R : Header->phis()) {
     // Skip phi-like recipes that generate their backedege values themselves.
-    if (isa<VPWidenPHIRecipe>(&R))
+    if (vputils::isPhiThatGeneratesBackedge(R))
       continue;
 
     if (isa<VPWidenPointerInductionRecipe>(&R) ||