Fix Comments

Author: mmarjieh

mlir/include/mlir/Dialect/SCF/Transforms/Passes.td

@@ -128,7 +128,7 @@ def SCFParallelForToNestedFors : Pass<"scf-parallel-for-to-nested-fors"> {
   let summary = "Convert SCF parallel for loops to nested SCF for loops";
   let constructor = "mlir::createParallelForToNestedForsPass()";
   let description = [{
-    This pass transforms SCF.ParallelOp operations into a nest of SCF.ForOp
+    This pass transforms SCF::ParallelOp operations into a nest of SCF::ForOp
     operations. The transformation is useful for cases where the parallel loop
     can be expressed as a series of sequential iterations, allowing for more
     fine-grained control over the loop execution.

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

@@ -13,6 +13,7 @@
 #ifndef MLIR_DIALECT_SCF_TRANSFORMS_TRANSFORMS_H_
 #define MLIR_DIALECT_SCF_TRANSFORMS_TRANSFORMS_H_
 
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -43,10 +44,9 @@ LogicalResult forallToParallelLoop(RewriterBase &rewriter, ForallOp forallOp,
                                    ParallelOp *result = nullptr);
 
 /// Try converting scf.forall into an scf.parallel loop.
-/// The conversion is only supported for forall operations with no results.
-LogicalResult parallelForToNestedFors(RewriterBase &rewriter,
-                                      ParallelOp parallelOp,
-                                      ForOp *result = nullptr);
+/// The conversion is only supported for parallel operations with no results.
+FailureOr<scf::LoopNest> parallelForToNestedFors(RewriterBase &rewriter,
+                                                 ParallelOp parallelOp);
 
 /// Fuses all adjacent scf.parallel operations with identical bounds and step
 /// into one scf.parallel operations. Uses a naive aliasing and dependency

mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp

@@ -176,14 +176,15 @@ DiagnosedSilenceableFailure transform::ParallelForToNestedForOps::apply(
     return diag;
   }
 
-  scf::ForOp opResult;
-  if (failed(scf::parallelForToNestedFors(rewriter, target, &opResult))) {
+  FailureOr<scf::LoopNest> loopNest =
+      scf::parallelForToNestedFors(rewriter, target);
+  if (failed(loopNest)) {
     DiagnosedSilenceableFailure diag =
         emitSilenceableError() << "failed to convert parallel into nested fors";
     return diag;
   }
 
-  results.set(cast<OpResult>(getTransformed()[0]), {opResult});
+  results.set(cast<OpResult>(getTransformed()[0]), {loopNest->loops.front()});
   return DiagnosedSilenceableFailure::success();
 }
 

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

@@ -23,51 +23,40 @@ namespace mlir {
 
 using namespace mlir;
 
-LogicalResult mlir::scf::parallelForToNestedFors(RewriterBase &rewriter,
-                                                 scf::ParallelOp parallelOp,
-                                                 scf::ForOp *result) {
+FailureOr<scf::LoopNest>
+mlir::scf::parallelForToNestedFors(RewriterBase &rewriter,
+                                   scf::ParallelOp parallelOp) {
 
   if (!parallelOp.getResults().empty()) {
-    parallelOp->emitError("Currently ScfParallel to ScfFor conversion "
-                          "doesn't support ScfParallel with results.");
+    parallelOp->emitError("Currently scf.parallel to scf.for conversion "
+                          "doesn't support scf.parallel with results.");
     return failure();
   }
 
   rewriter.setInsertionPoint(parallelOp);
 
   Location loc = parallelOp.getLoc();
-  auto lowerBounds = parallelOp.getLowerBound();
-  auto upperBounds = parallelOp.getUpperBound();
-  auto steps = parallelOp.getStep();
+  SmallVector<Value> lowerBounds = parallelOp.getLowerBound();
+  SmallVector<Value> upperBounds = parallelOp.getUpperBound();
+  SmallVector<Value> steps = parallelOp.getStep();
 
   assert(lowerBounds.size() == upperBounds.size() &&
          lowerBounds.size() == steps.size() &&
          "Mismatched parallel loop bounds");
 
   SmallVector<Value> ivs;
-  auto loopNest =
+  scf::LoopNest loopNest =
       scf::buildLoopNest(rewriter, loc, lowerBounds, upperBounds, steps);
 
-  auto oldInductionVars = parallelOp.getInductionVars();
-  auto newInductionVars = llvm::map_to_vector(
+  SmallVector<Value> newInductionVars = llvm::map_to_vector(
       loopNest.loops, [](scf::ForOp forOp) { return forOp.getInductionVar(); });
-  assert(oldInductionVars.size() == newInductionVars.size() &&
-         "Mismatched induction variables");
-  for (auto [oldIV, newIV] : llvm::zip(oldInductionVars, newInductionVars))
-    oldIV.replaceAllUsesWith(newIV);
-
-  auto *linearizedBody = loopNest.loops.back().getBody();
-  Block &parallelBody = *parallelOp.getBody();
-  for (Operation &op : llvm::make_early_inc_range(parallelBody)) {
-    // Skip the terminator of the parallelOp body.
-    if (&op == parallelBody.getTerminator())
-      continue;
-    op.moveBefore(linearizedBody->getTerminator());
-  }
+  Block *linearizedBody = loopNest.loops.back().getBody();
+  Block *parallelBody = parallelOp.getBody();
+  rewriter.eraseOp(parallelBody->getTerminator());
+  rewriter.inlineBlockBefore(parallelBody, linearizedBody->getTerminator(),
+                             newInductionVars);
   rewriter.eraseOp(parallelOp);
-  if (result)
-    *result = loopNest.loops.front();
-  return success();
+  return loopNest;
 }
 
 namespace {

mlir/test/Dialect/SCF/parallel-to-nested-fors.mlir

@@ -67,7 +67,7 @@ func.func private @callee(%i: index, %j: index) -> i32
 
 func.func @two_iters_with_reduce(%lb1: index, %lb2: index, %ub1: index, %ub2: index, %step1: index, %step2: index) -> i32 {
   %c0 = arith.constant 0 : i32
-  // expected-error@+1 {{Currently ScfParallel to ScfFor conversion doesn't support ScfParallel with results}}
+  // expected-error@+1 {{Currently scf.parallel to scf.for conversion doesn't support scf.parallel with results}}
   %0 = scf.parallel (%i, %j) = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) init (%c0) -> i32 {
     %curr = func.call @callee(%i, %j) : (index, index) -> i32
     scf.reduce(%curr : i32) {
@@ -77,4 +77,4 @@ func.func @two_iters_with_reduce(%lb1: index, %lb2: index, %ub1: index, %ub2: in
     }
   }
   return %0 : i32
-}
+}