[mlir] Return scf.parallel ops resulted from tiling.

Differential Revision: https://reviews.llvm.org/D96024

diff --git a/mlir/include/mlir/Dialect/SCF/Transforms.h b/mlir/include/mlir/Dialect/SCF/Transforms.h
index 3164d33..456eb4e 100644 (file)
--- a/mlir/include/mlir/Dialect/SCF/Transforms.h
+++ b/mlir/include/mlir/Dialect/SCF/Transforms.h
@@ -44,7 +44,11 @@ void naivelyFuseParallelOps(Region &region);
 ///                                           min(tileSize[1], %arg3-%j1))
 ///                                        step (%arg4, %arg5)
 /// The old loop is replaced with the new one.
-void tileParallelLoop(ParallelOp op, llvm::ArrayRef<int64_t> tileSizes);
+///
+/// The function returns the resulting ParallelOps, i.e. {outer_loop_op,
+/// inner_loop_op}.
+std::pair<ParallelOp, ParallelOp>
+tileParallelLoop(ParallelOp op, llvm::ArrayRef<int64_t> tileSizes);
 
 /// Populates patterns for SCF structural type conversions and sets up the
 /// provided ConversionTarget with the appropriate legality configuration for

diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp
index 542c1c2..cdb4afa 100644 (file)
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp
@@ -37,7 +37,8 @@ using namespace mlir::scf;
 /// %i0 + j0 and %i1 + %j1.
 //
 /// The old loop is replaced with the new one.
-void mlir::scf::tileParallelLoop(ParallelOp op, ArrayRef<int64_t> tileSizes) {
+std::pair<ParallelOp, ParallelOp>
+mlir::scf::tileParallelLoop(ParallelOp op, ArrayRef<int64_t> tileSizes) {
   OpBuilder b(op);
   auto zero = b.create<ConstantIndexOp>(op.getLoc(), 0);
   SmallVector<Value, 2> tileSizeConstants;
@@ -125,6 +126,7 @@ void mlir::scf::tileParallelLoop(ParallelOp op, ArrayRef<int64_t> tileSizes) {
   }
 
   op.erase();
+  return std::make_pair(outerLoop, innerLoop);
 }
 
 namespace {