[MLIR] Add output argument to affineParallelize utility

[bondhugula]

Summary

Add output argument to affineParallelize utility. NFC.

[vimal-polymagelabs]

LGTM

mlir/include/mlir/Dialect/Affine/Utils.h

 
 using ReductionLoopMap = DenseMap<Operation *, SmallVector<LoopReduction, 2>>;
 
 /// Replaces a parallel affine.for op with a 1-d affine.parallel op. `forOp`'s
 /// body is taken by the affine.parallel op and the former is erased.
 /// (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
 /// set to true if the op was folded or erased. This hoisting could lead to
 /// significant code expansion in some cases.
 LogicalResult hoistAffineIfOp(AffineIfOp ifOp, bool *folded = nullptr);
 
 /// Holds parameters to perform n-D vectorization on a single loop nest.

mlir/lib/Dialect/Affine/Utils/Utils.cpp

elseBlock->getOperations().splice(
       elseBlock->begin(), hoistOverOpClone->getBlock()->getOperations(),
       Block::iterator(hoistOverOpClone));
 
   return hoistedIfOp;
 }
 
 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())
     return failure();
 
   Location loc = forOp.getLoc();
   OpBuilder outsideBuilder(forOp);
   AffineMap lowerBoundMap = forOp.getLowerBoundMap();
mlir::affine::affineParallelize(AffineForOp forOp,
   // operation itself (now moved outside of the loop) and erase the block
   // arguments that correspond to reductions. Note that the loop always has one
   // "main" induction variable whenc coming from a non-parallel for.
   unsigned numIVs = 1;
   yieldOp->setOperands(reducedValues);
   newPloop.getBody()->eraseArguments(numIVs, numReductions);
 
   forOp.erase();
+  if (resOp)
+    *resOp = newPloop;
   return success();
 }
 
 // Returns success if any hoisting happened.
 LogicalResult mlir::affine::hoistAffineIfOp(AffineIfOp ifOp, bool *folded) {
   // Bail out early if the ifOp returns a result.  TODO: Consider how to
   // properly support this case.
   if (ifOp.getNumResults() != 0)