From 8a56a730f26c52eae38b7614edcc6c37ea033f48 Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday@polymagelabs.com>
Date: Mon, 29 May 2023 08:52:40 +0530
Subject: [PATCH] [MLIR] Add output argument to affineParallelize utility

Add output argument to affineParallelize utility. NFC.

Differential Revision: https://reviews.llvm.org/D151636
---
 mlir/include/mlir/Dialect/Affine/Utils.h | 9 +++++----
 mlir/lib/Dialect/Affine/Utils/Utils.cpp  | 5 ++++-
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Affine/Utils.h b/mlir/include/mlir/Dialect/Affine/Utils.h
index 8e54a02..ca52f17 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 4e02b61..d567093 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();
 }
 
-- 
2.7.4