[llvm] [LoopPeel] Peel last iteration to enable load widening (PR #173420)

[Guy David via llvm-commits]

================
@@ -895,6 +1098,70 @@ void llvm::computePeelCount(Loop *L, unsigned LoopSize,
   }
 }
 
+bool llvm::widenLoadsAfterPeel(Loop &L, ScalarEvolution &SE,
+                               const TargetTransformInfo &TTI,
+                               DominatorTree &DT) {
+  BasicBlock *Latch = L.getLoopLatch();
+  assert(Latch && "Loop should have a latch after peeling");
+  const DataLayout &DL = L.getHeader()->getDataLayout();
+  bool Changed = false;
+
+  for (BasicBlock *BB : L.blocks()) {
+    if (!DT.dominates(BB, Latch))
+      continue;
+    SmallVector<LoadGroup> Groups =
+        findLoadGroupsForWidening(BB, L, SE, DL, TTI);
+    for (const LoadGroup &Group : Groups) {
+      LoadInst *InsertPoint = Group.FirstLoad;
+      IRBuilder<> Builder(InsertPoint);
+      Value *BasePtr = Group.BasePtr;
+      int64_t FirstOffset = Group.Loads[0].second.getSExtValue();
+      unsigned AddrSpace = Group.FirstLoad->getPointerAddressSpace();
+      // If the first load doesn't start at offset 0, we need to adjust.
+      if (FirstOffset != 0) {
+        Value *OrigPtr = Group.BasePtr;
+        BasePtr = Builder.CreatePtrAdd(
+            OrigPtr,
+            ConstantInt::get(Builder.getIndexTy(DL, AddrSpace), FirstOffset));
+      }
+      // Merge AA metadata from all loads using intersection for correctness.
+      AAMDNodes AATags = InsertPoint->getAAMetadata();
+      for (const auto &[Load, Offset] : Group.Loads) {
+        if (Load != InsertPoint)
+          AATags = AATags.merge(Load->getAAMetadata());
+      }
+      // Create the wider load.
+      LoadInst *WideLoad = Builder.CreateLoad(Group.WideType, BasePtr);
+      unsigned SizeInBits = WideLoad->getType()->getScalarSizeInBits();
+      if (AATags)
+        WideLoad->setAAMetadata(AATags);
+      // For each original load, extract the corresponding bytes.
+      for (const auto &[Load, Offset] : Group.Loads) {
+        unsigned LoadBits = Load->getType()->getScalarSizeInBits();
+        Value *Extracted = WideLoad;
+        uint64_t BitPosition = (Offset.getSExtValue() - FirstOffset) * 8;
+        uint64_t BitOffset =
+            DL.getElementBitOffset(BitPosition, LoadBits, SizeInBits);
+        if (BitOffset != 0)
+          Extracted = Builder.CreateLShr(
+              Extracted, ConstantInt::get(WideLoad->getType(), BitOffset));
+        if (LoadBits < SizeInBits)
+          Extracted = Builder.CreateTrunc(Extracted, Load->getType());
+        Load->replaceAllUsesWith(Extracted);
+      }
+      // Delete the original loads.
+      for (auto &[Load, Offset] : Group.Loads)
+        Load->eraseFromParent();
----------------
guy-david wrote:

Formatter actually changes this back into a regular for loop.

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