[llvm] [Transforms][Utils] Add LoopSplit for iteration-space loop splitting (PR #217232)

[Florian Hahn via llvm-commits]

================
@@ -0,0 +1,564 @@
+//===- LoopSplit.cpp - Split a loop's iteration space ---------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Splits a counted loop's iteration space into a chain of per-partition
+// sub-loops. See LoopSplit.h for the high-level usage guidelines.
+//
+// Structure produced for partitions [S0,E0], [S1,E1], ... where E is the loop's
+// last iteration and each clamped end sel_i is min(E_i, E), or max descending:
+//
+//   guard0:                            ; every S_i and sel_i is computed here
+//     if (S0 <= sel0) goto preheader0 else goto guard1
+//   loop0: ...                         ; latch iterates while i < sel0
+//   exit0 -> guard1
+//   guard1:
+//     if (S1 <= sel1) goto preheader1 else goto guard2
+//   loop1: ...                         ; latch iterates while i < sel1
+//   exit1 -> guard2
+//     ...
+//   final.exit:
+//
+// Each guard holds the "S_i <= sel_i" check and skips an empty partition by
+// falling through to the next guard. All S_i/sel_i are materialized once in
+// guard0, and the end clamp keeps the "runs at least once" iteration in the
+// right partition.
+//
+// The latch keeps iterating while the value the next iteration would use is
+// still in the partition. That is written as the strict "i < sel_i" on the
+// induction PHI rather than "i + 1 <= sel_i" on the step value; the two agree
+// because legality analysis has established that the space does not wrap, and
+// the strict form never forms i + 1, so it remains a real test even when sel_i
+// is the last value of the type, where the inclusive one would be a tautology
+// and the partition would never exit.
+//
+// A descending (step -1) loop uses the same structure mirrored: partitions run
+// high-to-low and the clamp and predicates flip (>=/>).
+//
+// Usage guidelines:
+//  - Caller bounds must not wrap the induction type. The clamp absorbs a bound
+//    past the runtime trip count, and legality analysis reserves the one step
+//    past the induction start that an empty partition needs, but a bound
+//    reaching any further wraps in the bound arithmetic and cannot be repaired
+//    here.
+//  - Bounds must be loop-invariant: they are expanded in guard0, so a bound
+//    depending on a value defined inside the loop cannot be placed.
+//  - The partitions must tile the original iteration space exactly -- same
+//    iterations, same order -- so the split preserves program behaviour.
+//
+// The transform is structural: inside a partition it only seeds the induction
+// PHI with that partition's start and replaces the latch test. It never
+// rebuilds a value that flows between partitions, so no SSA reconstruction is
+// needed.
+//
+// Not yet supported, and rejected during legality analysis: loop-carried
+// values, values that escape the loop (exit values), non-unit and non-integer
+// inductions, top-tested loops, and multiple exits. Also rejected is an
+// induction start at the extreme of the iteration direction, which leaves
+// nowhere to put a boundary. An induction *end* at that extreme is fine,
+// because the latch stays strict.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/LoopSplit.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/ScalarEvolution.h"
+#include "llvm/Analysis/ScalarEvolutionExpressions.h"
+#include "llvm/Analysis/ScalarEvolutionPatternMatch.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/ProfDataUtils.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/LoopUtils.h"
+#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
+#include "llvm/Transforms/Utils/ValueMapper.h"
+#include <optional>
+
+using namespace llvm;
+using namespace llvm::SCEVPatternMatch;
+
+#define DEBUG_TYPE "loop-split"
+
+//===----------------------------------------------------------------------===//
+// LoopSplit - construction, partition list, induction analysis
+//===----------------------------------------------------------------------===//
+
+/// Per-split() scratch shared by the phase helpers; lives for one split() call.
+/// Everything derived from the induction lives on LoopSplit itself, filled in
+/// by legality analysis; this holds only what the transform creates. Partition
+/// 0 reuses the original loop's blocks, which live in Partitions[0].
+struct LoopSplit::SplitState {
+  // Where the partition chain converges.
+  BasicBlock *FinalExit = nullptr;
+
+  // Parent of the new blocks, if any.
+  Loop *OuterLoop = nullptr;
+
+  // The loop's induction variable.
+  PHINode *Induction = nullptr;
+
+  // !llvm.loop on the loop before splitting.
+  MDNode *OrigLoopID = nullptr;
+};
+
+// Record a new partition with the given inclusive iteration range.
+void LoopSplit::addPartition(const SCEV *Start, const SCEV *End) {
+  assert(InductionEnd && "addPartition() requires prior legality analysis");
+  // The bounds are combined with the induction end and expanded in its type. A
+  // mismatch would otherwise surface either as a bare "Operand types don't
+  // match!" from inside ScalarEvolution, or worse, as a silent cast.
+  assert(Start->getType() == InductionEnd->getType() &&
+         End->getType() == InductionEnd->getType() &&
+         "partition bounds must have the induction type");
+  if (Partitions.empty()) {
+    assert((Start == InductionStart ||
+            SE->isKnownPredicate(ICmpInst::ICMP_EQ, Start, InductionStart)) &&
+           "first partition Start must match the induction start");
+  }
+  Partitions.emplace_back(Start, End);
+}
+
+// Return the induction's add-recurrence, or null unless the induction is an
+// integer with a unit step.
+static const SCEVAddRecExpr *analyzeInduction(PHINode *Induction,
+                                              ScalarEvolution *SE) {
+  // Partition bounds are integer arithmetic on the induction type, so a loop
+  // whose only induction is a pointer is out of scope.
+  if (!Induction->getType()->isIntegerTy())
+    return nullptr;
+  const SCEV *IndSCEV = SE->getSCEV(Induction);
+  // Match an affine add-recurrence and capture its constant step; accept a unit
+  // step in either direction: +1 (ascending) or -1 (descending).
+  const APInt *Step;
+  if (!match(IndSCEV, m_scev_AffineAddRec(m_SCEV(), m_scev_APInt(Step))))
+    return nullptr;
+  if (!Step->isOne() && !Step->isAllOnes())
+    return nullptr;
+  return cast<SCEVAddRecExpr>(IndSCEV);
+}
+
+// Decide whether the iteration ordering is signed or unsigned; returns the
+// signedness, or nullopt if it cannot be proven.
+static std::optional<bool> computeSignedness(ScalarEvolution &SE, Loop *L,
+                                             const SCEVAddRecExpr *IndAR) {
+  ICmpInst::Predicate P = L->getLatchCmpInst()->getPredicate();
+  // A relational predicate gives the ordering directly; for eq/ne fall back to
+  // the recurrence's no-wrap flags.
+  if (ICmpInst::isRelational(P))
+    return ICmpInst::isSigned(P);
+  if (IndAR->hasNoSignedWrap() && IndAR->hasNoUnsignedWrap()) {
+    const ConstantRange UR = SE.getUnsignedRange(IndAR);
+    const ConstantRange SR = SE.getSignedRange(IndAR);
+    if (!SR.isFullSet())
+      return true;
+    if (!UR.isFullSet())
+      return false;
+    LLVM_DEBUG(dbgs() << DEBUG_TYPE
+               ": cannot prove iteration ordering signedness\n");
+    return std::nullopt;
+  }
+  if (IndAR->hasNoSignedWrap())
+    return true;
+  if (IndAR->hasNoUnsignedWrap())
+    return false;
+  LLVM_DEBUG(dbgs() << DEBUG_TYPE
+             ": cannot prove iteration ordering signedness\n");
+  return std::nullopt;
+}
+
+// Prove \p Pred between \p LHS and \p RHS on entry to \p L, with loop guards
+// folded in so a bound fixed by a dominating condition is seen.
+static bool isEntryGuardedByCond(ScalarEvolution &SE, Loop *L,
+                                 ICmpInst::Predicate Pred, const SCEV *LHS,
+                                 const SCEV *RHS) {
+  auto Guards = ScalarEvolution::LoopGuards::collect(L, SE);
+  return SE.isLoopEntryGuardedByCond(L, Pred, SE.applyLoopGuards(LHS, Guards),
+                                     SE.applyLoopGuards(RHS, Guards));
+}
+
+// Latch "keep iterating" predicate, comparing the induction PHI against the
+// partition end: `i < sel` ascending, `i > sel` descending.
+static ICmpInst::Predicate continuePredicate(bool Signed, bool Descending) {
+  ICmpInst::Predicate P = Descending ? ICmpInst::ICMP_UGT : ICmpInst::ICMP_ULT;
+  return Signed ? ICmpInst::getSignedPredicate(P) : P;
+}
+
+// Guard "enter this partition" predicate: the latch test made non-strict, so
+// `start <= sel` ascending and `start >= sel` descending. Also used during
+// legality analysis to prove the iteration space is monotonic.
+static ICmpInst::Predicate guardPredicate(bool Signed, bool Descending) {
+  return ICmpInst::getNonStrictPredicate(continuePredicate(Signed, Descending));
+}
+
+struct LoopSplitAnalysis {
----------------
fhahn wrote:

put in anonymous namespace

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