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

[Guy David via llvm-commits]

================
@@ -749,13 +753,194 @@ 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) sorted by offset.
+  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) {
+  // Verify all loads use the same address space.
+  unsigned AddrSpace = Loads[0].first->getPointerAddressSpace();
+  for (const auto &[Load, Offset] : Loads) {
+    if (Load->getPointerAddressSpace() != AddrSpace)
+      return std::nullopt;
+  }
+
+  // Find the span of the loaded data.
+  int64_t Left = std::numeric_limits<int64_t>::max();
+  int64_t Right = std::numeric_limits<int64_t>::min();
+  for (const auto &[Load, Offset] : Loads) {
+    int64_t OffsetVal = Offset.getSExtValue();
+    int64_t StoreSize =
+        static_cast<int64_t>(DL.getTypeStoreSize(Load->getType()));
+    if (OffsetVal > std::numeric_limits<int64_t>::max() - StoreSize)
+      return std::nullopt;
+    Left = std::min(Left, OffsetVal);
+    Right = std::max(Right, OffsetVal + StoreSize);
+  }
+  assert((Left < Right) && "Invalid load group span");
+  uint64_t TotalBytes = Right - Left;
+  uint64_t TotalBits = TotalBytes * 8;
+  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, no benefit from widening.
+  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 and without penalty.
+  unsigned Fast = 0;
+  if (!TTI.allowsMisalignedMemoryAccesses(L.getHeader()->getContext(), WideBits,
+                                          AddrSpace, Align(1), &Fast) ||
+      !Fast)
+    return std::nullopt;
+  // Validate pointer stride across iterations.
+  const SCEVConstant *ConstStep = nullptr;
+  if (!match(SE.getSCEV(Base),
+             m_scev_AffineAddRec(m_SCEV(), m_SCEVConstant(ConstStep),
+                                 m_SpecificLoop(&L))))
+    return std::nullopt;
+  int64_t StepVal = ConstStep->getValue()->getSExtValue();
+  if (StepVal != static_cast<int64_t>(TotalBytes))
+    return std::nullopt;
+
+  LoadInst *FirstLoad = Loads[0].first;
+  // Cost model: compare cost of individual loads vs wide load + extraction ops.
+  // Sort loads by offset first since we need this for cost calculation.
+  llvm::sort(Loads, [](const auto &A, const auto &B) {
+    return A.second.slt(B.second);
+  });
+  int64_t FirstOffset = Loads[0].second.getSExtValue();
+
+  TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput;
+
+  // Cost of original individual loads.
+  InstructionCost OriginalCost = 0;
+  for (const auto &[Load, Offset] : Loads) {
+    OriginalCost += TTI.getMemoryOpCost(Instruction::Load, Load->getType(),
+                                        Load->getAlign(), AddrSpace, CostKind);
+  }
+
+  // Cost of wide load + extraction operations (shift + trunc for each load).
----------------
guy-david wrote:

I thought about adding a helper in DataLayout but it couldn't find a clean refactor, mainly because `scalarizeExtExtract` uses a vector type while here there's none.

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