[llvm] [LoopFlatten] Add option to version loops instead of widening IVs (PR #166156)

[John Brawn via llvm-commits]

================
@@ -835,14 +839,47 @@ static bool DoFlattenLoopPair(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
   return true;
 }
 
+static bool VersionLoop(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
+                        ScalarEvolution *SE, const LoopAccessInfo &LAI) {
+
+  // Version the loop. The overflow check isn't a runtime pointer check, so we
+  // pass an empty list of runtime pointer checks, causing LoopVersioning to
+  // emit 'false' as the branch condition, and add our own check afterwards.
+  BasicBlock *CheckBlock = FI.OuterLoop->getLoopPreheader();
+  ArrayRef<RuntimePointerCheck> Checks(nullptr, nullptr);
+  LoopVersioning LVer(LAI, Checks, FI.OuterLoop, LI, DT, SE);
+  LVer.versionLoop();
+
+  // Check for overflow by calculating the new tripcount using
+  // umul_with_overflow and then checking if it overflowed.
+  BranchInst *Br = dyn_cast<BranchInst>(CheckBlock->getTerminator());
+  if (!Br || !Br->isConditional())
+    return false;
+  if (!match(Br->getCondition(), m_Zero()))
+    return false;
+  IRBuilder<> Builder(Br);
+  Value *Call = Builder.CreateIntrinsic(Intrinsic::umul_with_overflow,
+                                        FI.OuterTripCount->getType(),
+                                        {FI.OuterTripCount, FI.InnerTripCount},
+                                        /*FMFSource=*/nullptr, "flatten.mul");
+  FI.NewTripCount = Builder.CreateExtractValue(Call, 0, "flatten.tripcount");
+  Value *Overflow = Builder.CreateExtractValue(Call, 1, "flatten.overflow");
+  Br->setCondition(Overflow);
+  return true;
+}
+
 static bool CanWidenIV(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
                        ScalarEvolution *SE, AssumptionCache *AC,
-                       const TargetTransformInfo *TTI) {
+                       const TargetTransformInfo *TTI,
+                       const LoopAccessInfo &LAI) {
   if (!WidenIV) {
     LLVM_DEBUG(dbgs() << "Widening the IVs is disabled\n");
     return false;
   }
 
+  if (VersionLoopsOverWiden && VersionLoop(FI, DT, LI, SE, LAI))
----------------
john-brawn-arm wrote:

Versioning the loop inside of CanWidenIV looks a bit strange, I think it makes more sense to do this in FlattenLoopPair, i.e. something like
```
if (VersionLoopsOverWiden && VersionLoop(FI, DT, LI, SE, LAI))
  CanFlatten = true;
else
  CanFlatten = CanWidenIV(FI, DT, LI, SE, AC, TTI);
```


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