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

[Caroline Newcombe via llvm-commits]

================
@@ -0,0 +1,542 @@
+//===- 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 isLegal() 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 isLegal() 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 by isLegal(): loop-carried values, values
+// that escape the loop (exit values), non-unit and non-integer inductions,
----------------
cenewcombe wrote:

We have a Fortran workload where this transformation is the remaining blocker on vectorizing an OpenMP reduction, so I've been following this with interest.

I tried this PR against it and it stops at `loop has exit values`, since the loop carries a scalar accumulator across iterations and its result is used after the loop. I then tried #205995 as it was before the revert, and it splits the same loop correctly and rebuilds the accumulator across the partitions.

Do you plan to restore exit values after landing this PR? Feel free to tag me if you do put up a patch, I'd be happy to share the Fortran reproducer, or to test patches against our workload, if either would be useful.

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