[llvm] [VPlan] Expand AddRecs in VPSCEVExpander (PR #209921)

[Florian Hahn via llvm-commits]

================
@@ -1059,6 +1059,50 @@ VPValue *VPSCEVExpander::tryToExpand(const SCEV *S) {
                                              ResultTy, DL);
     return Result;
   }
+  case scAddRecExpr: {
+    // AddRecs never appear in the vector loop; the AR's loop would correspond
+    // to an outer loop outside the vector loop, and its header would be modeled
+    // as a VPIRBasicBlock.
+    auto *AR = cast<SCEVAddRecExpr>(S);
+    VPlan &Plan = Builder.getPlan();
+
+    // If a canonical IV to re-use is present, it would be in the Plan's entry.
+    // We cannot create a phi in the Plan's entry, which would be required in
+    // the absence of a canonical IV to re-use or if AR is non-affine, because
+    // its predecessors are not modeled. This code seems to be unreachable for
+    // AddRecs of pointer type.
+    if (!AR->isAffine() || AR->getType()->isPointerTy())
+      return nullptr;
+    auto FoundCanIV =
+        find_if(Plan.getEntry()->phis(), [&](const VPRecipeBase &R) {
+          if (!SE.isSCEVable(cast<VPIRPhi>(R).getIRPhi().getType()))
+            return false;
+          const SCEV *Candidate = SE.getSCEV(&cast<VPIRPhi>(R).getIRPhi());
+          return match(Candidate,
+                       m_scev_AffineAddRec(m_scev_Zero(), m_scev_One(),
+                                           m_SpecificLoop(AR->getLoop()))) &&
+                 Candidate->getType() == AR->getType();
+        });
+    if (FoundCanIV == Plan.getEntry()->phis().end())
+      return nullptr;
+
+    VPValue *CanonicalIV =
+        Plan.getOrAddLiveIn(&cast<VPIRPhi>(FoundCanIV)->getIRPhi());
+    VPValue *Start;
+    Start = tryToExpand(AR->getStart());
+    if (!Start)
+      return nullptr;
+    VPValue *Step = tryToExpand(AR->getStepRecurrence(SE));
+    if (!Step)
+      return nullptr;
+
+    // {X,+,F} --> X + {0,+,F}
+    // {0,+,F} --> {0,+,1} * F
+    VPIRFlags::WrapFlagsTy NWFlags = {AR->hasNoUnsignedWrap(), false};
----------------
fhahn wrote:

I think unconditionally dropping `NSW` is causing a regression in one of the tests. Can we use SCEV to preserve when safe?

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