[llvm] [LV] Use SCEV to compute final value of complex induction variables (PR #195059)

[Ramkumar Ramachandra via llvm-commits]

================
@@ -1098,54 +1098,75 @@ static VPValue *tryToComputeEndValueForInduction(VPWidenInductionRecipe *WideIV,
 
 /// Attempts to optimize the induction variable exit values for users in the
 /// exit block coming from the latch in the original scalar loop.
-static VPValue *
-optimizeLatchExitInductionUser(VPlan &Plan, VPValue *Op,
-                               DenseMap<VPValue *, VPValue *> &EndValues,
-                               PredicatedScalarEvolution &PSE) {
+static VPValue *optimizeLatchExitInductionUser(
+    VPlan &Plan, VPValue *Op, DenseMap<VPValue *, VPValue *> &EndValues,
+    PredicatedScalarEvolution &PSE, VPValue *ResumeTC, const Loop *L) {
   VPValue *Incoming;
   if (!match(Op, m_ExtractLastLaneOfLastPart(m_VPValue(Incoming))))
     return nullptr;
 
-  VPWidenInductionRecipe *WideIV = getOptimizableIVOf(Incoming, PSE);
-  if (!WideIV)
-    return nullptr;
+  auto *ExtractR = cast<VPInstruction>(Op);
+  if (VPWidenInductionRecipe *WideIV = getOptimizableIVOf(Incoming, PSE)) {
+    VPValue *EndValue = EndValues.lookup(WideIV);
+    assert(EndValue && "Must have computed the end value up front");
+
+    // `getOptimizableIVOf()` always returns the pre-incremented IV, so if it
+    // changed it means the exit is using the incremented value, so we don't
+    // need to subtract the step.
+    if (Incoming != WideIV)
+      return EndValue;
+
+    // Otherwise, subtract the step from the EndValue.
+    VPBuilder B(ExtractR);
+    VPValue *Step = WideIV->getStepValue();
+    Type *ScalarTy = WideIV->getScalarType();
+    if (ScalarTy->isIntegerTy())
+      return B.createSub(EndValue, Step, DebugLoc::getUnknown(), "ind.escape");
+    if (ScalarTy->isPointerTy()) {
+      Type *StepTy = Step->getScalarType();
+      auto *Zero = Plan.getZero(StepTy);
+      return B.createPtrAdd(EndValue, B.createSub(Zero, Step),
+                            DebugLoc::getUnknown(), "ind.escape");
+    }
+    if (ScalarTy->isFloatingPointTy()) {
+      const auto &ID = WideIV->getInductionDescriptor();
+      return B.createNaryOp(
+          ID.getInductionBinOp()->getOpcode() == Instruction::FAdd
+              ? Instruction::FSub
+              : Instruction::FAdd,
+          {EndValue, Step}, {ID.getInductionBinOp()->getFastMathFlags()});
+    }
+    llvm_unreachable("all possible induction types must be handled");
+  }
 
-  VPValue *EndValue = EndValues.lookup(WideIV);
-  assert(EndValue && "Must have computed the end value up front");
+  const SCEV *IncomingSCEV = vputils::getSCEVExprForVPValue(Incoming, PSE, L);
+  const SCEV *Start, *Step;
+  if (!match(IncomingSCEV, m_scev_AffineAddRec(m_SCEV(Start), m_SCEV(Step),
+                                               m_SpecificLoop(L))))
+    return nullptr;
 
-  // `getOptimizableIVOf()` always returns the pre-incremented IV, so if it
-  // changed it means the exit is using the incremented value, so we don't
-  // need to subtract the step.
-  if (Incoming != WideIV)
-    return EndValue;
+  // TODO: Start value can be defined be a VPExpandSCEVRecipe after
+  // VPDerivedIVRecipe supports a general VPValue as the start value.
+  VPIRValue *StartIRV = vputils::getVPIRValueForSCEVExpr(Plan, Start);
+  if (!StartIRV)
+    return nullptr;
 
-  // Otherwise, subtract the step from the EndValue.
-  auto *ExtractR = cast<VPInstruction>(Op);
-  VPBuilder B(ExtractR);
-  VPValue *Step = WideIV->getStepValue();
-  Type *ScalarTy = WideIV->getScalarType();
-  if (ScalarTy->isIntegerTy())
-    return B.createSub(EndValue, Step, DebugLoc::getUnknown(), "ind.escape");
-  if (ScalarTy->isPointerTy()) {
-    Type *StepTy = Step->getScalarType();
-    auto *Zero = Plan.getZero(StepTy);
-    return B.createPtrAdd(EndValue, B.createSub(Zero, Step),
-                          DebugLoc::getUnknown(), "ind.escape");
-  }
-  if (ScalarTy->isFloatingPointTy()) {
-    const auto &ID = WideIV->getInductionDescriptor();
-    return B.createNaryOp(
-        ID.getInductionBinOp()->getOpcode() == Instruction::FAdd
-            ? Instruction::FSub
-            : Instruction::FAdd,
-        {EndValue, Step}, {ID.getInductionBinOp()->getFastMathFlags()});
-  }
-  llvm_unreachable("all possible induction types must be handled");
-  return nullptr;
+  Type *StartTy = StartIRV->getType();
+  assert(StartTy->isIntOrPtrTy() && "The type must be SCEVable");
+  InductionDescriptor::InductionKind Kind =
+      StartTy->isPointerTy() ? InductionDescriptor::IK_PtrInduction
+                             : InductionDescriptor::IK_IntInduction;
+  VPValue *StepVPV = vputils::getOrCreateVPValueForSCEVExpr(Plan, Step);
----------------
artagnon wrote:

I don't think we have test coverage for when this creates an ExpandSCEV? I wonder if it's possible to craft a test where Step is not a VPIRValue? If not, would be good to change this to get the IRV instead, like Start?

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