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

[Luke Lau via llvm-commits]

================
@@ -888,6 +889,67 @@ 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();
+
+    // We cannot create a phi in a VPIRBasicBlock, which would be required in
+    // the absence of a canonical IV to re-use, because its predecessors are not
+    // modeled in the Plan: fall back to the IR expander.
+    PHINode *ARCanIV = AR->getLoop()->getCanonicalInductionVariable();
+    if (!ARCanIV)
+      return vputils::getOrCreateVPValueForSCEVExpr(Plan, AR);
+
+    // Find a canonical IV to re-use. In case a BasicBlock identical to the one
+    // referred to by AR's loop header is present in the Plan, we can proceed.
+    VPIRBasicBlock *Header = nullptr;
+    for (auto *VPIRBB : VPBlockUtils::blocksOnly<VPIRBasicBlock>(
+             vp_depth_first_shallow(Plan.getEntry()))) {
+      if (VPIRBB->getIRBasicBlock() == AR->getLoop()->getHeader()) {
+        Header = VPIRBB;
+        break;
+      }
+    }
+
+    // The AR's loop refers to a loop that doesn't exist in the Plan: fall back
+    // to the IR expander.
+    if (!Header)
+      return vputils::getOrCreateVPValueForSCEVExpr(Plan, AR);
+
+    auto FoundCanIV = find_if(*Header, [ARCanIV](VPRecipeBase &R) {
+      auto *IRPhi = dyn_cast<VPIRPhi>(&R);
+      return IRPhi && &IRPhi->getIRPhi() == ARCanIV;
+    });
+    assert(FoundCanIV != Header->end() &&
+           "VPIRPhi in identical VPIRBasicBlock must be present");
+    VPValue *CanonicalIV = FoundCanIV->getVPSingleValue();
+
+    VPValue *Start;
+    Start = tryToExpand(AR->getStart());
+    if (!Start)
+      return nullptr;
+    VPValue *Step = tryToExpand(AR->getStepRecurrence(SE));
+    if (!Step)
+      return nullptr;
----------------
lukel97 wrote:

Is this correct if the step isn't loop invariant? Do we need to check if the addrec is affine first?

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