[llvm] [LAA] Be more careful when evaluating AddRecs at symbolic max BTC. (PR #128061)

[Florian Hahn via llvm-commits]

================
@@ -206,11 +288,37 @@ std::pair<const SCEV *, const SCEV *> llvm::getStartAndEndForAccess(
   const SCEV *ScStart;
   const SCEV *ScEnd;
 
+  auto &DL = Lp->getHeader()->getDataLayout();
+  Type *IdxTy = DL.getIndexType(PtrExpr->getType());
+  const SCEV *EltSizeSCEV = SE->getStoreSizeOfExpr(IdxTy, AccessTy);
   if (SE->isLoopInvariant(PtrExpr, Lp)) {
     ScStart = ScEnd = PtrExpr;
   } else if (auto *AR = dyn_cast<SCEVAddRecExpr>(PtrExpr)) {
     ScStart = AR->getStart();
-    ScEnd = AR->evaluateAtIteration(MaxBECount, *SE);
+    if (!isa<SCEVCouldNotCompute>(BTC))
+      // Evaluating AR at an exact BTC is safe: LAA separately checks that
+      // accesses cannot wrap in the loop. If evaluating AR at BTC wraps, then
+      // the loop either triggers UB when executing a memory access with a
+      // poison pointer or the wrapping/poisoned pointer is not used.
+      ScEnd = AR->evaluateAtIteration(BTC, *SE);
+    else {
+      // Evaluating AR at MaxBTC may wrap and create an expression that is less
+      // than the start of the AddRec due to wrapping (for example consider
+      // MaxBTC = -2). If that's the case, set ScEnd to -(EltSize + 1). ScEnd
----------------
fhahn wrote:

If there is computable BTC, we must take the backedge exactly BTC times if the pointer expression would wrap then the original loop is guarnateed to have UB.

If we only have a symbolic max BTC, we take the backedge up to BTC times, but the original loop could exit before we wrap and trigger UB. If we would generate start and end pointers based on the symbolic max before the loop, those may wrap and cause incorrect runtime check results.

With the computable BTC, the start/end pointers can only wrap if the loop has UB.

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