[Mlir-commits] [mlir] [mlir] Add splitForOpAtBound utility and use it in loop unrolling (PR #215108)

[Matthias Springer]

================
@@ -363,6 +364,112 @@ void mlir::generateUnrolledLoop(
   loopBodyBlock->getTerminator()->setOperands(lastYielded);
 }
 
+/// Splits `forOp` into two consecutive loops at `splitPoint`.
+FailureOr<std::pair<scf::ForOp, scf::ForOp>>
+mlir::splitForOpAtPoint(scf::ForOp forOp, Value splitPoint) {
+  if (splitPoint.getType() != forOp.getLowerBound().getType())
+    return failure();
+
+  // When a bound is constant, require a valid lattice-aligned split point.
+  // When it is dynamic, emit the same checks as runtime asserts.
+  OpBuilder runtimeBuilder(forOp);
+  Location loc = forOp.getLoc();
+  bool isUnsigned = forOp.getUnsignedCmp();
+  Value lbVal = forOp.getLowerBound();
+  Value ubVal = forOp.getUpperBound();
+  Value stepVal = forOp.getStep();
+  // Fail at runtime if `cond` is false.
+  auto emitAssert = [&](Value cond, StringRef msg) {
----------------
matthias-springer wrote:

Can you leave away the runtime asserts for now? We don't emit those for other loop helpers as far as I know. Instead, document the expectations in the header file.

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