[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) {
+    forOp->getContext()->getOrLoadDialect<cf::ControlFlowDialect>();
+    cf::AssertOp::create(runtimeBuilder, loc, cond,
+                         runtimeBuilder.getStringAttr(msg));
+  };
+  // Compare according to the loop's signedness.
+  auto emitICmp = [&](arith::CmpIPredicate signedPred,
+                      arith::CmpIPredicate unsignedPred, Value lhs, Value rhs) {
+    return arith::CmpIOp::create(
+        runtimeBuilder, loc, isUnsigned ? unsignedPred : signedPred, lhs, rhs);
+  };
+  // Check the split point statically when constant, otherwise at runtime.
+  auto checkSplitPoint = [&](auto getBound) -> LogicalResult {
+    auto lb = getBound(lbVal);
+    auto ub = getBound(ubVal);
+    auto step = getBound(stepVal);
+    auto split = getBound(splitPoint);
+    // lowerBound <= splitPoint
+    if (lb && split) {
+      if (*lb > *split)
+        return failure();
+    } else {
+      emitAssert(emitICmp(arith::CmpIPredicate::sle, arith::CmpIPredicate::ule,
+                          lbVal, splitPoint),
+                 "splitForOpAtPoint: split point below lower bound");
+    }
+    // splitPoint < upperBound
+    if (split && ub) {
+      if (*split >= *ub)
+        return failure();
+    } else {
+      emitAssert(emitICmp(arith::CmpIPredicate::slt, arith::CmpIPredicate::ult,
+                          splitPoint, ubVal),
+                 "splitForOpAtPoint: split point not below upper bound");
+    }
+    // step > 0
+    if (step) {
+      if (*step <= 0)
+        return failure();
+    } else {
+      Value zero = arith::ConstantOp::create(
+          runtimeBuilder, loc,
+          runtimeBuilder.getIntegerAttr(stepVal.getType(), 0));
+      emitAssert(emitICmp(arith::CmpIPredicate::sgt, arith::CmpIPredicate::ugt,
+                          stepVal, zero),
+                 "splitForOpAtPoint: step must be positive");
+    }
+    // splitPoint == lowerBound + k * step
+    if (lb && step && split) {
+      if ((*split - *lb) % *step != 0)
+        return failure();
+    } else {
+      Value zero = arith::ConstantOp::create(
----------------
matthias-springer wrote:

nit: You can use `getValueOrCreateConstantIntOp(getIntegerAttr(...))`

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