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

[Mel Chen via llvm-commits]

================
@@ -1097,51 +1097,75 @@ static VPValue *tryToComputeEndValueForInduction(VPWidenInductionRecipe *WideIV,
 /// exit block coming from the latch in the original scalar loop.
 static VPValue *optimizeLatchExitInductionUser(
     VPlan &Plan, VPTypeAnalysis &TypeInfo, VPValue *Op,
-    DenseMap<VPValue *, VPValue *> &EndValues, PredicatedScalarEvolution &PSE) {
+    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 = TypeInfo.inferScalarType(WideIV);
+    if (ScalarTy->isIntegerTy())
+      return B.createSub(EndValue, Step, DebugLoc::getUnknown(), "ind.escape");
+    if (ScalarTy->isPointerTy()) {
+      Type *StepTy = TypeInfo.inferScalarType(Step);
+      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;
+  auto *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 = TypeInfo.inferScalarType(WideIV);
-  if (ScalarTy->isIntegerTy())
-    return B.createSub(EndValue, Step, DebugLoc::getUnknown(), "ind.escape");
-  if (ScalarTy->isPointerTy()) {
-    Type *StepTy = TypeInfo.inferScalarType(Step);
-    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);
+  VPBuilder Builder(ExtractR);
+  Type *TCTy = TypeInfo.inferScalarType(ResumeTC);
+  VPValue *It = Builder.createSub(ResumeTC, Plan.getConstantInt(TCTy, 1),
----------------
Mel-Chen wrote:

Rename to `ExitCount`
2ad53e77cc001125ba240ad860584705f2d978ae

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