[MLIR][SCF] Add loops as parameter to LoopTerminator callback when using CustomOp. (#161386)

This PR adds to the generateLoopTerminatorFn callback the loops
generated by GenerateLoopHeaderFn. This is needed to correctly set the
insertion point with scf.forall ops.

Author: sebvince

mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h

@@ -183,6 +183,7 @@ struct SCFTilingOptions {
       ArrayRef<OpFoldResult> givenTileSizes, ValueRange destinationTensors)>;
 
   // Type of the callback function that generates the loop terminator.
+  // - `loops` : generated loops from the GenerateLoopHeaderFn callback
   // - `tiledResults` : Tiles of the result computed for the iteration space
   //                    tile.
   // - `resultOffsets` : For each of the `tiledResults`, the offset at which
@@ -193,7 +194,8 @@ struct SCFTilingOptions {
   //                   tensor.
   // Returns the `CustomLoopHeaderInfo` object (described above)
   using GenerateLoopTerminatorFn = std::function<LogicalResult(
-      RewriterBase &rewriter, Location loc, ValueRange tiledResults,
+      RewriterBase &rewriter, Location loc, ArrayRef<LoopLikeOpInterface> loops,
+      ValueRange tiledResults,
       ArrayRef<SmallVector<OpFoldResult>> resultOffsets,
       ArrayRef<SmallVector<OpFoldResult>> resultSizes,
       ValueRange destinationTensors)>;

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

@@ -665,8 +665,8 @@ generateLoopNestUsingCustomOp(
     return failure();
   }
 
-  if (failed(generateLoopTerminatorFn(rewriter, loc, tiledResults,
-                                      resultOffsets, resultSizes,
+  if (failed(generateLoopTerminatorFn(rewriter, loc, loopHeaderInfo->loops,
+                                      tiledResults, resultOffsets, resultSizes,
                                       loopHeaderInfo->destinationTensors))) {
     return failure();
   }

mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp

@@ -581,7 +581,8 @@ DiagnosedSilenceableFailure transform::TestTileUsingCustomLoopOp::apply(
   };
 
   scf::SCFTilingOptions::GenerateLoopTerminatorFn terminatorFn =
-      [&](RewriterBase &rewriter, Location loc, ValueRange tiledResults,
+      [&](RewriterBase &rewriter, Location loc,
+          ArrayRef<LoopLikeOpInterface> loops, ValueRange tiledResults,
           ArrayRef<SmallVector<OpFoldResult>> resultOffsets,
           ArrayRef<SmallVector<OpFoldResult>> resultSizes,
           ValueRange destinationTensors) -> LogicalResult {