[llvm] [SCEV] Try to prove no-wrap for AddRecs via BTC. (PR #131538)

[Florian Hahn via llvm-commits]

================
@@ -5090,6 +5090,38 @@ ScalarEvolution::proveNoWrapViaConstantRanges(const SCEVAddRecExpr *AR) {
   return Result;
 }
 
+/// Return true if \p AR is known to not wrap via the loop's backedge-taken
+/// count.
+static SCEV::NoWrapFlags proveNoWrapViaBTC(const SCEVAddRecExpr *AR,
+                                           ScalarEvolution &SE) {
+  SCEV::NoWrapFlags Result = SCEV::FlagAnyWrap;
+  if (AR->hasNoUnsignedWrap() && AR->hasNoSignedWrap())
+    return Result;
+
+  const Loop *L = AR->getLoop();
+  const SCEV *BTC = SE.getBackedgeTakenCount(L);
+  if (isa<SCEVCouldNotCompute>(BTC) || !AR->getStepRecurrence(SE)->isOne())
+    return Result;
+
+  Type *WTy = SE.getWiderType(AR->getType(), BTC->getType());
+  // If AR's type is wider than BTC, we can zero extend BTC, otherwise bail out.
+  if (WTy != AR->getType() || !WTy->isIntegerTy())
+    return Result;
+
+  const SCEV *ExtBTC = SE.getNoopOrZeroExtend(BTC, WTy);
+  // AR has a step of 1, it is NUW/NSW if Start + BTC >= Start.
+  const SCEV *Add = SE.getAddExpr(AR->getStart(), ExtBTC);
+  if (!AR->hasNoUnsignedWrap() &&
+      SE.willNotOverflow(Instruction::Add, false, AR->getStart(), ExtBTC))
+    Result = ScalarEvolution::setFlags(Result, SCEV::FlagNUW);
+  if (!AR->hasNoSignedWrap() &&
+      SE.willNotOverflow(Instruction::Add, true, AR->getStart(), ExtBTC) &&
+      SE.isKnownPredicate(CmpInst::ICMP_SGE, Add, AR->getStart()))
----------------
fhahn wrote:

There's an issue here with `willNotOverflow`. When passing `IsSigned=true`, both input values will be sign-extended, so if the BTC is -1 we won't wrap, even though we should. 

I think what would need to happen is for willNotOverflow to zero-extend the passed BTC. I removed handling for signed for now. I'll put up a follow-up with a test case if it is actually useful in practice.

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