[MLIR] Add lowering for affine.parallel to scf.parallel

[yash]

Summary

[MLIR] adding affine parallel conversion to scf parallel in lower affine pass

[rriddle]

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

369

nit: expandAffineMap can fail, please check error conditions.

374

nit: Just use a for-each loop over op.steps() instead.

[yash]

Added check file.

[yash]

Minor Updates

[bondhugula]

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

360

Nit: use backticks instead of quotes.

370–385

This is all lacking comments - please add some.

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

626–630

The test cases are insufficient - please include one with non-trivial affine map bounds.

[yash]

Added non trivial cases for test file

[yash]

Minor update.

[yash]

Commented and used for-each loop over op.steps()

[bondhugula]

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

377

Please avoid names like step1. Also, avoid auto here.

[bondhugula]

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

373

Still not checking if expandAffineMap returns no value, here and below.

[bondhugula]

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

675

No need to capture A5, it's unused.

[bondhugula]

Can you please update the commit title to "[MLIR] Add lowering for affine.parallel to scf.parallel"? (Updating the commit via arc won't automatically update the title here - please hit edit here to update as well.)

[yash]

Checking error conditions for expandAffineMap

[yash]

Changing commit title

[yash]

Checking error conditions for expandAffineMap.

[bondhugula]

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

379

lowerBoundValue should actually be called lowerBound and vice versa. Likewise for upper.

385

If the expandAffineMap gives you no value there, you would just be passing empty lowerBound and upperBound. And this could crash.

[yash]

Renamed variable and put asserts.

[bondhugula]

Ping @ftynse - do you have feedback here?

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

640

Can you use a better / more descriptive name than dot1?

[ftynse]

Please address the comments before submitting.

@bondhugula, sorry, I thought this was relanding of a rolled back change

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

377

Nit: upperBoundValue sounds like it's a _single_ value whereas is a vector of values. This got me confused down the line where the loop is created from upperBoundValue, lowerBoundValue (singular) and steps (plural).

382

Please call steps.reserve before calling push_back in a loop when you know the number of elements to be inserted.

383–384

Nit: would it be possible to construct a ConstantOp and pass it an attribute without extracting its value? Something like rewriter.create<ConstantOp>(loc, step)? Loop bounds are already of index type.

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

640

+1. Also for @dot, I assumed it was a dot product.

[yash]

Renamed from dot to affine_parallel.
Renamed from lowerBoundValue to lowerBoundTuple to make it plural.
Used steps.reserve.
Not able to use rewriter.create<ConstantOp>(loc, step).

[yash]

Comment.

[ftynse]

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

380

Actually, do you expect it to never fail (it will fail on quasi-affine maps, for example)? Otherwise, you can just return failure() from the pattern.

[bondhugula]

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

380

Did you mean quasi-affine or semi-affine? expandAffineMap returns an Optional but doesn't document when it is expected to have a value nor is it obvious from either the declaration or even immediately from the implementation! Goes to show how important documentation is! Do you know who the author of this is - could we have the doc comment updated? Return values like these must be documented without fail.

[bondhugula]

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

380

In theory, we could expand any affine map, but I'm guessing semi-affine maps (where the RHS of the floordiv/ceildiv/mod is a symbol) are perhaps unsupported by expandAffineMap and so this would fail in those cases. So, for completeness sake, we should return failure for now and add a comment that in theory, it's always guaranteed to succeed if expandAffineMap is extended.

[yash]

Return failure if expandAffineMap fails.

[yash]

Minor update.

[bondhugula]

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

675

Nit: You can drop the %{{.*}} as well.

[yash]

Minor change.

mlir/include/mlir/Conversion/AffineToStandard/AffineToStandard.h

 #define MLIR_CONVERSION_AFFINETOSTANDARD_AFFINETOSTANDARD_H
 
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
 class AffineExpr;
 class AffineForOp;
 class AffineMap;
+class AffineParallelOp;
 class Location;
 struct LogicalResult;
 class MLIRContext;
 class OpBuilder;
 class RewritePattern;
 class Value;
 class ValueRange;
 

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

LogicalResult matchAndRewrite(AffineForOp op,
     auto f = rewriter.create<scf::ForOp>(loc, lowerBound, upperBound, step);
     rewriter.eraseBlock(f.getBody());
     rewriter.inlineRegionBefore(op.region(), f.region(), f.region().end());
     rewriter.eraseOp(op);
     return success();
   }
 };
 
