[mlir] parallel loop canonicalization

[gysit]

Summary

The patch introduces a canonicalization pattern for parallel loops. The pattern removes single-iteration loop dimensions if the loop bounds and steps are constants.

[herhut]

Thank you for splitting this out! Just some nits but otherwise this looks great! Thanks for adding this.

mlir/lib/Dialect/SCF/SCF.cpp

728

nit: You can also use op.getInductionVars()which makes the intent clearer.

mlir/test/Dialect/SCF/canonicalize.mlir

11

It would be great to add the case of an empty iteration space. Either with another copy that does not get rewritten or by introducing one in this example.

[gysit]

mlir/test/Dialect/SCF/canonicalize.mlir

11

Does it make sense to have a canonicalization pattern that remove loops with an empty iteration space entirely. I was circumventing this bcs I was not sure if I can simply replace all uses of the loop result values by the init values. And to be honest, I also do not believe it is a very relevant "optimization".

[gysit]

• added additional test

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

23

I don't think most of these are necessary.

755

Why clone instead of inline?

[gysit]

• cleanup includes

[gysit]

mlir/lib/Dialect/SCF/SCF.cpp

23

You are right.

755

I didn't find an inline signature that takes a mapping parameter. Luckily the clone is in a code path that will not execute often... If there is a better solution for the entire block argument mapping I am keen to learn it.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

755

It isn't clear to me that you need a mapping here, instead of mapping.map(std::get<3>(bounds), std::get<0>(bounds)); couldn't you just do std::get<3>(bounds).replaceAllUsesWith(std::get<0>(bounds)) and inline the region?
(you may have to remap other block args, but it may be handled by the inline region code)

Alternatively, we should extend the inlineRegion methods with one that fits the needs here.

[gysit]

mlir/lib/Dialect/SCF/SCF.cpp

755

My understanding (from an earlier review) is that it is not a good idea to use replaceAllUsesWith with the PatternRewriter since it cannot track these changes (which is needed to rollback?). But I may be wrong there and things may have changed since this discussion. Let me know if there is a misunderstanding on my side.

Having an inline signature with a mapping sounds like a useful extension to me!

[gysit]

mlir/lib/Dialect/SCF/SCF.cpp

755

I learned canonicalization patterns are never reverted on success. I thus assume I can use replaceAllUsesWith but only after line 738 where the pattern may return failure()? Unfortunately I also need to erase all remapped arguments. The following code works for me:

for (auto iv : llvm::enumerate(llvm::reverse(op.getInductionVars()))) {
  if (auto lowerBound = mapping.lookupOrNull(iv.value())) {
    iv.value().replaceAllUsesWith(lowerBound);
    op.getBody()->eraseArgument(op.getNumLoops() - iv.index() - 1);
  }
}
rewriter.inlineRegionBefore(op.region(), newOp.region(),
                            newOp.region().begin());

If I do not erase the remapped arguments I get an invalid loop with the original number of induction variables / block arguments. Let me know if there is a nicer way to do this and if I should update the patch.

[ftynse]

mlir/lib/Dialect/SCF/SCF.cpp

755

I'd say they are _currently_ never reverted on success if you run the canonicalization pass. If somebody decides to put this pattern into the legalization infra (e.g., applyAnalysisConversion) which can use these patterns just fine, they can be reverted...

That being said, if you use eraseArgument, you can use replaceAllUsesWith and the rest of in-place modifications because none of those would be compatible with the rollback...

@rriddle will be the most knowledgeable here

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

755

I would not rely on the fact that they aren't currently reverted. All patterns generally have the same restrictions irregardless of the driver, i.e. you should not really be mutating outside of the pattern rewriter. If you code patterns, especially canonicalization patterns, to the quirks of a specific driver it breaks the ability to use it anywhere else and it also makes it much much more difficult if anyone ever wants to change the way the driver works in the future.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

755

That's interesting: limiting canonicalization to "cancellable" transformations may incur extra costs (as we see here): do we have an RAUW exposed in a way that the driver could change the behavior to have efficient canonicalizations? (there may be tricky interactions though, I'm not sure how it'd play out)

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

755

It's not just about cancellable or not, even the greedy driver keeps a mapping of things which can easily become invalidated. There have been many instances of crashes/asan failures because a pattern is doing something without informing the driver. For dialect conversion it is largely about cancellability, for greedy it is largely about allowing the driver to remove mappings for things that are being invalidated.

[gysit]

• Added a comment that reflects the discussions about clone (can be reverted) vs inline (performance).
• The current implementation is revertible at the cost of cloning the loop body.
• Use std::tie to improve the code readability.

[herhut]

I think the conclusion is that we want it to be cancelable. So this is ready to land.

mlir/test/Dialect/SCF/canonicalize.mlir

11

