[mlir][linalg] Support tile+peel with TiledLoopOp

[springerm]

Summary

Only scf.for was supported until now.

Depends On D110214

[pifon2a]

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

300

Would it make sense to create a method that does loop peeling for options.peeledLoops?

307–309

can this assertion be hoisted from the loop?

312–314

nit: what about moving these assertions into peelLoop func?

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

583–591

neat!

[springerm]

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

307–309

It could, but then I would have to duplicate the if-check in the preceding line.

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

 
 mlir::linalg::LinalgBaseTilingPattern::LinalgBaseTilingPattern(
     MLIRContext *context, LinalgTilingOptions options,
     LinalgTransformationFilter filter, PatternBenefit benefit)
     : RewritePattern(MatchAnyOpTypeTag(), benefit, context), filter(filter),
       options(options) {}
 
 /// Try to peel a loop `op` and return the new result.
-// TODO: Only scf.for loops are supported at the moment.
+// TODO: Add support for scf.parallel and affine.for loops.
 static SmallVector<Value, 4> peelLoop(RewriterBase &rewriter, Operation *op) {
   return llvm::TypeSwitch<Operation *, SmallVector<Value, 4>>(op)
       .Case<scf::ForOp>([&](scf::ForOp forOp) {
         scf::ForOp partialIteration;
         if (succeeded(scf::peelAndCanonicalizeForLoop(rewriter, forOp,
                                                       partialIteration)))
           return partialIteration->getResults();
         assert(!partialIteration && "expected that loop was not peeled");
         return forOp->getResults();
       })
       .Default([&](Operation *op) { return op->getResults(); });
 }
 
