[mlir][linalg] Extend tiled_loop to SCF conversion to generate scf.parallel.

[pifon2a]

remove new line

[pifon2a]

Update variable names.

[tpopp]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

736–744

At this point, the op should be limited to being entirely bufferized or not at all bufferized. Can you add this check to the verify method also? It could be as a followup as there's no very method currently.

I say this because partial bufferization along with this method would be very confusing and could easily lead to incorrect assumptions/bugs

mlir/lib/Dialect/Linalg/Transforms/Loops.cpp

498–499

clang tidy checks

519

This system relies on the assumption that a tiled loop with iteration patterns [parallel, sequential, parallel, sequential] can be reordered to [parallel, parallel, sequential, sequential]. Is that true and specified?

533

Maybe use builder here. This is relying on the caller of this never creating a new OpBuilfer from the rewriter passed to it (relying on two objects actually being the same). That could be easy to break and annoying to debug.

535–536

I don't think this is enough. If one op uses the result of another op in the body, I think it will still point to the original op rather than the new one.

ie
%0 = add %a, %b
%1 = negate %0

Will result in
%new0 = add %a, %b
%new1 = negate %0

Won't it?

[pifon2a]

Address the comments.

[pifon2a]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

736–744

Unfortunately, comprehensive bufferization relies on tiled_loop with mixed buffers-tensors. It adds bufferized memref operands to the tiled_loop without removing the tensor ones and then canonicalization kicks in to clean this up. It does this because it has to be in-place.

mlir/lib/Dialect/Linalg/Transforms/Loops.cpp

519

I think it is true for the tiling that we support, since we are generating a perfectly nested loop where loop ctrl vars don't depend on each other. I will add a comment.

533

good catch, thanks.

535–536

builder.clone(op, bvm) also maps old results to new ones.

[tpopp]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

736–744

What do you think about a name that specifies the "only" fact? hasNoTensorSemantics or hasOnlyBufferSemantics.

[pifon2a]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

736–744

hasBufferSemantics is what's used by mlir/Dialect/Linalg/IR/LinalgInterfaces.td

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

OpOperand* findInputOperand(Value value) {
       return it.getBase();
     }
 
     OpOperand* findOutputOperand(Value value) {
       OperandRange::iterator it = llvm::find(outputs(), value);
       if (it == outputs().end()) return nullptr;
       return it.getBase();
     }
+
+    /// Return whether the op has only MemRef input and outputs.
+    bool hasBufferSemantics() {
+      Operation* op = this->getOperation();
+      return op->getNumResults() == 0 &&
+             llvm::all_of(op->getOpOperands(), [&](OpOperand & operand) {
+               return !operand.get().getType().template isa<ShapedType>() ||
+                      operand.get().getType().template isa<MemRefType>();
+             });
+    }
+
+    /// Return whether the loop dimension is parallel or not.
+    bool isParallelDimension(unsigned dim) {
+      StringAttr attr = this->iterator_types()[dim].cast<StringAttr>();
+      return attr.getValue() == getParallelIteratorTypeName();
+    }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 def Linalg_IndexOp : Linalg_Op<"index", [NoSideEffect]>,
     Arguments<(ins Confined<I64Attr, [IntMinValue<0>]>:$dim)>,

mlir/lib/Dialect/Linalg/Transforms/Loops.cpp

if (!isa<LinalgOp>(op))
       return failure();
     if (!linalgOpToLoopsImpl<LoopType>(rewriter, linalgOp))
       return failure();
     rewriter.eraseOp(op);
     return success();
   }
 };
 
