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

[Florian Hahn via llvm-commits]

================
@@ -746,13 +751,148 @@ static bool violatesLegacyMultiExitLoopCheck(Loop *L) {
     });
 }
 
+namespace {
+// Represents a group of loads in a loop that can be combined into a wider one.
+struct LoadGroup {
+  // Base object being read.
+  Value *BasePtr;
+  // First load instruction in the program order.
+  LoadInst *FirstLoad;
+  // Pairs of (load instruction, offset from base).
+  SmallVector<std::pair<LoadInst *, APInt>, 4> Loads;
+  // An applicable wider integer type to load as.
+  Type *WideType;
+};
+
+// Helper to compute load group span and validate for widening.
+static std::optional<LoadGroup> tryFormLoadGroupForWidening(
+    Value *Base, SmallVectorImpl<std::pair<LoadInst *, APInt>> &Loads, Loop &L,
+    ScalarEvolution &SE, const DataLayout &DL, const TargetTransformInfo &TTI) {
+  // Find the span of the loaded data.
+  int64_t Left = INT64_MAX;
+  int64_t Right = INT64_MIN;
+  for (const auto &[Load, Offset] : Loads) {
+    Left = std::min(Left, Offset.getSExtValue());
+    Right = std::max(
+        Right, Offset.getSExtValue() +
+                   static_cast<int64_t>(DL.getTypeStoreSize(Load->getType())));
+  }
+  assert((Left < Right) && "Invalid load group span");
+  uint64_t TotalBytes = Right - Left;
+  uint64_t TotalBits = TotalBytes * 8;
+  // Powers of two are already natural for most targets.
+  if (isPowerOf2_64(TotalBits))
+    return std::nullopt;
+  Type *WideType =
+      DL.getSmallestLegalIntType(L.getHeader()->getContext(), TotalBits);
+  if (!WideType)
+    return std::nullopt;
+  unsigned WideBits = WideType->getIntegerBitWidth();
+  // Total size is already natural for the target.
+  if (WideBits == TotalBits)
+    return std::nullopt;
+  // Peeling doubles dereferenceable bytes, ensure wide type fits.
+  if (WideBits > TotalBits * 2)
+    return std::nullopt;
+  // Check alignment is unconstrained.
+  unsigned Fast = 0;
+  if (!TTI.allowsMisalignedMemoryAccesses(L.getHeader()->getContext(), WideBits,
+                                          DL.getDefaultGlobalsAddressSpace(),
+                                          Align(1), &Fast) &&
+      !Fast)
+    return std::nullopt;
+  // Validate pointer stride across iterations.
+  const SCEV *PtrSCEV = SE.getSCEV(Base);
+  const SCEVAddRecExpr *AR = dyn_cast<SCEVAddRecExpr>(PtrSCEV);
+  if (!AR || AR->getLoop() != &L)
+    return std::nullopt;
+  const SCEV *Step = AR->getStepRecurrence(SE);
+  auto *ConstStep = dyn_cast<SCEVConstant>(Step);
+  if (!ConstStep)
+    return std::nullopt;
+  int64_t StepVal = ConstStep->getValue()->getSExtValue();
+  if (StepVal != static_cast<int64_t>(TotalBytes))
+    return std::nullopt;
+
+  LoadInst *FirstLoad = Loads[0].first;
+  llvm::sort(Loads, [](const auto &A, const auto &B) {
+    return A.second.slt(B.second);
+  });
+  return LoadGroup{Base, FirstLoad, std::move(Loads), WideType};
+}
+
+// Find groups of consecutive loads in a basic block for peeling purposes
+static SmallVector<LoadGroup>
+findLoadGroupsForWidening(BasicBlock *BB, Loop &L, ScalarEvolution &SE,
+                          const DataLayout &DL,
+                          const TargetTransformInfo &TTI) {
+  SmallVector<LoadGroup> Groups;
+  // Mapping from base pointer to loads instructions and their offset from the
+  // base.
+  DenseMap<Value *, SmallVector<std::pair<LoadInst *, APInt>>> LoadsByBase;
+
+  auto ProcessCollectedLoads = [&]() {
+    for (auto &[Base, Loads] : LoadsByBase) {
+      if (Loads.size() <= 1)
+        continue;
+      if (auto Group = tryFormLoadGroupForWidening(Base, Loads, L, SE, DL, TTI))
+        Groups.emplace_back(*std::move(Group));
+    }
+    LoadsByBase.clear();
+  };
+
+  for (Instruction &I : *BB) {
+    if (auto *Load = dyn_cast<LoadInst>(&I)) {
+      if (Load->isVolatile() || Load->isAtomic() ||
+          !Load->getType()->isIntegerTy() ||
+          Load->getPointerAddressSpace() != DL.getDefaultGlobalsAddressSpace())
+        continue;
+      Value *Ptr = Load->getPointerOperand();
+      APInt Offset(DL.getIndexTypeSizeInBits(Ptr->getType()), 0);
+      Value *ActualBase = Ptr->stripAndAccumulateConstantOffsets(
+          DL, Offset, /*AllowNonInbounds=*/false);
+      LoadsByBase[ActualBase].emplace_back(Load, Offset);
+    } else if (I.mayHaveSideEffects())
+      ProcessCollectedLoads();
+  }
+  ProcessCollectedLoads();
+  return Groups;
+}
+
+// Returns 1 if peeling the last iteration would enable widening load groups to
+// natural sizes. Returns 0 otherwise.
+static unsigned peelLastForLoadWidening(Loop &L, ScalarEvolution &SE,
+                                        const DataLayout &DL,
+                                        const TargetTransformInfo &TTI,
+                                        DominatorTree &DT) {
+  if (!EnablePeelForLoadWidening)
+    return 0;
+  if (!L.isInnermost())
+    return 0;
+  if (!canPeelLastIteration(L, SE))
----------------
fhahn wrote:

can we combine this with the other peeling-from-last-iteration logic to avoid repated calls ?

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