[llvm] [VPlan] Add transform to replace VPWidenCanonicalIV with wide IV. (PR #194267)

[Florian Hahn via llvm-commits]

================
@@ -691,6 +691,73 @@ static void removeRedundantCanonicalIVs(VPlan &Plan) {
   }
 }
 
+void VPlanTransforms::replaceWidenCanonicalIVWithWidenIV(
+    VPlan &Plan, ScalarEvolution &SE, const TargetTransformInfo &TTI,
+    TargetTransformInfo::TargetCostKind CostKind, ElementCount VF,
+    unsigned UF) {
+  VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
+  if (!LoopRegion || Plan.hasScalarVFOnly())
+    return;
+
+  VPValue *CanonicalIV = LoopRegion->getCanonicalIV();
+  auto *WideCanIV = vputils::findUserOf<VPWidenCanonicalIVRecipe>(CanonicalIV);
+  if (!WideCanIV)
+    return;
+
+  // If a canonical VPWidenIntOrFpInductionRecipe already exists, reuse it.
+  VPBasicBlock *Header = LoopRegion->getEntryBasicBlock();
+  for (VPRecipeBase &Phi : Header->phis()) {
+    auto *WideIV = dyn_cast<VPWidenIntOrFpInductionRecipe>(&Phi);
+    if (!WideIV || !WideIV->isCanonical())
+      continue;
+    // The canonical wide IV is used to compute the header mask, hence all
+    // lanes will be used. Drop poison-generating flags.
+    WideIV->dropPoisonGeneratingFlags();
+    WideCanIV->replaceAllUsesWith(WideIV);
+    WideCanIV->eraseFromParent();
+    return;
+  }
+
+  // Otherwise, introduce a new VPWidenIntOrFpInductionRecipe if profitable.
+  if (vputils::onlyFirstLaneUsed(WideCanIV) ||
+      vputils::onlyScalarValuesUsed(WideCanIV))
+    return;
+
+  Type *CanIVTy = LoopRegion->getCanonicalIVType();
+  auto *VecTy = VectorType::get(CanIVTy, VF);
+  InstructionCost BroadcastCost = TTI.getShuffleCost(
+      TargetTransformInfo::SK_Broadcast, VecTy, VecTy, {}, CostKind);
+  InstructionCost PHICost = TTI.getCFInstrCost(Instruction::PHI, CostKind);
+  if (PHICost > BroadcastCost)
+    return;
+
+  // Bail out if the wide induction phis increase the expected spill cost.
+  SmallPtrSet<const Value *, 1> ValuesToIgnore;
+  VPRegisterUsage UnrolledBase =
+      calculateRegisterUsageForPlan(Plan, {VF}, TTI, ValuesToIgnore)[0];
+  for (auto &Pair : UnrolledBase.MaxLocalUsers)
+    Pair.second *= UF;
+  unsigned RegClass = TTI.getRegisterClassForType(/*Vector=*/true, VecTy);
+  VPRegisterUsage Projected = UnrolledBase;
+  Projected.MaxLocalUsers[RegClass] += UF;
+  if (Projected.spillCost(TTI, CostKind) >
+      UnrolledBase.spillCost(TTI, CostKind))
+    return;
+
+  Constant *Zero = ConstantInt::get(CanIVTy, 0);
+  InductionDescriptor ID(Zero, InductionDescriptor::IK_IntInduction,
+                         SE.getOne(CanIVTy));
+  VPIRValue *StartV = Plan.getZero(CanIVTy);
+  VPValue *StepV = Plan.getConstantInt(CanIVTy, 1);
+  auto *NewWideIV = new VPWidenIntOrFpInductionRecipe(
+      /*IV=*/nullptr, StartV, StepV, &Plan.getVF(), ID,
+      VPIRFlags::WrapFlagsTy(/*HasNUW=*/false, /*HasNSW=*/false),
+      WideCanIV->getDebugLoc());
----------------
fhahn wrote:

Yes, ScalarIVSteps produces scalars for all lanes, so all lanes would need to be packed into a vector again. This would be much worse

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