+/// Converts tiled_loop to SCF loop nests. All parallel dimensions are collected
+/// into an scf.parallel loop and all sequential dimensions will result in the
+/// nested scf.for loop nest. The pattern assumes that a tiled loop with
+/// iterator_types ["reduction", "parallel", "reduction"] can be reordered. It
+/// is true for the tiling that is currently suppported by Linalg.
 struct TiledLoopToSCFPattern : public OpRewritePattern<TiledLoopOp> {
   using OpRewritePattern<TiledLoopOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(TiledLoopOp tiledLoop,
                                 PatternRewriter &rewriter) const override {
-    Location loc = tiledLoop.getLoc();
-
     // Fail conversion if the `tiled_loop` has not been bufferized.
-    if (!llvm::all_of(tiledLoop.outputs(), [&](Value arg) {
-          return arg.getType().isa<MemRefType>();
-        }))
+    if (!tiledLoop.hasBufferSemantics())
       return failure();
 
-    // TODO: Build loop nest with `scf.for` and `scf.parallel` depending on the
-    // iterator type.
-    scf::buildLoopNest(rewriter, loc, tiledLoop.lowerBound(),
-                       tiledLoop.upperBound(), tiledLoop.step(),
-                       [&](OpBuilder &builder, Location loc, ValueRange ivs) {
-                         // Move body without its terminator.
-                         SmallVector<Value> newBlockArgs;
-                         newBlockArgs.append(ivs.begin(), ivs.end());
-                         newBlockArgs.append(tiledLoop.inputs().begin(),
-                                             tiledLoop.inputs().end());
-                         newBlockArgs.append(tiledLoop.outputs().begin(),
-                                             tiledLoop.outputs().end());
-                         Block *newBody = rewriter.getInsertionBlock();
-                         rewriter.mergeBlocks(tiledLoop.getBody(), newBody,
-                                              newBlockArgs);
-                         rewriter.eraseOp(newBody->getTerminator());
-                       });
+    // Collect loop control parameters for parallel and sequential dimensions.
+    SmallVector<Value, 3> seqLBs, seqUBs, seqSteps, seqIVs;
+    SmallVector<Value, 3> parLBs, parUBs, parSteps, parIVs;
+    for (auto en : llvm::enumerate(
+             llvm::zip(tiledLoop.lowerBound(), tiledLoop.upperBound(),
+                       tiledLoop.step(), tiledLoop.getInductionVars()))) {
+      Value lb, ub, step, iv;
+      std::tie(lb, ub, step, iv) = en.value();
+      if (tiledLoop.isParallelDimension(en.index())) {
+        parLBs.push_back(lb);
+        parUBs.push_back(ub);
+        parSteps.push_back(step);
+        parIVs.push_back(iv);
+      } else {
+        seqLBs.push_back(lb);
+        seqUBs.push_back(ub);
+        seqSteps.push_back(step);
+        seqIVs.push_back(iv);
+      }
+    }
+
+    Location loc = tiledLoop.getLoc();
+    auto generateForLoopNestAndCloneBody = [&](OpBuilder &builder, Location loc,
+                                               ValueRange ivs) {
+      BlockAndValueMapping bvm;
+      bvm.map(parIVs, ivs);
+      bvm.map(tiledLoop.getRegionInputArgs(), tiledLoop.inputs());
+      bvm.map(tiledLoop.getRegionOutputArgs(), tiledLoop.outputs());
+
+      // If not all dimensions of the tiled loop are parallel, an scf.for loop
+      // nest is generated.
+      if (!seqIVs.empty()) {
+        scf::LoopNest nest =
+            scf::buildLoopNest(builder, loc, seqLBs, seqUBs, seqSteps,
+                               [&](OpBuilder &builder, Location loc,
+                                   ValueRange ivs) { bvm.map(seqIVs, ivs); });
+        builder.setInsertionPointToStart(nest.loops.back().getBody());
+      }
+      for (auto &op : tiledLoop.getBody()->without_terminator())
+        builder.clone(op, bvm);
+    };
+
+    if (parIVs.empty())
+      generateForLoopNestAndCloneBody(rewriter, loc, llvm::None);
+    else
+      rewriter.create<scf::ParallelOp>(loc, parLBs, parUBs, parSteps,
+                                       generateForLoopNestAndCloneBody);
     rewriter.eraseOp(tiledLoop);
     return success();
   }
 };
 
 /// Local folding pattern for AffineApplyOp that we can apply greedily.
 /// This replaces AffineApplyOp by the proper value in cases where the
 /// associated map is trivial.

