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

[Ashutosh Nema via llvm-commits]

================
@@ -0,0 +1,143 @@
+//===- LoopSplitPass.cpp - Test driver for LoopSplit ----------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass drives LoopSplit from `opt` for testing. For every eligible loop it
+// builds partitions from the -loop-split-points offsets and splits the loop.
+// Which loops are eligible is chosen by -loop-split-depth; the default is the
+// innermost ones.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/LoopSplitPass.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/ScalarEvolution.h"
+#include "llvm/Analysis/ScalarEvolutionExpressions.h"
+#include "llvm/Analysis/ScalarEvolutionPatternMatch.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/Function.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Transforms/Utils/LoopSplit.h"
+
+using namespace llvm;
+using namespace llvm::SCEVPatternMatch;
+
+#define DEBUG_TYPE "loop-split"
+
+static cl::list<unsigned>
+    SplitPoints("loop-split-points",
+                cl::desc("Iteration offsets (relative to the induction start) "
+                         "at which to split each loop"),
+                cl::CommaSeparated);
+
+static cl::opt<unsigned> SplitDepth(
+    "loop-split-depth",
+    cl::desc(
+        "Split the loops at this nesting depth (1 is outermost) instead of "
+        "the innermost ones"),
+    cl::init(0));
+
+// Build the partition list for \p L from the command-line split offsets and run
+// the transform. Returns true if the loop was split.
+static bool splitLoop(Loop *L, ScalarEvolution &SE, DominatorTree &DT,
+                      LoopInfo &LI) {
+  LoopSplit LS(L, &LI, &SE, &DT);
+  if (!LS.isLegal()) {
+    LLVM_DEBUG(dbgs() << DEBUG_TYPE ": loop is not legal for splitting\n");
+    return false;
+  }
+
+  // isLegal() has already established this shape.
+  const SCEV *IndVarSCEV = SE.getSCEV(LS.getInductionVariable());
+  const SCEV *Start;
+  const APInt *StepC;
+  [[maybe_unused]] bool Matched = match(
+      IndVarSCEV, m_scev_AffineAddRec(m_SCEV(Start), m_scev_APInt(StepC)));
+  assert(Matched && "isLegal() guarantees a unit-step affine induction");
+
+  const SCEV *BTC = SE.getBackedgeTakenCount(L);
+  const SCEV *End = LS.getInductionEnd();
+  Type *Ty = Start->getType();
+  unsigned BitWidth = Ty->getIntegerBitWidth();
+  // The backedge-taken count is a separate expression and need not share the
+  // induction's width, so coerce it before doing arithmetic in that type.
+  const SCEV *Count = SE.getTruncateOrZeroExtend(BTC, Ty);
+
+  // Build boundaries in iteration order, stepping away from Start by each
+  // offset (down for a descending loop). Each offset opens a new partition at
+  // iteration `Start +/- offset`; the previous partition ends one step before.
+  bool Descending = StepC->isAllOnes();
+
+  // Boundaries must be increasing and distinct to tile the space, so sort and
+  // unique the offsets. Drop any that do not fit the induction type; truncating
+  // would reorder them and the partitions would overlap.
+  SmallVector<unsigned, 4> Offsets;
+  for (unsigned Offset : SplitPoints)
+    if (BitWidth >= 32 || Offset < (1u << BitWidth))
----------------
nema-ashutosh wrote:

BitWidth >= 32 is there on purpose so we never evaluate 1u << BitWidth for wide induction types, since that would be UB.

Since -loop-split-points is unsigned, anything you can pass from the CLI already fits in i32 and above, so we do not need the extra check there. For narrow types like i8, the other branch drops offsets that would truncate and mess up partition order. That is what oversized-split-offset.ll is testing.

I will add an i128 case to cover the wide type path.

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