[llvm] [SCEV] Infer loop max trip count from memory accesses (PR #70361)

[Shilei Tian via llvm-commits]

================
@@ -8253,6 +8261,210 @@ ScalarEvolution::getSmallConstantTripMultiple(const Loop *L,
   return getSmallConstantTripMultiple(L, ExitCount);
 }
 
+/// Collect all load/store instructions that must be executed in every iteration
+/// of loop \p L .
+static void
+collectExecLoadStoreInsideLoop(const Loop *L, DominatorTree &DT,
+                               SmallVector<Instruction *, 4> &MemInsts) {
+  // It is difficult to tell if the load/store instruction is executed on every
+  // iteration inside an irregular loop.
+  if (!L->isLoopSimplifyForm() || !L->isInnermost())
+    return;
+
+  const BasicBlock *LoopLatch = L->getLoopLatch();
+  assert(LoopLatch && "normal form loop doesn't have a latch");
+  assert(L->getExitingBlock() == LoopLatch);
+
+  // We will not continue if sanitizer is enabled.
+  const Function *F = LoopLatch->getParent();
+  if (F->hasFnAttribute(Attribute::SanitizeAddress) ||
+      F->hasFnAttribute(Attribute::SanitizeThread) ||
+      F->hasFnAttribute(Attribute::SanitizeMemory) ||
+      F->hasFnAttribute(Attribute::SanitizeHWAddress) ||
+      F->hasFnAttribute(Attribute::SanitizeMemTag))
+    return;
+
+  for (auto *BB : L->getBlocks()) {
+    // We need to make sure that max execution time of MemAccessBB in loop
+    // represents latch max excution time. The BB below should be skipped:
+    //            Entry
+    //              │
+    //        ┌─────▼─────┐
+    //        │Loop Header◄─────┐
+    //        └──┬──────┬─┘     │
+    //           │      │       │
+    //  ┌────────▼──┐ ┌─▼─────┐ │
+    //  │MemAccessBB│ │OtherBB│ │
+    //  └────────┬──┘ └─┬─────┘ │
+    //           │      │       │
+    //         ┌─▼──────▼─┐     │
+    //         │Loop Latch├─────┘
+    //         └────┬─────┘
+    //              ▼
+    //             Exit
+    if (!DT.dominates(BB, LoopLatch))
+      continue;
+
+    for (Instruction &I : *BB) {
+      if (isa<LoadInst>(&I) || isa<StoreInst>(&I))
+        MemInsts.push_back(&I);
+    }
+  }
+}
+
+/// Return a SCEV representing the memory size of pointer \p V .
+static const SCEV *getCertainSizeOfMem(const SCEV *V, Type *RTy,
+                                       const DataLayout &DL,
+                                       const TargetLibraryInfo &TLI,
+                                       ScalarEvolution *SE) {
+  const SCEVUnknown *PtrBase = dyn_cast<SCEVUnknown>(V);
+  if (!PtrBase)
+    return nullptr;
+  Value *Ptr = PtrBase->getValue();
+  uint64_t Size = 0;
+  if (!llvm::getObjectSize(Ptr, Size, DL, &TLI))
+    return nullptr;
+  return SE->getConstant(RTy, Size);
+}
+
+/// Get the range of given index represented by \p AddRec.
+static const SCEV *getIndexRange(const SCEVAddRecExpr *AddRec,
+                                 ScalarEvolution *SE) {
+  const SCEV *Range = SE->getConstant(SE->getUnsignedRangeMax(AddRec) -
+                                      SE->getUnsignedRangeMin(AddRec));
+  const SCEV *Step = AddRec->getStepRecurrence(*SE);
+  if (SE->isKnownNegative(Step))
+    Step = SE->getNegativeSCEV(Step);
+  return SE->getUDivCeilSCEV(Range, Step);
+}
+
+/// Get the underlying SCEVAddExpr from a cast expression if possible.
+const SCEV *peelCastExpr(const SCEVCastExpr *S, ScalarEvolution *SE) {
+  const SCEV *Op = S->getOperand();
+  if (isa<SCEVCouldNotCompute>(Op))
+    return Op;
+  if (isa<SCEVAddExpr>(Op))
+    return Op;
+  if (isa<SCEVCastExpr>(Op))
+    return peelCastExpr(cast<SCEVCastExpr>(Op), SE);
+  return SE->getCouldNotCompute();
+}
+
+static Value *peelExt(Value *V) {
+  if (isa<ZExtInst>(V) || isa<SExtInst>(V))
+    return peelExt(cast<Instruction>(V)->getOperand(0));
+  return V;
+}
+
+static bool isIndexInductionVariable(PHINode *InductionVar, Value *Index) {
----------------
shiltian wrote:

This might be too conservative. It can't handle uses such as `array[i * 2]`. What is the best way to do that?

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