I agree it is not very relevant but yes, you could replace the results with the init values in such case.

mlir/include/mlir/Dialect/SCF/SCFOps.td

def ParallelOp : SCF_Op<"parallel",
 
   let extraClassDeclaration = [{
     ValueRange getInductionVars() {
       return getBody()->getArguments();
     }
     unsigned getNumLoops() { return step().size(); }
     unsigned getNumReductions() { return initVals().size(); }
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 def ReduceOp : SCF_Op<"reduce", [HasParent<"ParallelOp">]> {
   let summary = "reduce operation for parallel for";
   let description = [{
     "scf.reduce" is an operation occurring inside "scf.parallel" operations.
     It consists of one block with two arguments which have the same type as the
     operand of "scf.reduce".

mlir/lib/Dialect/SCF/SCF.cpp

 //===- SCF.cpp - Structured Control Flow Operations -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/Value.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/MathExtras.h"
 
 using namespace mlir;
 using namespace mlir::scf;
 
 //===----------------------------------------------------------------------===//
 // SCFDialect
 //===----------------------------------------------------------------------===//
ParallelOp mlir::scf::getParallelForInductionVarOwner(Value val) {
   auto ivArg = val.dyn_cast<BlockArgument>();
   if (!ivArg)
     return ParallelOp();
   assert(ivArg.getOwner() && "unlinked block argument");
   auto *containingOp = ivArg.getOwner()->getParentOp();
   return dyn_cast<ParallelOp>(containingOp);
 }
 
+namespace {
+// Collapse loop dimensions that perform a single iteration.
+struct CollapseSingleIterationLoops : public OpRewritePattern<ParallelOp> {
+  using OpRewritePattern<ParallelOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ParallelOp op,
+                                PatternRewriter &rewriter) const override {
+    BlockAndValueMapping mapping;
+    // Compute new loop bounds that omit all single-iteration loop dimensions.
+    SmallVector<Value, 2> newLowerBounds;
+    SmallVector<Value, 2> newUpperBounds;
+    SmallVector<Value, 2> newSteps;
+    newLowerBounds.reserve(op.lowerBound().size());
+    newUpperBounds.reserve(op.upperBound().size());
+    newSteps.reserve(op.step().size());
+    for (auto bounds : llvm::zip(op.lowerBound(), op.upperBound(), op.step(),
+                                 op.getBody()->getArguments())) {
+      auto lowerBound = dyn_cast_or_null<ConstantIndexOp>(
+          std::get<0>(bounds).getDefiningOp());
+      auto upperBound = dyn_cast_or_null<ConstantIndexOp>(
+          std::get<1>(bounds).getDefiningOp());
+      auto step = dyn_cast_or_null<ConstantIndexOp>(
+          std::get<2>(bounds).getDefiningOp());
+      // Replace the loop induction variable by the lower bound if the loop
+      // performs a single iteration. Otherwise, copy the loop bounds.
+      if (lowerBound && upperBound && step &&
+          (upperBound.getValue() - lowerBound.getValue()) > 0 &&
+          (upperBound.getValue() - lowerBound.getValue()) <= step.getValue()) {
+        mapping.map(std::get<3>(bounds), std::get<0>(bounds));
+      } else {
+        newLowerBounds.push_back(std::get<0>(bounds));
+        newUpperBounds.push_back(std::get<1>(bounds));
+        newSteps.push_back(std::get<2>(bounds));
+      }
+    }
+    // Exit if all or none of the loop dimensions performs a single iteration.
+    if (newLowerBounds.size() == 0 ||
+        newLowerBounds.size() == op.lowerBound().size())
+      return failure();
+    // Replace the parallel loop by lower-dimensional parallel loop.
+    auto newOp =
+        rewriter.create<ParallelOp>(op.getLoc(), newLowerBounds, newUpperBounds,
+                                    newSteps, op.initVals(), nullptr);
+    rewriter.cloneRegionBefore(op.region(), newOp.region(),
+                               newOp.region().begin(), mapping);
+    rewriter.replaceOp(op, newOp.getResults());
+    return success();
+  }
+};
+} // namespace
+
+void ParallelOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+                                             MLIRContext *context) {
+  results.insert<CollapseSingleIterationLoops>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // ReduceOp
 //===----------------------------------------------------------------------===//
 
 void ReduceOp::build(
     OpBuilder &builder, OperationState &result, Value operand,
     function_ref<void(OpBuilder &, Location, Value, Value)> bodyBuilderFn) {
   auto type = operand.getType();

mlir/test/Dialect/SCF/canonicalize.mlir

+// RUN: mlir-opt %s -pass-pipeline='func(canonicalize)' | FileCheck %s
+
+func @single_iteration(%A: memref<?x?x?xi32>) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c3 = constant 3 : index
+  %c6 = constant 6 : index
+  %c7 = constant 7 : index
+  %c10 = constant 10 : index
+  scf.parallel (%i0, %i1, %i2) = (%c0, %c3, %c7) to (%c1, %c6, %c10) step (%c1, %c2, %c3) {
+    %c42 = constant 42 : i32
+    store %c42, %A[%i0, %i1, %i2] : memref<?x?x?xi32>
+    scf.yield
+  }
+  return
+}
+
+// CHECK-LABEL:   func @single_iteration(
+// CHECK-SAME:                        [[ARG0:%.*]]: memref<?x?x?xi32>) {
+// CHECK:           [[C0:%.*]] = constant 0 : index
+// CHECK:           [[C2:%.*]] = constant 2 : index
+// CHECK:           [[C3:%.*]] = constant 3 : index
+// CHECK:           [[C6:%.*]] = constant 6 : index
+// CHECK:           [[C7:%.*]] = constant 7 : index
+// CHECK:           [[C42:%.*]] = constant 42 : i32
+// CHECK:           scf.parallel ([[V0:%.*]]) = ([[C3]]) to ([[C6]]) step ([[C2]]) {
+// CHECK:             store [[C42]], [[ARG0]]{{\[}}[[C0]], [[V0]], [[C7]]] : memref<?x?x?xi32>
+// CHECK:             scf.yield
+// CHECK:           }
+// CHECK:           return

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

func @parallel_many_dims() {
   %c7 = constant 7 : index
   %c8 = constant 8 : index
   %c9 = constant 9 : index
   %c10 = constant 10 : index
   %c11 = constant 11 : index
   %c12 = constant 12 : index
   %c13 = constant 13 : index
   %c14 = constant 14 : index
+  %c15 = constant 15 : index
+  %c26 = constant 26 : index
 
   scf.parallel (%i0, %i1, %i2, %i3, %i4) = (%c0, %c3, %c6, %c9, %c12) to (%c2, %c5, %c8, %c11, %c14)
                                           step (%c1, %c4, %c7, %c10, %c13) {
     %result = "magic.op"(%i0, %i1, %i2, %i3, %i4): (index, index, index, index, index) -> index
   }
   return
 }
 
 // CHECK-LABEL: func @parallel_many_dims() {
 // CHECK:         [[C6:%.*]] = constant 6 : index
-// CHECK:         [[C7:%.*]] = constant 7 : index
 // CHECK:         [[C9:%.*]] = constant 9 : index
 // CHECK:         [[C10:%.*]] = constant 10 : index
 // CHECK:         [[C12:%.*]] = constant 12 : index
-// CHECK:         [[C13:%.*]] = constant 13 : index
-// CHECK:         [[C3:%.*]] = constant 3 : index
 // CHECK:         [[C0:%.*]] = constant 0 : index
 // CHECK:         [[C1:%.*]] = constant 1 : index
 // CHECK:         [[C2:%.*]] = constant 2 : index
-// CHECK:         scf.parallel ([[NEW_I0:%.*]], [[NEW_I1:%.*]], [[NEW_I2:%.*]]) = ([[C0]], [[C0]], [[C0]]) to ([[C2]], [[C1]], [[C1]]) step ([[C1]], [[C1]], [[C1]]) {
+// CHECK:         [[C3:%.*]] = constant 3 : index
+// CHECK:         scf.parallel ([[NEW_I0:%.*]]) = ([[C0]]) to ([[C2]]) step ([[C1]]) {
 // CHECK:           [[I0:%.*]] = remi_signed [[NEW_I0]], [[C2]] : index
-// CHECK:           [[VAL_16:%.*]] = muli [[NEW_I1]], [[C13]] : index
-// CHECK:           [[I4:%.*]] = addi [[VAL_16]], [[C12]] : index
-// CHECK:           [[VAL_18:%.*]] = muli [[NEW_I0]], [[C10]] : index
-// CHECK:           [[I3:%.*]] = addi [[VAL_18]], [[C9]] : index
-// CHECK:           [[VAL_20:%.*]] = muli [[NEW_I2]], [[C7]] : index
-// CHECK:           [[I2:%.*]] = addi [[VAL_20]], [[C6]] : index
-// CHECK:           "magic.op"([[I0]], [[C3]], [[I2]], [[I3]], [[I4]]) : (index, index, index, index, index) -> index
+// CHECK:           [[V18:%.*]] = muli [[NEW_I0]], [[C10]] : index
+// CHECK:           [[I3:%.*]] = addi [[V18]], [[C9]] : index
+// CHECK:           "magic.op"([[I0]], [[C3]], [[C6]], [[I3]], [[C12]]) : (index, index, index, index, index) -> index
 // CHECK:           scf.yield
 // CHECK-NEXT:    }
 // CHECK-NEXT:    return