[Mlir-commits] [mlir] [flang] [mlir][TilingInterface] Use `LoopLikeOpInterface` in tiling using SCF to unify tiling with `scf.for` and `scf.forall`. (PR #77874)

[Nicolas Vasilache]

================
@@ -584,6 +588,63 @@ ForOp::replaceWithAdditionalYields(RewriterBase &rewriter,
   return cast<LoopLikeOpInterface>(newLoop.getOperation());
 }
 
+FailureOr<LoopLikeOpInterface> ForOp::yieldTiledValuesAndReplace(
+    RewriterBase &rewriter, ValueRange newInitOperands,
+    const YieldTiledValuesFn &yieldTiledValuesFn) {
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(getOperation());
+
+  auto inits = llvm::to_vector(getInitArgs());
+  inits.append(newInitOperands.begin(), newInitOperands.end());
+  auto newLoop = rewriter.create<ForOp>(
+      getLoc(), getLowerBound(), getUpperBound(), getStep(), inits,
+      [](OpBuilder &, Location, Value, ValueRange) {});
+
+  // Move the loop body to the new op.
+  rewriter.mergeBlocks(getBody(), newLoop.getBody(),
+                       newLoop.getBody()->getArguments().take_front(
+                           getBody()->getNumArguments()));
+
+  auto yieldOp = cast<scf::YieldOp>(newLoop.getBody()->getTerminator());
+  rewriter.setInsertionPoint(yieldOp);
+
+  SmallVector<Value> tiledValues;
+  SmallVector<SmallVector<OpFoldResult>> resultOffsets, resultSizes;
+  ValueRange newRegionIterArgs =
+      newLoop.getRegionIterArgs().take_back(newInitOperands.size());
+  if (failed(yieldTiledValuesFn(rewriter, getLoc(), newLoop.getInductionVar(),
+                                newRegionIterArgs, tiledValues, resultOffsets,
+                                resultSizes))) {
+    return rewriter.notifyMatchFailure(getOperation(),
----------------
nicolasvasilache wrote:

Careful here as we have already created IR, so depending how this is used (today and in the future) we may or may not violate invariants (e.g. rewriter returns failure but IR has been modified).

Not 100% sure this is/will be a problem but:
1. Consider cleaning up the newLoop here ?
2. Also consider adding documentation to the lambda (e.g. either it must succeed or clean up after itself)?

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