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

[Mel Chen via llvm-commits]

================
@@ -1119,50 +1119,81 @@ static VPValue *tryToComputeEndValueForInduction(VPWidenInductionRecipe *WideIV,
 /// exit block coming from the latch in the original scalar loop.
 static VPValue *optimizeLatchExitInductionUser(
     VPlan &Plan, VPTypeAnalysis &TypeInfo, VPBlockBase *PredVPBB, VPValue *Op,
-    DenseMap<VPValue *, VPValue *> &EndValues, PredicatedScalarEvolution &PSE) {
+    DenseMap<VPValue *, VPValue *> &EndValues, PredicatedScalarEvolution &PSE,
+    VPValue *ResumeTC, const Loop *L) {
   VPValue *Incoming;
-  VPWidenInductionRecipe *WideIV = nullptr;
-  if (match(Op, m_ExtractLastLaneOfLastPart(m_VPValue(Incoming))))
-    WideIV = getOptimizableIVOf(Incoming, PSE);
-
-  if (!WideIV)
+  if (!match(Op, m_ExtractLastLaneOfLastPart(m_VPValue(Incoming))))
     return nullptr;
 
-  VPValue *EndValue = EndValues.lookup(WideIV);
-  assert(EndValue && "Must have computed the end value up front");
+  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;
+    // `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(cast<VPBasicBlock>(PredVPBB)->getTerminator());
-  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;
+    // Otherwise, subtract the step from the EndValue.
+    VPBuilder B(cast<VPBasicBlock>(PredVPBB)->getTerminator());
+    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");
+  }
+
+  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;
+
+  VPValue *StartVPV = vputils::getOrCreateVPValueForSCEVExpr(Plan, Start);
+  auto *StartIRV = dyn_cast<VPIRValue>(StartVPV);
+  if (!StartIRV) {
+    VPRecipeBase *Def = StartVPV->getDefiningRecipe();
+    assert(Def && "The value must be defined by VPExpandSCEVRecipe");
+    assert(StartVPV->getNumUsers() == 0 &&
+           "Newly created VPExpandSCEVRecipe should have no users");
+    Def->eraseFromParent();
+    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(cast<VPInstruction>(Op));
+  Type *TCTy = TypeInfo.inferScalarType(ResumeTC);
+  VPValue *It = Builder.createSub(ResumeTC, Plan.getConstantInt(TCTy, 1),
+                                  DebugLoc::getUnknown());
+  Type *StepTy = TypeInfo.inferScalarType(StepVPV);
+  It =
+      Builder.createScalarZExtOrTrunc(It, StepTy, TCTy, DebugLoc::getUnknown());
+  return Builder.createDerivedIV(Kind, /*FPBinOp=*/nullptr, StartIRV, It,
+                                 StepVPV);
----------------
Mel-Chen wrote:

I’ve considered using SCEV analysis to replace the current analysis of IV live-out optimization, including the  inserting ExitingIVValue. However, I’d like to determine if this approach is profitable first. If the SCEV-based method proves feasible, it could indeed replace getOptimizableIVOf.

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