+/// Try to peel a TiledLoopOp and return the new result.
+static SmallVector<Value, 4> peelLoop(RewriterBase &rewriter,
+                                      TiledLoopOp tiledLoop, int64_t idx) {
+  TiledLoopOp result;
+  if (succeeded(peelAndCanonicalizeTiledLoop(rewriter, tiledLoop, idx, result)))
+    return result->getResults();
+  assert(!result && "expected that loop was not peeled");
+  return tiledLoop->getResults();
+}
+
 LogicalResult mlir::linalg::LinalgBaseTilingPattern::matchAndRewriteBase(
     Operation *op, PatternRewriter &rewriter, TiledLinalgOp &result) const {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
   if (!linalgOp)
     return failure();
   if (failed(filter.checkAndNotify(rewriter, linalgOp)))
     return failure();
 
   Optional<TiledLinalgOp> res = tileLinalgOp(rewriter, linalgOp, options);
 
   if (!res)
     return failure();
   // Clear filter to stop recursive pattern application.
   filter.replaceLinalgTransformationFilter(rewriter, res->op);
 
   // Peel loops.
   for (int64_t loop : options.peeledLoops) {
     assert(loop < static_cast<int64_t>(res->loops.size()) &&
            "requested peeling of non-existing loop");
+    SmallVector<Value, 4> loopResults;
     Operation *loopOp = res->loops[loop];
-    SmallVector<Value, 4> loopResults = peelLoop(rewriter, loopOp);
+    if (options.loopType == LinalgTilingLoopType::TiledLoops) {
+      assert(llvm::all_of(
+                 res->loops,
+                 [&](Operation *op) { return op == res->loops.front(); }) &&
+             "expected that all loop ops are the same TiledLoopOp");
+      auto tiledLoopOp = dyn_cast<TiledLoopOp>(loopOp);
+      assert(tiledLoopOp && "expected TiledLoopOp");
+      assert(loop < static_cast<int64_t>(tiledLoopOp.iterator_types().size()) &&
+             "requested peeling of non-existing loop");
+      loopResults = peelLoop(rewriter, tiledLoopOp, loop);
+    } else {
+      loopResults = peelLoop(rewriter, loopOp);
+    }
+
     // The result of the loop nest may change with peeling.
     if (res->tensorResults.size() == loopOp->getNumResults() &&
         std::equal(res->tensorResults.begin(), res->tensorResults.end(),
                    loopOp->getResults().begin()))
       res->tensorResults = loopResults;
   }
 
   // Consider padding on the fly only if the op has tensor semantics.

mlir/test/Dialect/Linalg/tile-and-peel-tensors.mlir

 // RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-pattern tile-sizes=256,128,512 peeled-loops=0" -canonicalize | \
 // RUN:     FileCheck %s -check-prefix=CHECK-PEEL-0
 
 // RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-pattern tile-sizes=256,128,512 peeled-loops=1,2" -canonicalize | \
 // RUN:     FileCheck %s -check-prefix=CHECK-PEEL-12
 
+// RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-pattern tile-sizes=256,128,512 loop-type=tiled_loop peeled-loops=0" -canonicalize | \
+// RUN:     FileCheck %s -check-prefix=CHECK-TILED-LOOP-PEEL-0
+
+// RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-pattern tile-sizes=256,128,512 loop-type=tiled_loop peeled-loops=0,1" -canonicalize | \
+// RUN:     FileCheck %s -check-prefix=CHECK-TILED-LOOP-PEEL-01
+
 //     CHECK-PEEL-0: func @matmul_static_tensor
 // CHECK-PEEL-0-DAG:   %[[c0:.*]] = constant 0 : index
 // CHECK-PEEL-0-DAG:   %[[c128:.*]] = constant 128 : index
 // CHECK-PEEL-0-DAG:   %[[c256:.*]] = constant 256 : index
 // CHECK-PEEL-0-DAG:   %[[c512:.*]] = constant 512 : index
 // CHECK-PEEL-0-DAG:   %[[c1280:.*]] = constant 1280 : index
 // CHECK-PEEL-0-DAG:   %[[c1600:.*]] = constant 1600 : index
 // CHECK-PEEL-0-DAG:   %[[c1700:.*]] = constant 1700 : index

 //     CHECK-PEEL-12:         linalg.matmul ins({{.*}} : tensor<?x512xf32>, tensor<512x128xf32>) outs({{.*}} : tensor<?x128xf32>)
 //     CHECK-PEEL-12:       }
 //     CHECK-PEEL-12:       linalg.matmul ins({{.*}} : tensor<?x64xf32>, tensor<64x128xf32>) outs({{.*}} : tensor<?x128xf32>)
 //     CHECK-PEEL-12:     }
 //     CHECK-PEEL-12:     scf.for %{{.*}} = %[[c0]] to %[[c1600]] step %[[c512]] {{.*}} {
 //     CHECK-PEEL-12:       linalg.matmul ins({{.*}} : tensor<?x?xf32>, tensor<?x36xf32>) outs({{.*}} : tensor<?x36xf32>)
 //     CHECK-PEEL-12:     }
 //     CHECK-PEEL-12:   }
+
+//     CHECK-TILED-LOOP-PEEL-0: func @matmul_static_tensor
+// CHECK-TILED-LOOP-PEEL-0-DAG:   %[[c0:.*]] = constant 0 : index
+// CHECK-TILED-LOOP-PEEL-0-DAG:   %[[c128:.*]] = constant 128 : index
+// CHECK-TILED-LOOP-PEEL-0-DAG:   %[[c256:.*]] = constant 256 : index
+// CHECK-TILED-LOOP-PEEL-0-DAG:   %[[c512:.*]] = constant 512 : index
+// CHECK-TILED-LOOP-PEEL-0-DAG:   %[[c1280:.*]] = constant 1280 : index
+// CHECK-TILED-LOOP-PEEL-0-DAG:   %[[c1500:.*]] = constant 1500 : index
+// CHECK-TILED-LOOP-PEEL-0-DAG:   %[[c1600:.*]] = constant 1600 : index
+// CHECK-TILED-LOOP-PEEL-0-DAG:   %[[c1700:.*]] = constant 1700 : index
+//     CHECK-TILED-LOOP-PEEL-0:   linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) = (%[[c0]], %[[c0]], %[[c0]]) to (%[[c1280]], %[[c1700]], %[[c1600]]) step (%[[c256]], %[[c128]], %[[c512]])
+//     CHECK-TILED-LOOP-PEEL-0:     linalg.matmul ins({{.*}} : tensor<256x?xf32>, tensor<?x?xf32>) outs({{.*}} : tensor<256x?xf32>)
+//     CHECK-TILED-LOOP-PEEL-0:   }
+//     CHECK-TILED-LOOP-PEEL-0:   linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) = (%[[c1280]], %[[c0]], %[[c0]]) to (%[[c1500]], %[[c1700]], %[[c1600]]) step (%[[c256]], %[[c128]], %[[c512]])
+//     CHECK-TILED-LOOP-PEEL-0:     linalg.matmul ins({{.*}} : tensor<?x?xf32>, tensor<?x?xf32>) outs({{.*}} : tensor<?x?xf32>)
+//     CHECK-TILED-LOOP-PEEL-0:   }
+
+//     CHECK-TILED-LOOP-PEEL-01: func @matmul_static_tensor
+// CHECK-TILED-LOOP-PEEL-01-DAG:   %[[c0:.*]] = constant 0 : index
+// CHECK-TILED-LOOP-PEEL-01-DAG:   %[[c128:.*]] = constant 128 : index
+// CHECK-TILED-LOOP-PEEL-01-DAG:   %[[c256:.*]] = constant 256 : index
+// CHECK-TILED-LOOP-PEEL-01-DAG:   %[[c512:.*]] = constant 512 : index
+// CHECK-TILED-LOOP-PEEL-01-DAG:   %[[c1280:.*]] = constant 1280 : index
+// CHECK-TILED-LOOP-PEEL-01-DAG:   %[[c1500:.*]] = constant 1500 : index
+// CHECK-TILED-LOOP-PEEL-01-DAG:   %[[c1600:.*]] = constant 1600 : index
+// CHECK-TILED-LOOP-PEEL-01-DAG:   %[[c1664:.*]] = constant 1664 : index
+// CHECK-TILED-LOOP-PEEL-01-DAG:   %[[c1700:.*]] = constant 1700 : index
+//     CHECK-TILED-LOOP-PEEL-01:   linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) = (%[[c0]], %[[c0]], %[[c0]]) to (%[[c1280]], %[[c1664]], %[[c1600]]) step (%[[c256]], %[[c128]], %[[c512]])
+//     CHECK-TILED-LOOP-PEEL-01:     linalg.matmul ins({{.*}} : tensor<256x?xf32>, tensor<?x128xf32>) outs({{.*}} : tensor<256x128xf32>)
+//     CHECK-TILED-LOOP-PEEL-01:   }
+//     CHECK-TILED-LOOP-PEEL-01:   linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) = (%[[c0]], %[[c1664]], %[[c0]]) to (%[[c1280]], %[[c1700]], %[[c1600]]) step (%[[c256]], %[[c128]], %[[c512]])
+//     CHECK-TILED-LOOP-PEEL-01:     linalg.matmul ins({{.*}} : tensor<256x?xf32>, tensor<?x?xf32>) outs({{.*}} : tensor<256x?xf32>)
+//     CHECK-TILED-LOOP-PEEL-01:   }
+//     CHECK-TILED-LOOP-PEEL-01:   linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) = (%[[c1280]], %[[c0]], %[[c0]]) to (%[[c1500]], %[[c1700]], %[[c1600]]) step (%[[c256]], %[[c128]], %[[c512]])
+//     CHECK-TILED-LOOP-PEEL-01:     linalg.matmul ins({{.*}} : tensor<?x?xf32>, tensor<?x?xf32>) outs({{.*}} : tensor<?x?xf32>)
+//     CHECK-TILED-LOOP-PEEL-01:   }
 func @matmul_static_tensor(%arg0: tensor<1500x1600xf32>, %arg1: tensor<1600x1700xf32>)
     -> tensor<1500x1700xf32> {
   %out = linalg.init_tensor [1500, 1700] : tensor<1500x1700xf32>
   %r = linalg.matmul {__internal_linalg_transform__ = "tile"}
       ins(%arg0, %arg1: tensor<1500x1600xf32>, tensor<1600x1700xf32>)
       outs(%out: tensor<1500x1700xf32>) -> tensor<1500x1700xf32>
   return %r : tensor<1500x1700xf32>
 }

 //     CHECK-PEEL-12:       }
 //     CHECK-PEEL-12:     }
 //     CHECK-PEEL-12:     scf.for %{{.*}} {
 //     CHECK-PEEL-12:       scf.for %{{.*}} = %[[c0]] to %{{.*}} step %[[c512]] {{.*}} {
 //     CHECK-PEEL-12:         linalg.matmul ins({{.*}} : tensor<?x?xf32>, tensor<?x?xf32>) outs({{.*}} : tensor<?x?xf32>)
 //     CHECK-PEEL-12:       }
 //     CHECK-PEEL-12:     }
 //     CHECK-PEEL-12:   }
+
+//     CHECK-TILED-LOOP-PEEL-0: func @matmul_dynamic_tensor
+// CHECK-TILED-LOOP-PEEL-0-DAG:   %[[c0:.*]] = constant 0 : index
+// CHECK-TILED-LOOP-PEEL-0-DAG:   %[[c128:.*]] = constant 128 : index
+// CHECK-TILED-LOOP-PEEL-0-DAG:   %[[c256:.*]] = constant 256 : index
+// CHECK-TILED-LOOP-PEEL-0-DAG:   %[[c512:.*]] = constant 512 : index
+//     CHECK-TILED-LOOP-PEEL-0:   linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) = (%[[c0]], %[[c0]], %[[c0]]) to (%{{.*}}, %{{.*}}, %{{.*}}) step (%[[c256]], %[[c128]], %[[c512]])
+//     CHECK-TILED-LOOP-PEEL-0:     linalg.matmul ins({{.*}} : tensor<256x?xf32>, tensor<?x?xf32>) outs({{.*}} : tensor<256x?xf32>)
+//     CHECK-TILED-LOOP-PEEL-0:   }
+//     CHECK-TILED-LOOP-PEEL-0:   linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) = (%{{.*}}, %[[c0]], %[[c0]]) to (%{{.*}}, %{{.*}}, %{{.*}}) step (%[[c256]], %[[c128]], %[[c512]])
+//     CHECK-TILED-LOOP-PEEL-0:     linalg.matmul ins({{.*}} : tensor<?x?xf32>, tensor<?x?xf32>) outs({{.*}} : tensor<?x?xf32>)
+//     CHECK-TILED-LOOP-PEEL-0:   }
+
+//     CHECK-TILED-LOOP-PEEL-01: func @matmul_dynamic_tensor
+// CHECK-TILED-LOOP-PEEL-01-DAG:   %[[c0:.*]] = constant 0 : index
+// CHECK-TILED-LOOP-PEEL-01-DAG:   %[[c128:.*]] = constant 128 : index
+// CHECK-TILED-LOOP-PEEL-01-DAG:   %[[c256:.*]] = constant 256 : index
+// CHECK-TILED-LOOP-PEEL-01-DAG:   %[[c512:.*]] = constant 512 : index
+//     CHECK-TILED-LOOP-PEEL-01:   linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) = (%[[c0]], %[[c0]], %[[c0]]) to (%{{.*}}, %{{.*}}, %{{.*}}) step (%[[c256]], %[[c128]], %[[c512]])
+//     CHECK-TILED-LOOP-PEEL-01:     linalg.matmul ins({{.*}} : tensor<256x?xf32>, tensor<?x128xf32>) outs({{.*}} : tensor<256x128xf32>)
+//     CHECK-TILED-LOOP-PEEL-01:   }
+//     CHECK-TILED-LOOP-PEEL-01:   linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) = (%[[c0]], %{{.*}}, %[[c0]]) to (%{{.*}}, %{{.*}}, %{{.*}}) step (%[[c256]], %[[c128]], %[[c512]])
+//     CHECK-TILED-LOOP-PEEL-01:     linalg.matmul ins({{.*}} : tensor<256x?xf32>, tensor<?x?xf32>) outs({{.*}} : tensor<256x?xf32>)
+//     CHECK-TILED-LOOP-PEEL-01:   }
+//     CHECK-TILED-LOOP-PEEL-01:   linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) = (%{{.*}}, %[[c0]], %[[c0]]) to (%{{.*}}, %{{.*}}, %{{.*}}) step (%[[c256]], %[[c128]], %[[c512]])
+//     CHECK-TILED-LOOP-PEEL-01:     linalg.matmul ins({{.*}} : tensor<?x?xf32>, tensor<?x?xf32>) outs({{.*}} : tensor<?x?xf32>)
+//     CHECK-TILED-LOOP-PEEL-01:   }
 func @matmul_dynamic_tensor(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>)
     -> tensor<?x?xf32> {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
   %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
   %out = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32>
   %r = linalg.matmul {__internal_linalg_transform__ = "tile"}
       ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
       outs(%out: tensor<?x?xf32>) -> tensor<?x?xf32>
   return %r : tensor<?x?xf32>
 }

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

struct TestLinalgTransforms
   ListOption<unsigned> testTiledLoopPeeling{
       *this, "test-tiled-loop-peeling",
       llvm::cl::desc("Test peeling of linalg.tiled_loop ops"),
       llvm::cl::OneOrMore, llvm::cl::MiscFlags::CommaSeparated};
   Option<bool> skipPartial{
       *this, "skip-partial",
       llvm::cl::desc("Skip loops inside partial iterations during peeling"),
       llvm::cl::init(false)};
+  Option<std::string> loopType{
+      *this, "loop-type",
+      llvm::cl::desc("Specify the type of loops to generate: for, parallel or "
+                     "tiled_loop"),
+      llvm::cl::init("for")};
 };
 } // end anonymous namespace
 
 static void applyPatterns(FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   RewritePatternSet patterns(ctx);
 
   //===--------------------------------------------------------------------===//

 
 // For now, just assume it is the zero of type.
 // In the future, it should be the zero of type + op.
 static Value getNeutralOfLinalgOp(OpBuilder &b, OpOperand &op) {
   auto t = getElementTypeOrSelf(op.get());
   return b.create<ConstantOp>(op.getOwner()->getLoc(), t, b.getZeroAttr(t));
 }
 
-static void applyTilePattern(FuncOp funcOp, ArrayRef<int64_t> tileSizes,
-                             bool padTiles, ArrayRef<int64_t> peeledLoops,
+static void applyTilePattern(FuncOp funcOp, std::string loopType,
+                             ArrayRef<int64_t> tileSizes, bool padTiles,
+                             ArrayRef<int64_t> peeledLoops,
                              bool scalarizeDynamicDims) {
   MLIRContext *context = funcOp.getContext();
   RewritePatternSet tilingPattern(context);
-  auto linalgTilingOptions =
-      linalg::LinalgTilingOptions().setPeeledLoops(peeledLoops);
+  LinalgTilingLoopType type =
+      llvm::StringSwitch<LinalgTilingLoopType>(loopType)
+          .Case("for", LinalgTilingLoopType::Loops)
+          .Case("affine", LinalgTilingLoopType::AffineLoops)
+          .Case("parallel", LinalgTilingLoopType::ParallelLoops)
+          .Case("tiled_loop", LinalgTilingLoopType::TiledLoops);
+  auto linalgTilingOptions = linalg::LinalgTilingOptions()
+                                 .setPeeledLoops(peeledLoops)
+                                 .setLoopType(type);
   if (scalarizeDynamicDims) {
     linalgTilingOptions.scalarizeDynamicDims();
     assert(tileSizes.empty() &&
            "tileSizes and scalarizeDynamicDims is mutually exclusive");
   } else {
     linalgTilingOptions.setTileSizes(tileSizes);
   }
   if (padTiles)
void TestLinalgTransforms::runOnFunction() {
   if (testGeneralizePadTensor)
     return applyGeneralizePadTensorPatterns(getFunction());
   if (testSwapSubTensorPadTensor)
     return applyExtractSliceOfPadTensorSwapPattern(getFunction());
   if (testTiledLoopPeeling.hasValue())
     return applyTiledLoopPeelingPattern(getFunction(), testTiledLoopPeeling,
                                         skipPartial);
   if (testTilePattern)
-    return applyTilePattern(getFunction(), tileSizes, padTiles, peeledLoops,
-                            /*scalarizeDynamicDims=*/false);
+    return applyTilePattern(getFunction(), loopType, tileSizes, padTiles,
+                            peeledLoops, /*scalarizeDynamicDims=*/false);
   if (testTileScalarizeDynamicDims)
-    return applyTilePattern(getFunction(), tileSizes, padTiles,
+    return applyTilePattern(getFunction(), loopType, tileSizes, padTiles,
                             /*peeledLoops=*/{}, /*scalarizeDynamicDims=*/true);
   if (testHoistPadding) {
     getFunction().walk([&](linalg::PadTensorOp padTensorOp) {
       (void)linalg::hoistPaddingOnTensors(padTensorOp, testHoistPadding);
     });
   }
   if (testInterchangePattern.hasValue())
     return applyInterchangePattern(getFunction(), testInterchangePattern);