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

[Ramkumar Ramachandra 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;
+  }
----------------
artagnon wrote:

I just tried it out and the following patch:

```diff
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index 5e6d1bbcd5a7..b956683c5b2c 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -658,36 +658,28 @@ static void removeRedundantInductionCasts(VPlan &Plan) {
 /// Try to replace VPWidenCanonicalIVRecipes with a widened canonical IV
 /// recipe, if it exists.
 static void removeRedundantCanonicalIVs(VPlan &Plan) {
+  if (Plan.hasScalarVFOnly())
+    return;
   VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
   VPValue *CanonicalIV = LoopRegion->getCanonicalIV();
-  auto *WidenNewIV = vputils::findUserOf<VPWidenCanonicalIVRecipe>(CanonicalIV);
-
-  if (!WidenNewIV)
+  auto *WideCanIV = vputils::findUserOf<VPWidenCanonicalIVRecipe>(CanonicalIV);
+  if (!WideCanIV)
     return;

   VPBasicBlock *HeaderVPBB = LoopRegion->getEntryBasicBlock();
   for (VPRecipeBase &Phi : HeaderVPBB->phis()) {
-    auto *WidenOriginalIV = dyn_cast<VPWidenIntOrFpInductionRecipe>(&Phi);
-
-    if (!WidenOriginalIV || !WidenOriginalIV->isCanonical())
+    auto *WideInd = dyn_cast<VPWidenIntOrFpInductionRecipe>(&Phi);
+    if (!WideInd || !WideInd->isCanonical())
       continue;

-    // Replace WidenNewIV with WidenOriginalIV if WidenOriginalIV provides
-    // everything WidenNewIV's users need. That is, WidenOriginalIV will
-    // generate a vector phi or all users of WidenNewIV demand the first lane
-    // only.
-    if (Plan.hasScalarVFOnly() ||
-        !vputils::onlyScalarValuesUsed(WidenOriginalIV) ||
-        vputils::onlyFirstLaneUsed(WidenNewIV)) {
-      // We are replacing a wide canonical iv with a suitable wide induction.
-      // This is used to compute header mask, hence all lanes will be used and
-      // we need to drop wrap flags only applying to lanes guranteed to execute
-      // in the original scalar loop.
-      WidenOriginalIV->dropPoisonGeneratingFlags();
-      WidenNewIV->replaceAllUsesWith(WidenOriginalIV);
-      WidenNewIV->eraseFromParent();
-      return;
-    }
+    // We are replacing a wide canonical iv with a suitable wide induction.
+    // This is used to compute header mask, hence all lanes will be used and
+    // we need to drop wrap flags only applying to lanes guranteed to execute
+    // in the original scalar loop.
+    WideInd->dropPoisonGeneratingFlags();
+    WideCanIV->replaceAllUsesWith(WideInd);
+    WideCanIV->eraseFromParent();
+    return;
   }
 }
```

leads to minor regressions from adds not being CSE'd. I think the right thing to do is my preparatory patch to lower first-lane wide IVs to SIVSteps, expand the WideCanIV, finally making the proposed diff possible without regressions.

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