[mlir] parallel loop canonicalization
authorTobias Gysi <tobias.gysi@inf.ethz.ch>
Fri, 26 Jun 2020 07:47:01 +0000 (09:47 +0200)
committerTobias Gysi <tobias.gysi@inf.ethz.ch>
Fri, 26 Jun 2020 07:57:08 +0000 (09:57 +0200)
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.

Reviewers: herhut, ftynse

Reviewed By: herhut

Subscribers: mehdi_amini, rriddle, jpienaar, shauheen, antiagainst, nicolasvasilache, arpith-jacob, mgester, lucyrfox, aartbik, liufengdb, stephenneuendorffer, Joonsoo, grosul1, Kayjukh, jurahul, msifontes

Tags: #mlir

Differential Revision: https://reviews.llvm.org/D82191

mlir/include/mlir/Dialect/SCF/SCFOps.td
mlir/lib/Dialect/SCF/SCF.cpp
mlir/test/Dialect/SCF/canonicalize.mlir [new file with mode: 0644]
mlir/test/Transforms/parallel-loop-collapsing.mlir

index 4205831..863f8d4 100644 (file)
@@ -346,6 +346,8 @@ def ParallelOp : SCF_Op<"parallel",
     unsigned getNumLoops() { return step().size(); }
     unsigned getNumReductions() { return initVals().size(); }
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 def ReduceOp : SCF_Op<"reduce", [HasParent<"ParallelOp">]> {
index 4484fdf..a85977d 100644 (file)
@@ -8,17 +8,8 @@
 
 #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/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/BlockAndValueMapping.h"
 #include "mlir/IR/PatternMatch.h"
-#include "mlir/IR/StandardTypes.h"
-#include "mlir/IR/Value.h"
-#include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/InliningUtils.h"
 
 using namespace mlir;
@@ -742,6 +733,68 @@ ParallelOp mlir::scf::getParallelForInductionVarOwner(Value val) {
   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 dim : llvm::zip(op.lowerBound(), op.upperBound(), op.step(),
+                              op.getInductionVars())) {
+      Value lowerBound, upperBound, step, iv;
+      std::tie(lowerBound, upperBound, step, iv) = dim;
+      // Collect the statically known loop bounds.
+      auto lowerBoundConstant =
+          dyn_cast_or_null<ConstantIndexOp>(lowerBound.getDefiningOp());
+      auto upperBoundConstant =
+          dyn_cast_or_null<ConstantIndexOp>(upperBound.getDefiningOp());
+      auto stepConstant =
+          dyn_cast_or_null<ConstantIndexOp>(step.getDefiningOp());
+      // Replace the loop induction variable by the lower bound if the loop
+      // performs a single iteration. Otherwise, copy the loop bounds.
+      if (lowerBoundConstant && upperBoundConstant && stepConstant &&
+          (upperBoundConstant.getValue() - lowerBoundConstant.getValue()) > 0 &&
+          (upperBoundConstant.getValue() - lowerBoundConstant.getValue()) <=
+              stepConstant.getValue()) {
+        mapping.map(iv, lowerBound);
+      } else {
+        newLowerBounds.push_back(lowerBound);
+        newUpperBounds.push_back(upperBound);
+        newSteps.push_back(step);
+      }
+    }
+    // Exit if all or none of the loop dimensions perform 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);
+    // Clone the loop body and remap the block arguments of the collapsed loops
+    // (inlining does not support a cancellable block argument mapping).
+    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
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
new file mode 100644 (file)
index 0000000..fc98dab
--- /dev/null
@@ -0,0 +1,55 @@
+// 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
+
+// -----
+
+func @no_iteration(%A: memref<?x?xi32>) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  scf.parallel (%i0, %i1) = (%c0, %c0) to (%c1, %c0) step (%c1, %c1) {
+    %c42 = constant 42 : i32
+    store %c42, %A[%i0, %i1] : memref<?x?xi32>
+    scf.yield
+  }
+  return
+}
+
+// CHECK-LABEL:   func @no_iteration(
+// CHECK-SAME:                        [[ARG0:%.*]]: memref<?x?xi32>) {
+// CHECK:           [[C0:%.*]] = constant 0 : index
+// CHECK:           [[C1:%.*]] = constant 1 : index
+// CHECK:           [[C42:%.*]] = constant 42 : i32
+// CHECK:           scf.parallel ([[V1:%.*]]) = ([[C0]]) to ([[C0]]) step ([[C1]]) {
+// CHECK:             store [[C42]], [[ARG0]]{{\[}}[[C0]], [[V1]]] : memref<?x?xi32>
+// CHECK:             scf.yield
+// CHECK:           }
+// CHECK:           return
index a124d18..36b6a2c 100644 (file)
@@ -16,6 +16,8 @@ func @parallel_many_dims() {
   %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) {
@@ -26,24 +28,18 @@ func @parallel_many_dims() {
 
 // 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