[Mlir-commits] [mlir] 8a56a73 - [MLIR] Add output argument to affineParallelize utility

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2023-05-29T21:11:00+05:30
New Revision: 8a56a730f26c52eae38b7614edcc6c37ea033f48

URL: https://github.com/llvm/llvm-project/commit/8a56a730f26c52eae38b7614edcc6c37ea033f48
DIFF: https://github.com/llvm/llvm-project/commit/8a56a730f26c52eae38b7614edcc6c37ea033f48.diff

LOG: [MLIR] Add output argument to affineParallelize utility

Add output argument to affineParallelize utility. NFC.

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Affine/Utils.h
    mlir/lib/Dialect/Affine/Utils/Utils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Affine/Utils.h b/mlir/include/mlir/Dialect/Affine/Utils.h
index 8e54a02a89105..ca52f1771737f 100644
--- a/mlir/include/mlir/Dialect/Affine/Utils.h
+++ b/mlir/include/mlir/Dialect/Affine/Utils.h
@@ -44,10 +44,11 @@ using ReductionLoopMap = DenseMap<Operation *, SmallVector<LoopReduction, 2>>;
 /// (mlir::isLoopParallel can be used to detect a parallel affine.for op.) The
 /// reductions specified in `parallelReductions` are also parallelized.
 /// Parallelization will fail in the presence of loop iteration arguments that
-/// are not listed in `parallelReductions`.
-LogicalResult
-affineParallelize(AffineForOp forOp,
-                  ArrayRef<LoopReduction> parallelReductions = {});
+/// are not listed in `parallelReductions`. `resOp` if non-null is set to the
+/// newly created affine.parallel op.
+LogicalResult affineParallelize(AffineForOp forOp,
+                                ArrayRef<LoopReduction> parallelReductions = {},
+                                AffineParallelOp *resOp = nullptr);
 
 /// Hoists out affine.if/else to as high as possible, i.e., past all invariant
 /// affine.fors/parallel's. Returns success if any hoisting happened; folded` is

diff  --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 4e02b612b9bfe..d567093188e0c 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -344,7 +344,8 @@ static AffineIfOp hoistAffineIfOp(AffineIfOp ifOp, Operation *hoistOverOp) {
 
 LogicalResult
 mlir::affine::affineParallelize(AffineForOp forOp,
-                                ArrayRef<LoopReduction> parallelReductions) {
+                                ArrayRef<LoopReduction> parallelReductions,
+                                AffineParallelOp *resOp) {
   // Fail early if there are iter arguments that are not reductions.
   unsigned numReductions = parallelReductions.size();
   if (numReductions != forOp.getNumIterOperands())
@@ -398,6 +399,8 @@ mlir::affine::affineParallelize(AffineForOp forOp,
   newPloop.getBody()->eraseArguments(numIVs, numReductions);
 
   forOp.erase();
+  if (resOp)
+    *resOp = newPloop;
   return success();
 }