mlir/test/Dialect/Linalg/tiled-loop-to-scf.mlir

+// RUN: mlir-opt %s -convert-linalg-tiled-loops-to-scf --split-input-file | FileCheck %s
+
+
+#map0 = affine_map<(d0) -> (24, -d0 + 192)>
+#map1 = affine_map<(d0, d1)[s0] -> (d0 * 192 + s0 + d1)>
+#map2 = affine_map<(d0) -> (16, -d0 + 192)>
+
+func @tiled_loop(%A: memref<192x192xf32>,
+                 %B: memref<192x192xf32>,
+                 %C: memref<192x192xf32>) {
+  %cst = constant 0.000000e+00 : f32
+  %c24 = constant 24 : index
+  %c16 = constant 16 : index
+  %c0 = constant 0 : index
+  %c192 = constant 192 : index
+
+  linalg.tiled_loop (%i, %j) = (%c0, %c0) to (%c192, %c192) step (%c24, %c16)
+      ins (%A_ = %A: memref<192x192xf32>, %B_ = %B:  memref<192x192xf32>)
+      outs (%C_ = %C: memref<192x192xf32>) {
+    %0 = affine.min #map0(%i)
+    %1 = memref.subview %A_[%i, 0] [%0, 192] [1, 1]
+      : memref<192x192xf32> to memref<?x192xf32, #map1>
+    %2 = affine.min #map2(%j)
+    %3 = memref.subview %B_[0, %j] [192, %2] [1, 1]
+      : memref<192x192xf32> to memref<192x?xf32, #map1>
+    %4 = memref.subview %C_[%i, %j] [%0, %2] [1, 1]
+      : memref<192x192xf32> to memref<?x?xf32, #map1>
+    linalg.fill(%cst, %4) : f32, memref<?x?xf32, #map1>
+    linalg.matmul ins(%1, %3 : memref<?x192xf32, #map1>,
+                               memref<192x?xf32, #map1>)
+                  outs(%4 : memref<?x?xf32, #map1>)
+    linalg.yield
+  }
+  return
+}
+
+// CHECK-LABEL: @tiled_loop
+// CHECK-SAME:  %[[A:.*]]: memref<192x192xf32>, %[[B:.*]]: memref<192x192xf32>,
+// CHECK-SAME:  %[[C:.*]]: memref<192x192xf32>) {
+// CHECK:       %[[C24:.*]] = constant 24 : index
+// CHECK:       %[[C16:.*]] = constant 16 : index
+// CHECK:       %[[C0:.*]] = constant 0 : index
+// CHECK:       %[[C192:.*]] = constant 192 : index
+// CHECK:       scf.parallel (%[[I:.*]], %[[J:.*]]) = (%[[C0]], %[[C0]])
+// CHECK-SAME:      to (%[[C192]], %[[C192]]) step (%[[C24]], %[[C16]]) {
+// CHECK:         %[[A_sub:.*]] = memref.subview %[[A]][%[[I]]
+// CHECK:         %[[B_sub:.*]] = memref.subview %[[B]][0, %[[J]]]
+// CHECK:         %[[C_sub:.*]] = memref.subview %[[C]][%[[I]]
+// CHECK:         linalg.fill
+// CHECK:         linalg.matmul
+
+// -----
+
+func @tiled_loop_reduction(%A: memref<192x192xf32>,
+                           %B: memref<192x192xf32>,
+                           %C: memref<f32>) {
+   %c24 = constant 24 : index
+   %c16 = constant 16 : index
+   %c0 = constant 0 : index
+   %c192 = constant 192 : index
+   %cst = constant 0.000000e+00 : f32
+
+  linalg.tiled_loop (%i, %j) = (%c0, %c0) to (%c192, %c192) step (%c24, %c16)
+      ins (%A_ = %A: memref<192x192xf32>, %B_ = %B:  memref<192x192xf32>)
+      outs (%C_ = %C: memref<f32>)
+      iterators["reduction", "reduction"] {
+    linalg.fill(%cst, %A_) : f32, memref<192x192xf32>
+    linalg.yield
+  }
+  return
+}
+
+// CHECK-LABEL: @tiled_loop_reduction
+// CHECK:       %[[C24:.*]] = constant 24 : index
+// CHECK:       %[[C16:.*]] = constant 16 : index
+// CHECK:       %[[C0:.*]] = constant 0 : index
+// CHECK:       %[[C192:.*]] = constant 192 : index
+// CHECK:       scf.for %{{.*}} = %[[C0]] to %[[C192]] step %[[C24]]
+// CHECK:         scf.for %{{.*}} = %[[C0]] to %[[C192]] step %[[C16]]
+// CHECK:           linalg.fill
+
+// -----
+
+#strided_1d = affine_map<(d0)[s0] -> (d0 + s0)>
+#strided_2d = affine_map<(d0, d1)[s0] -> (d0 * 8 + s0 + d1)>
+
+func @tiled_loop_row_reduction(%A: memref<10x8xf32>,
+                               %B: memref<8xf32>) {
+   %c0 = constant 0 : index
+   %c2 = constant 2 : index
+   %c4 = constant 4 : index
+   %c8 = constant 8 : index
+   %c10 = constant 10 : index
+   %cst = constant 0.000000e+00 : f32
+
+  linalg.tiled_loop (%i, %j) = (%c0, %c0) to (%c10, %c8) step (%c2, %c4)
+      ins (%A_ = %A: memref<10x8xf32>)
+      outs (%B_ = %B: memref<8xf32>)
+      iterators["reduction", "parallel"] {
+    %A_sub = memref.subview %A_[%i, %j][2, 4][1, 1]
+      : memref<10x8xf32> to memref<2x4xf32, #strided_2d>
+    %B_sub = memref.subview %B_[%j][4][1]
+      : memref<8xf32> to memref<4xf32, #strided_1d>
+    linalg.generic {
+        indexing_maps = [affine_map<(i, j) -> (i, j)>,
+                         affine_map<(i, j) -> (j)>],
+        iterator_types = ["reduction", "parallel"]}
+        ins(%A_sub : memref<2x4xf32, #strided_2d>)
+        outs(%B_sub : memref<4xf32, #strided_1d>) {
+      ^bb(%a: f32, %b: f32) :
+        %0 = addf %a, %b: f32
+        linalg.yield %0 : f32
+      }
+    linalg.yield
+  }
+  return
+}
+
+// CHECK-LABEL: @tiled_loop_row_reduction
+
+// CHECK-DAG: %[[C0:.*]] = constant 0 : index
+// CHECK-DAG: %[[C2:.*]] = constant 2 : index
+// CHECK-DAG: %[[C4:.*]] = constant 4 : index
+// CHECK-DAG: %[[C8:.*]] = constant 8 : index
+// CHECK-DAG: %[[C10:.*]] = constant 10 : index
+
+// CHECK:     scf.parallel (%[[J:.*]]) = (%[[C0]]) to (%[[C8]]) step (%[[C4]])
+// CHECK-NEXT:  scf.for %[[I:.*]] = %[[C0]] to %[[C10]] step %[[C2]]
+// CHECK-NEXT:    memref.subview %arg{{[0-9]+}}[%[[I]], %[[J]]] [2, 4] [1, 1]
+// CHECK-SAME:      : memref<10x8xf32> to memref<2x4xf32, #map{{[0-9]+}}>
+// CHECK-NEXT:    memref.subview %arg{{[0-9]+}}[%[[J]]] [4] [1]
+// CHECK-SAME:      : memref<8xf32> to memref<4xf32, #map{{[0-9]+}}>
+
+// -----
+
+#strided_1d = affine_map<(d0)[s0] -> (d0 + s0)>
+#strided_2d = affine_map<(d0, d1)[s0] -> (d0 * 8 + s0 + d1)>
+
+func @tiled_loop_col_reduction(%A: memref<10x8xf32>,
+                               %B: memref<10xf32>) {
+   %c0 = constant 0 : index
+   %c2 = constant 2 : index
+   %c4 = constant 4 : index
+   %c8 = constant 8 : index
+   %c10 = constant 10 : index
+   %cst = constant 0.000000e+00 : f32
+
+  linalg.tiled_loop (%i, %j) = (%c0, %c0) to (%c10, %c8) step (%c2, %c4)
+      ins (%A_ = %A: memref<10x8xf32>)
+      outs (%B_ = %B: memref<10xf32>)
+      iterators["parallel", "reduction"] {
+    %A_sub = memref.subview %A_[%i, %j][2, 4][1, 1]
+      : memref<10x8xf32> to memref<2x4xf32, #strided_2d>
+    %B_sub = memref.subview %B_[%i][2][1]
+      : memref<10xf32> to memref<2xf32, #strided_1d>
+    linalg.generic {
+        indexing_maps = [affine_map<(i, j) -> (i, j)>,
+                         affine_map<(i, j) -> (i)>],
+        iterator_types = ["parallel", "reduction"]}
+        ins(%A_sub : memref<2x4xf32, #strided_2d>)
+        outs(%B_sub : memref<2xf32, #strided_1d>) {
+      ^bb(%a: f32, %b: f32) :
+        %0 = addf %a, %b: f32
+        linalg.yield %0 : f32
+      }
+    linalg.yield
+  }
+  return
+}
+
+// CHECK-LABEL: @tiled_loop_col_reduction
+
+// CHECK-DAG: %[[C0:.*]] = constant 0 : index
+// CHECK-DAG: %[[C2:.*]] = constant 2 : index
+// CHECK-DAG: %[[C4:.*]] = constant 4 : index
+// CHECK-DAG: %[[C8:.*]] = constant 8 : index
+// CHECK-DAG: %[[C10:.*]] = constant 10 : index
+
+// CHECK:     scf.parallel (%[[I:.*]]) = (%[[C0]]) to (%[[C10]]) step (%[[C2]])
+// CHECK-NEXT:  scf.for %[[J:.*]] = %[[C0]] to %[[C8]] step %[[C4]]
+// CHECK-NEXT:    memref.subview %arg{{[0-9]+}}[%[[I]], %[[J]]] [2, 4] [1, 1]
+// CHECK-SAME:      : memref<10x8xf32> to memref<2x4xf32, #map{{[0-9]+}}>
+// CHECK-NEXT:    memref.subview %arg{{[0-9]+}}[%[[I]]] [2] [1]
+// CHECK-SAME:      : memref<10xf32> to memref<2xf32, #map{{[0-9]+}}>

mlir/test/Dialect/Linalg/tiled-loops.mlir

-// RUN: mlir-opt %s -convert-linalg-tiled-loops-to-scf | FileCheck %s
-
-
-#map0 = affine_map<(d0) -> (24, -d0 + 192)>
-#map1 = affine_map<(d0, d1)[s0] -> (d0 * 192 + s0 + d1)>
-#map2 = affine_map<(d0) -> (16, -d0 + 192)>
-
-func @tiled_loop(%A: memref<192x192xf32>,
-                 %B: memref<192x192xf32>,
-                 %C: memref<192x192xf32>) {
-  %cst = constant 0.000000e+00 : f32
-  %c24 = constant 24 : index
-  %c16 = constant 16 : index
-  %c0 = constant 0 : index
-  %c192 = constant 192 : index
-
-  linalg.tiled_loop (%i, %j) = (%c0, %c0) to (%c192, %c192) step (%c24, %c16)
-      ins (%A_ = %A: memref<192x192xf32>, %B_ = %B:  memref<192x192xf32>)
-      outs (%C_ = %C: memref<192x192xf32>) {
-    %0 = affine.min #map0(%i)
-    %1 = memref.subview %A_[%i, 0] [%0, 192] [1, 1]
-      : memref<192x192xf32> to memref<?x192xf32, #map1>
-    %2 = affine.min #map2(%j)
-    %3 = memref.subview %B_[0, %j] [192, %2] [1, 1]
-      : memref<192x192xf32> to memref<192x?xf32, #map1>
-    %4 = memref.subview %C_[%i, %j] [%0, %2] [1, 1]
-      : memref<192x192xf32> to memref<?x?xf32, #map1>
-    linalg.fill(%cst, %4) : f32, memref<?x?xf32, #map1>
-    linalg.matmul ins(%1, %3 : memref<?x192xf32, #map1>,
-                               memref<192x?xf32, #map1>)
-                  outs(%4 : memref<?x?xf32, #map1>)
-    linalg.yield
-  }
-  return
-}
-
-// CHECK-LABEL: @tiled_loop
-// CHECK-SAME:  %[[A:.*]]: memref<192x192xf32>, %[[B:.*]]: memref<192x192xf32>,
-// CHECK-SAME:  %[[C:.*]]: memref<192x192xf32>) {
-// CHECK:       %[[C24:.*]] = constant 24 : index
-// CHECK:       %[[C16:.*]] = constant 16 : index
-// CHECK:       %[[C0:.*]] = constant 0 : index
-// CHECK:       %[[C192:.*]] = constant 192 : index
-// CHECK:       scf.for %[[I:.*]] = %[[C0]] to %[[C192]] step %[[C24]] {
-// CHECK:         scf.for %[[J:.*]] = %[[C0]] to %[[C192]] step %[[C16]] {
-// CHECK:           %[[A_sub:.*]] = memref.subview %[[A]][%[[I]]
-// CHECK:           %[[B_sub:.*]] = memref.subview %[[B]][0, %[[J]]]
-// CHECK:           %[[C_sub:.*]] = memref.subview %[[C]][%[[I]]
-// CHECK:           linalg.fill
-// CHECK:           linalg.matmul
-
-
-func @tiled_loop_reduction(%A: memref<192x192xf32>,
-                           %B: memref<192x192xf32>,
-                           %C: memref<f32>) {
-   %c24 = constant 24 : index
-   %c16 = constant 16 : index
-   %c0 = constant 0 : index
-   %c192 = constant 192 : index
-   %cst = constant 0.000000e+00 : f32
-
-  linalg.tiled_loop (%i, %j) = (%c0, %c0) to (%c192, %c192) step (%c24, %c16)
-      ins (%A_ = %A: memref<192x192xf32>, %B_ = %B:  memref<192x192xf32>)
-      outs (%C_ = %C: memref<f32>)
-      iterators["reduction", "reduction"] {
-    linalg.fill(%cst, %A_) : f32, memref<192x192xf32>
-    linalg.yield
-  }
-  return
-}
-
-// CHECK-LABEL: @tiled_loop_reduction
-// CHECK:       %[[C24:.*]] = constant 24 : index
-// CHECK:       %[[C16:.*]] = constant 16 : index
-// CHECK:       %[[C0:.*]] = constant 0 : index
-// CHECK:       %[[C192:.*]] = constant 192 : index
-// CHECK:       scf.for %{{.*}} = %[[C0]] to %[[C192]] step %[[C24]]
-// CHECK:         scf.for %{{.*}} = %[[C0]] to %[[C192]] step %[[C16]]
-// CHECK:           linalg.fill