+/// Convert an `affine.parallel` (loop nest) operation into a `scf.parallel`
+/// operation.
+class AffineParallelLowering : public OpRewritePattern<AffineParallelOp> {
+public:
+  using OpRewritePattern<AffineParallelOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AffineParallelOp op,
+                                PatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    SmallVector<Value, 8> steps;
+    SmallVector<Value, 8> upperBoundTuple;
+    SmallVector<Value, 8> lowerBoundTuple;
+    // Finding lower and upper bound by expanding the map expression.
+    // Checking if expandAffineMap is not giving NULL.
+    Optional<SmallVector<Value, 8>> upperBound = expandAffineMap(
+        rewriter, loc, op.upperBoundsMap(), op.getUpperBoundsOperands());
+    Optional<SmallVector<Value, 8>> lowerBound = expandAffineMap(
+        rewriter, loc, op.lowerBoundsMap(), op.getLowerBoundsOperands());
+    if (!lowerBound || !upperBound)
+      return failure();
+    upperBoundTuple = *upperBound;
+    lowerBoundTuple = *lowerBound;
+    steps.reserve(op.steps().size());
+    for (Attribute step : op.steps())
+      steps.push_back(rewriter.create<ConstantIndexOp>(
+          loc, step.cast<IntegerAttr>().getInt()));
+    // Creating empty scf.parallel op body with appropriate bounds.
+    auto parallelOp = rewriter.create<scf::ParallelOp>(loc, lowerBoundTuple,
+                                                       upperBoundTuple, steps);
+    rewriter.eraseBlock(parallelOp.getBody());
+    rewriter.inlineRegionBefore(op.region(), parallelOp.region(),
+                                parallelOp.region().end());
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
 class AffineIfLowering : public OpRewritePattern<AffineIfOp> {
 public:
   using OpRewritePattern<AffineIfOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(AffineIfOp op,
                                 PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
 
void mlir::populateAffineToStdConversionPatterns(
   // clang-format off
   patterns.insert<
       AffineApplyLowering,
       AffineDmaStartLowering,
       AffineDmaWaitLowering,
       AffineLoadLowering,
       AffineMinLowering,
       AffineMaxLowering,
+      AffineParallelLowering,
       AffinePrefetchLowering,
       AffineStoreLowering,
       AffineForLowering,
       AffineIfLowering,
       AffineYieldOpLowering>(ctx);
   // clang-format on
 }
 

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

func @affine_max(%arg0: index, %arg1: index) -> index{
   // CHECK: %[[Cm2:.*]] = constant -1
   // CHECK: %[[neg2:.*]] = muli %[[ARG0]], %[[Cm2:.*]]
   // CHECK: %[[second:.*]] = addi %[[ARG1]], %[[neg2]]
   // CHECK: %[[cmp:.*]] = cmpi "sgt", %[[first]], %[[second]]
   // CHECK: select %[[cmp]], %[[first]], %[[second]]
   %0 = affine.max affine_map<(d0,d1) -> (d0 - d1, d1 - d0)>(%arg0, %arg1)
   return %0 : index
 }
+
+// CHECK-LABEL: func @affine_parallel(
+// CHECK-SAME: %[[ARG0:.*]]: memref<100x100xf32>, %[[ARG1:.*]]: memref<100x100xf32>) {
+func @affine_parallel(%o: memref<100x100xf32>, %a: memref<100x100xf32>) {
+  affine.parallel (%i, %j) = (0, 0) to (100, 100) {
+  }
+  return
+}
+
+// CHECK-DAG:    %[[C100:.*]] = constant 100
+// CHECK-DAG:    %[[C100_1:.*]] = constant 100
+// CHECK-DAG:    %[[C0:.*]] = constant 0
+// CHECK-DAG:    %[[C0_1:.*]] = constant 0
+// CHECK-DAG:    %[[C1:.*]] = constant 1
+// CHECK-DAG:    %[[C1_1:.*]] = constant 1
+// CHECK-DAG:    scf.parallel (%arg2, %arg3) = (%[[C0]], %[[C0_1]]) to (%[[C100]], %[[C100_1]]) step (%[[C1]], %[[C1_1]]) {
+
+// CHECK-LABEL: func @affine_parallel_tiled(
+// CHECK-SAME: %[[ARG0:.*]]: memref<100x100xf32>, %[[ARG1:.*]]: memref<100x100xf32>, %[[ARG2:.*]]: memref<100x100xf32>) {
+func @affine_parallel_tiled(%o: memref<100x100xf32>, %a: memref<100x100xf32>, %b: memref<100x100xf32>) {
+  affine.parallel (%i0, %j0, %k0) = (0, 0, 0) to (100, 100, 100) step (10, 10, 10) {
+    affine.parallel (%i1, %j1, %k1) = (%i0, %j0, %k0) to (%i0 + 10, %j0 + 10, %k0 + 10) {
+      %0 = affine.load %a[%i1, %k1] : memref<100x100xf32>
+      %1 = affine.load %b[%k1, %j1] : memref<100x100xf32>
+      %2 = mulf %0, %1 : f32
+    }
+  }
+  return
+}
+
+// CHECK-DAG:     %[[C100:.*]] = constant 100
+// CHECK-DAG:     %[[C100_0:.*]] = constant 100
+// CHECK-DAG:     %[[C100_1:.*]] = constant 100
+// CHECK-DAG:     %[[C0:.*]] = constant 0
+// CHECK-DAG:     %[[C0_2:.*]] = constant 0
+// CHECK-DAG:     %[[C0_3:.*]] = constant 0
+// CHECK-DAG:     %[[C10:.*]] = constant 10
+// CHECK-DAG:     %[[C10_4:.*]] = constant 10
+// CHECK-DAG:     %[[C10_5:.*]] = constant 10
+// CHECK:         scf.parallel (%[[arg3:.*]], %[[arg4:.*]], %[[arg5:.*]]) = (%[[C0]], %[[C0_2]], %[[C0_3]]) to (%[[C100]], %[[C100_0]], %[[C100_1]]) step (%[[C10]], %[[C10_4]], %[[C10_5]]) {
+// CHECK-DAG:       %[[C10_6:.*]] = constant 10
+// CHECK-DAG:       %[[A0:.*]] = addi %[[arg3]], %[[C10_6]]
+// CHECK-DAG:       %[[C10_7:.*]] = constant 10
+// CHECK-DAG:       %[[A1:.*]] = addi %[[arg4]], %[[C10_7]]
+// CHECK-DAG:       %[[C10_8:.*]] = constant 10
+// CHECK-DAG:       %[[A2:.*]] = addi %[[arg5]], %[[C10_8]]
+// CHECK-DAG:       %[[C1:.*]] = constant 1
+// CHECK-DAG:       %[[C1_9:.*]] = constant 1
+// CHECK-DAG:       %[[C1_10:.*]] = constant 1
+// CHECK:           scf.parallel (%[[arg6:.*]], %[[arg7:.*]], %[[arg8:.*]]) = (%[[arg3]], %[[arg4]], %[[arg5]]) to (%[[A0]], %[[A1]], %[[A2]]) step (%[[C1]], %[[C1_9]], %[[C1_10]]) {
+// CHECK:             %[[A3:.*]] = load %[[ARG1]][%[[arg6]], %[[arg8]]] : memref<100x100xf32>
+// CHECK:             %[[A4:.*]] = load %[[ARG2]][%[[arg8]], %[[arg7]]] : memref<100x100xf32>
+// CHECK:             mulf %[[A3]], %[[A4]] : f32
+// CHECK:             scf.yield