[MLIR] Add parallel loop coalescing.

[tpopp]

Summary

This allows conversion of a ParallelLoop from N induction variables to
some nuber of induction variables less than N.

The first intended use of this is for the GPUDialect to convert
ParallelLoops to iterate over 3 dimensions so they can be launched as
GPU Kernels.

To implement this:
- Normalize each iteration space of the ParallelLoop
- Use the same induction variable in a new ParallelLoop for multiple
  original iterations.
- Split the new induction variable back into the original set of values
  inside the body of the ParallelLoop.

[tpopp]

Refactored some code to be more proper.

[tpopp]

Save loops.getLoc in a variable and use variable everywhere instead.

[lebedev.ri]

Would it *please* be possible for MLIR patches to be named as such, please?

[rriddle]

(Just some drive-by nits)

mlir/include/mlir/InitAllPasses.h

115

Let's keep these sorted.

mlir/include/mlir/Transforms/LoopUtils.h

229

Can you please add a comment here?

mlir/lib/Transforms/ParallelLoopCoalescing.cpp

23 ↗

(On Diff #251099)

Can you use pass options instead? https://mlir.llvm.org/docs/WritingAPass/#instance-specific-pass-options

mlir/lib/Transforms/Utils/LoopUtils.cpp

976–983

nit: Use /// for top-level comments

1000–1001

nit: Please drop all trivial braces.

1032

nit: Please use /// for top-level comments.

1119

nit: Cache the end iterator of the loop, and prefer pre-increment.

1152

Same here and below.

[bondhugula]

Could you please add a summary to the commit message - even if it's a couple of lines? On a side note, do you want to use the term 'collapse' instead of 'coalesce'? OpenMP uses collapse for such linearization, and coalesce could also imply fusion of loops, which this isn't. I do see that colaesceLoops already existed prior to this patch

mlir/lib/Transforms/ParallelLoopCoalescing.cpp

52 ↗

(On Diff #251099)

List initialize?

[mehdi_amini]

In D76363#1929217, @lebedev.ri wrote:

Would it *please* be possible for MLIR patches to be named as such, please?

Where is this coming from?
It'd be nice to clearly understand the goal and where is it coming from. In particular I don't understand why this isn't the role of separate tooling and it has to be encoded directly in commit messages. Are we going to prefix down to the last component manually? Is something like [LLVM][Backend][X86][Register allocator] to be expected?

[mehdi_amini]

In D76363#1929476, @bondhugula wrote:

Could you please add a summary to the commit message - even if it's a couple of lines?

+1, we document it here: https://mlir.llvm.org/getting_started/Contributing/#commit-messages

[tpopp]

In D76363#1930039, @mehdi_amini wrote:

In D76363#1929476, @bondhugula wrote:

Could you please add a summary to the commit message - even if it's a couple of lines?

+1, we document it here: https://mlir.llvm.org/getting_started/Contributing/#commit-messages

I'll add to the summary. Is there a reason arc diff didn't pick up my git commit message? Does it only create the summary based on the first arc diff and not any subsequent calls?

[tpopp]

Handle formatting and naming feedback.

[tpopp]

Use OptionList instead of llvm:🆑:list

[herhut]

Thanks. Looks great with some nits.

mlir/lib/Transforms/ParallelLoopCoalescing.cpp

44 ↗

(On Diff #251099)

Can you make this an OperationPass instead?

52 ↗

(On Diff #251099)

Use llvm::SmallVector instead?

mlir/lib/Transforms/Utils/LoopUtils.cpp

979

Maybe have a little struct here instead of a tuple? Or use std::tie at use sites to improve readability.

983

Move these closer to their first use.

1006

Maybe Value newLowerBound = isZeroBased ? lowerBound : boundsBuilder.create<ConstantIndexOp>(loc, 0)?

1012

Here, too?

1050

Mega-nit: The order lower, step, upper is strange...

1138

Why not newUpperBound = cst1 here?

1154

A comment what this computes would help readability.

1158

Should this be the normalized upper bound?

1160

It would read easier for me if updating previous was also done here except for the last case. Would that make sense?

1163

Normalized here, too?

[tpopp]

Handle herhut's comments and fix broken collapsing
logic that used the wrong upper bound value.

[tpopp]

mlir/lib/Transforms/Utils/LoopUtils.cpp

1138

No real reason. I thought it would be easier for debugging purposes if each string of calculations is fully unconnected from other calculations.

1158

Yes

1163

Yes

[tpopp]

mlir/lib/Transforms/Utils/LoopUtils.cpp

1160

I think this trades one mess for a different one because then it's just a different bounds check and not all indexing is happening at ivar_idx anymore.

[tpopp]

Change loop form inside of collapsePLoops.

[tpopp]

mlir/lib/Transforms/Utils/LoopUtils.cpp

1160

I tried to restructure it to be more readable.

[tpopp]

Use correct variable to fix undefined variable error.

[tpopp]

Rename variable for clang tidy reasons.

[herhut]

Thanks.

[bondhugula]

In D76363#1930727, @tpopp wrote:

In D76363#1930039, @mehdi_amini wrote:

In D76363#1929476, @bondhugula wrote:

Could you please add a summary to the commit message - even if it's a couple of lines?

+1, we document it here: https://mlir.llvm.org/getting_started/Contributing/#commit-messages

I'll add to the summary. Is there a reason arc diff didn't pick up my git commit message? Does it only create the summary based on the first arc diff and not any subsequent calls?

I've noticed this too. It doesn't pick up subsequent commit message updates. I just manually update the commit message here when needed.

[bondhugula]

Mostly minor suggestions on readability

mlir/lib/Transforms/Utils/LoopUtils.cpp

1045

OpBuilder innerBuilder(inner.getBody()) will be sufficient.

1109

But to spell this out? PLoops -> ParallelLoops?

1130

Nit: period at the end.

mlir/test/Transforms/parallel-loop-collapsing.mlir

5–14 ↗

(On Diff #252767)

Can you drop the extra indent? Also VAL_0 -> C6, VAL_1 -> C7, ...?

31 ↗

(On Diff #252767)

Can you use more descriptive names? VAL_10 -> IV0, VAL_11 -> IV1, ...

34 ↗

(On Diff #252767)

Drop the additional indent between the CHECK and the string?

49 ↗

(On Diff #252767)

CHECK-NEXT

50–51 ↗

(On Diff #252767)

CHECK-NEXT

mlir/test/Transforms/single-parallel-loop-collapsing.mlir

26 ↗

(On Diff #252767)

VAL_14 isn't used; no need to capture it.

32–33 ↗

(On Diff #252767)

CHECK-NEXT

34 ↗

(On Diff #252767)

Not needed.

35 ↗

(On Diff #252767)

Drop trailing blank lines.

mlir/include/mlir/InitAllPasses.h

#endif
   createLinalgTilingToParallelLoopsPass();
   createLinalgPromotionPass(0);
   createConvertLinalgToLoopsPass();
   createConvertLinalgToParallelLoopsPass();
   createConvertLinalgToAffineLoopsPass();
   createConvertLinalgToLLVMPass();
 
   // LoopOps
+  createParallelLoopCollapsingPass();
   createParallelLoopFusionPass();
   createParallelLoopSpecializationPass();
   createParallelLoopTilingPass();
 
   // QuantOps
   quant::createConvertSimulatedQuantPass();
   quant::createConvertConstPass();
   quantizer::createAddDefaultStatsPass();
   quantizer::createRemoveInstrumentationPass();
   quantizer::registerInferQuantizedTypesPass();
 
   // SPIR-V

mlir/include/mlir/Transforms/LoopUtils.h

 class AffineForOp;
 class FuncOp;
 class OpBuilder;
 class Value;
 struct MemRefRegion;
 
 namespace loop {
 class ForOp;
+class ParallelOp;
 } // end namespace loop
 
 /// Unrolls this for operation completely if the trip count is known to be
 /// constant. Returns failure otherwise.
 LogicalResult loopUnrollFull(AffineForOp forOp);
 
 /// Unrolls this for operation by the specified unroll factor. Returns failure
 /// if the loop cannot be unrolled either due to restrictions or due to invalid

 TileLoops extractFixedOuterLoops(loop::ForOp rootFOrOp,
                                  ArrayRef<int64_t> sizes);
 
 /// Replace a perfect nest of "for" loops with a single linearized loop. Assumes
 /// `loops` contains a list of perfectly nested loops with bounds and steps
 /// independent of any loop induction variable involved in the nest.
 void coalesceLoops(MutableArrayRef<loop::ForOp> loops);
 
+/// Take the ParallelLoop and for each set of dimension indices, combine them
+/// into a single dimension. combinedDimensions must contain each index into
+/// loops exactly once.
+void collapsePLoops(loop::ParallelOp loops,
+                    std::vector<std::vector<unsigned>> combinedDimensions);
+
 /// Maps `forOp` for execution on a parallel grid of virtual `processorIds` of
 /// size given by `numProcessors`. This is achieved by embedding the SSA values
 /// corresponding to `processorIds` and `numProcessors` into the bounds and step
 /// of the `forOp`. No check is performed on the legality of the rewrite, it is
 /// the caller's responsibility to ensure legality.
 ///
 /// Requires that `processorIds` and `numProcessors` have the same size and that
 /// for each idx, `processorIds`[idx] takes, at runtime, all values between 0

mlir/include/mlir/Transforms/Passes.h

 /// Creates a pass to perform tiling on loop nests.
 std::unique_ptr<OpPassBase<FuncOp>>
 createLoopTilingPass(uint64_t cacheSizeBytes);
 
 /// Creates a pass that transforms perfectly nested loops with independent
 /// bounds into a single loop.
 std::unique_ptr<OpPassBase<FuncOp>> createLoopCoalescingPass();
 
+/// Creates a pass that transforms a single ParallelLoop over N induction
+/// variables into another ParallelLoop over less than N induction variables.
+std::unique_ptr<OpPassBase<FuncOp>> createParallelLoopCollapsingPass();
+
 /// Performs packing (or explicit copying) of accessed memref regions into
 /// buffers in the specified faster memory space through either pointwise copies
 /// or DMA operations.
 std::unique_ptr<OpPassBase<FuncOp>> createAffineDataCopyGenerationPass(
     unsigned slowMemorySpace, unsigned fastMemorySpace,
     unsigned tagMemorySpace = 0, int minDmaTransferSize = 1024,
     uint64_t fastMemCapacityBytes = std::numeric_limits<uint64_t>::max());
 

mlir/lib/Transforms/CMakeLists.txt

add_mlir_library(MLIRTransforms
   LoopCoalescing.cpp
   LoopFusion.cpp
   LoopInvariantCodeMotion.cpp
   LoopTiling.cpp
   LoopUnrollAndJam.cpp
   LoopUnroll.cpp
   MemRefDataFlowOpt.cpp
   OpStats.cpp
+  ParallelLoopCollapsing.cpp
   PipelineDataTransfer.cpp
   SimplifyAffineStructures.cpp
   StripDebugInfo.cpp
   SymbolDCE.cpp
   Vectorize.cpp
   ViewOpGraph.cpp
   ViewRegionGraph.cpp
 

mlir/lib/Transforms/Utils/LoopUtils.cpp

 replaceAllUsesExcept(Value orig, Value replacement,
                      const SmallPtrSetImpl<Operation *> &exceptions) {
   for (auto &use : llvm::make_early_inc_range(orig.getUses())) {
     if (exceptions.count(use.getOwner()) == 0)
       use.set(replacement);
   }
 }
 
-// Transform a loop with a strictly positive step
-//   for %i = %lb to %ub step %s
-// into a 0-based loop with step 1
-//   for %ii = 0 to ceildiv(%ub - %lb, %s) step 1 {
-//     %i = %ii * %s + %lb
-// Insert the induction variable remapping in the body of `inner`, which is
-// expected to be either `loop` or another loop perfectly nested under `loop`.
-// Insert the definition of new bounds immediate before `outer`, which is
-// expected to be either `loop` or its parent in the loop nest.
-static void normalizeLoop(loop::ForOp loop, loop::ForOp outer,
-                          loop::ForOp inner) {
-  OpBuilder builder(outer);
-  Location loc = loop.getLoc();
-
+/// Return the new lower bound, upper bound, and step in that order. Insert any
+/// additional bounds calculations before the given builder and any additional
+/// conversion back to the original loop induction value inside the given Block.
+static std::tuple<Value, Value, Value>
+normalizeLoop(OpBuilder &boundsBuilder, OpBuilder &insideLoopBuilder,
+              Location loc, Value lowerBound, Value upperBound, Value step,
+              Value inductionVar) {
+  Value newLowerBound, newUpperBound, newStep;
   // Check if the loop is already known to have a constant zero lower bound or
   // a constant one step.
   bool isZeroBased = false;
   if (auto ubCst =
-          dyn_cast_or_null<ConstantIndexOp>(loop.lowerBound().getDefiningOp()))
+          dyn_cast_or_null<ConstantIndexOp>(lowerBound.getDefiningOp()))
     isZeroBased = ubCst.getValue() == 0;
 
   bool isStepOne = false;
-  if (auto stepCst =
-          dyn_cast_or_null<ConstantIndexOp>(loop.step().getDefiningOp()))
+  if (auto stepCst = dyn_cast_or_null<ConstantIndexOp>(step.getDefiningOp()))
     isStepOne = stepCst.getValue() == 1;
 
-  if (isZeroBased && isStepOne)
-    return;
 
   // Compute the number of iterations the loop executes: ceildiv(ub - lb, step)
   // assuming the step is strictly positive.  Update the bounds and the step
   // of the loop to go from 0 to the number of iterations, if necessary.
   // TODO(zinenko): introduce support for negative steps or emit dynamic asserts
   // on step positivity, whatever gets implemented first.
-  Value diff =
-      builder.create<SubIOp>(loc, loop.upperBound(), loop.lowerBound());
-  Value numIterations = ceilDivPositive(builder, loc, diff, loop.step());
-  loop.setUpperBound(numIterations);
-
-  Value lb = loop.lowerBound();
-  if (!isZeroBased) {
-    Value cst0 = builder.create<ConstantIndexOp>(loc, 0);
-    loop.setLowerBound(cst0);
-  }
+  if (isZeroBased && isStepOne)
+    return {lowerBound, upperBound, step};
 
-  Value step = loop.step();
-  if (!isStepOne) {
-    Value cst1 = builder.create<ConstantIndexOp>(loc, 1);
-    loop.setStep(cst1);
-  }
+  Value diff = boundsBuilder.create<SubIOp>(loc, upperBound, lowerBound);
+  newUpperBound = ceilDivPositive(boundsBuilder, loc, diff, step);
+
+  if (isZeroBased)
+    newLowerBound = lowerBound;
+  else
+    newLowerBound = boundsBuilder.create<ConstantIndexOp>(loc, 0);
+
+  if (isStepOne)
+    newStep = step;
+  else
+    newStep = boundsBuilder.create<ConstantIndexOp>(loc, 1);
 
   // Insert code computing the value of the original loop induction variable
   // from the "normalized" one.
-  builder.setInsertionPointToStart(inner.getBody());
   Value scaled =
-      isStepOne ? loop.getInductionVar()
-                : builder.create<MulIOp>(loc, loop.getInductionVar(), step);
+      isStepOne ? inductionVar
+                : insideLoopBuilder.create<MulIOp>(loc, inductionVar, step);
   Value shifted =
-      isZeroBased ? scaled : builder.create<AddIOp>(loc, scaled, lb);
+      isZeroBased ? scaled
+                  : insideLoopBuilder.create<AddIOp>(loc, scaled, lowerBound);
 
   SmallPtrSet<Operation *, 2> preserve{scaled.getDefiningOp(),
                                        shifted.getDefiningOp()};
-  replaceAllUsesExcept(loop.getInductionVar(), shifted, preserve);
+  replaceAllUsesExcept(inductionVar, shifted, preserve);
+  return {newLowerBound, newUpperBound, newStep};
+}
+
+/// Transform a loop with a strictly positive step
+///   for %i = %lb to %ub step %s
+/// into a 0-based loop with step 1
+///   for %ii = 0 to ceildiv(%ub - %lb, %s) step 1 {
+///     %i = %ii * %s + %lb
+/// Insert the induction variable remapping in the body of `inner`, which is
+/// expected to be either `loop` or another loop perfectly nested under `loop`.
+/// Insert the definition of new bounds immediate before `outer`, which is
+/// expected to be either `loop` or its parent in the loop nest.
+static void normalizeLoop(loop::ForOp loop, loop::ForOp outer,
+                          loop::ForOp inner) {
+  OpBuilder builder(outer);
+  OpBuilder innerBuilder(inner.getBody(), inner.getBody()->begin());
+  auto loopPieces =
+      normalizeLoop(builder, innerBuilder, loop.getLoc(), loop.lowerBound(),
+                    loop.upperBound(), loop.step(), loop.getInductionVar());
+
+  loop.setLowerBound(std::get<0>(loopPieces));
+  loop.setStep(std::get<2>(loopPieces));
+  loop.setUpperBound(std::get<1>(loopPieces));
 }
 
 void mlir::coalesceLoops(MutableArrayRef<loop::ForOp> loops) {
   if (loops.size() < 2)
     return;
 
   loop::ForOp innermost = loops.back();
   loop::ForOp outermost = loops.front();
void mlir::coalesceLoops(MutableArrayRef<loop::ForOp> loops) {
   loop::ForOp second = loops[1];
   innermost.getBody()->back().erase();
   outermost.getBody()->getOperations().splice(
       Block::iterator(second.getOperation()),
       innermost.getBody()->getOperations());
   second.erase();
 }
 
+void mlir::collapsePLoops(
+    loop::ParallelOp loops,
+    std::vector<std::vector<unsigned>> combinedDimensions) {
+  OpBuilder outsideBuilder(loops);
+  Location loc = loops.getLoc();
+
+  // Normalize ParallelOp's iteration pattern.
+  SmallVector<Value, 3> normalizedLowerBounds;
+  SmallVector<Value, 3> normalizedSteps;
+  SmallVector<Value, 3> normalizedUpperBounds;
+  for (unsigned i = 0, e = loops.getNumLoops(); i < e; ++i) {
+    OpBuilder insideLoopBuilder(loops.getBody(), loops.getBody()->begin());
+    auto resultBounds =
+        normalizeLoop(outsideBuilder, insideLoopBuilder, loc,
+                      loops.lowerBound()[i], loops.upperBound()[i],
+                      loops.step()[i], loops.getBody()->getArgument(i));
+
+    normalizedLowerBounds.push_back(std::get<0>(resultBounds));
+    normalizedUpperBounds.push_back(std::get<1>(resultBounds));
+    normalizedSteps.push_back(std::get<2>(resultBounds));
+  }
+
+  // Combine iteration spaces
+  SmallVector<Value, 3> lowerBounds;
+  SmallVector<Value, 3> steps;
+  SmallVector<Value, 3> upperBounds;
+  auto cst0 = outsideBuilder.create<ConstantIndexOp>(loc, 0);
+  auto cst1 = outsideBuilder.create<ConstantIndexOp>(loc, 1);
+  for (unsigned i = 0, e = combinedDimensions.size(); i < e; ++i) {
+    Value newUpperBound = outsideBuilder.create<ConstantIndexOp>(loc, 1);
+    for (auto idx : combinedDimensions[i]) {
+      newUpperBound = outsideBuilder.create<MulIOp>(loc, newUpperBound,
+                                                    normalizedUpperBounds[idx]);
+    }
+    lowerBounds.push_back(cst0);
+    steps.push_back(cst1);
+    upperBounds.push_back(newUpperBound);
+  }
+
+  // Create new ParallelLoop with conversions to the original induction values.
+  auto newPloop = outsideBuilder.create<loop::ParallelOp>(loc, lowerBounds,
+                                                          upperBounds, steps);
+  OpBuilder insideBuilder(newPloop.getBody(), newPloop.getBody()->begin());
+  for (unsigned i = 0, e = combinedDimensions.size(); i < e; ++i) {
+    Value previous = newPloop.getBody()->getArgument(i);
+    for (unsigned idx = 0, e = combinedDimensions[i].size(); idx < e; ++idx) {
+      unsigned ivar_idx = combinedDimensions[i][idx];
+      if (idx != 0)
+        previous = insideBuilder.create<SignedDivIOp>(
+            loc, previous, loops.upperBound()[ivar_idx]);
+
+      Value iv = (idx == e - 1)
+                     ? previous
+                     : insideBuilder.create<SignedRemIOp>(
+                           loc, previous, loops.upperBound()[ivar_idx]);
+      replaceAllUsesInRegionWith(loops.getBody()->getArgument(ivar_idx), iv,
+                                 loops.region());
+    }
+  }
+
+  // Replace the old loop with the new loop.
+  loops.getBody()->back().erase();
+  newPloop.getBody()->getOperations().splice(
+      Block::iterator(newPloop.getBody()->back()),
+      loops.getBody()->getOperations());
+  loops.erase();
+}
+
 void mlir::mapLoopToProcessorIds(loop::ForOp forOp, ArrayRef<Value> processorId,
                                  ArrayRef<Value> numProcessors) {
   assert(processorId.size() == numProcessors.size());
   if (processorId.empty())
     return;
 
   OpBuilder b(forOp);
   Location loc(forOp.getLoc());