[Mlir-commits] [mlir] [MLIR] Add continuous tiling to TileUsingForOp (PR #82792)

[Oleksandr Alex Zinenko]

================
@@ -679,6 +863,276 @@ mlir::scf::tileUsingSCF(RewriterBase &rewriter, TilingInterface op,
   return scf::SCFTilingResult{tilingResult->tiledOps, loops, replacements};
 }
 
+/// Implementation of continuous-tiling transformation of `op`
+/// that implements the `TilingInterface` using a sequence of
+/// `scf.for` loops to iterate over tiles of exponentially
+/// diminishing sizes.
+///
+/// The generated sequence of `scf.for` loops iterate over tiles of
+/// exponentially diminishing sizes as opposed to vanilla tiling scheme where
+/// only the tile sizes specified as transform parameters are used.
+/// continuous-tiling first applies regular tiling using the specified
+/// tile size, then halves the tile size and uses it to tile the leftover
+/// chunk. This process of halving the tile size and tiling the leftover
+/// chunk is repeated until tile size reaches 1. The transform parameter
+/// continuous_tiles controls which nested loop should be tiled. If all
+/// arguments in continuous_tiles are set to false, the result is identical
+/// to vanilla tiling transform.
+///
+/// When tiling a tensor of size M with tile size 8, it generates
+/// the following loop to tile
+///
+/// for (int i = 0; i < M; i += 8) {
+///   int size = min(8, M-i);
+///   // size is unknown at compile time because M is dynamic
+///   // compute using dynamic size (not optimal)
+/// }
+///
+/// In case of continuous tiling the above loop is converted to a chain of
+/// loops that attempt to tile the cases where (M-i < 8) using smaller
+/// tile sizes, as follows
+///
+/// // use tile size of 8 in the head loop as originally used
+/// for (int i = 0; i != (M - M % 8); i += 8) {
+///   int size = 8;
+///   // size is constant the tail is moved to the next loop
+///   // compute using static size (optimal)
+/// }
+/// // use halved tile size 4 to tile the remaining chunk
+/// for (int i = (M - M % 8); i != (M - M % 4); i += 4)
+///   int size = 4;
+/// // use halved tile size 2 to tile the remaining chunk
+/// for (int i = (M - M % 4); i != (M - M % 2); i += 2)
+///   int size = 2;
+/// // use halved tile size 1 to tile the remaining chunk
+/// for (int i = (M - M % 2); i != M; i += 1)
+///   int size = 1;
+///
+/// All tail-loops can be converted to IFs as they execute at-most once.
+///
+FailureOr<scf::SCFTilingResult>
+mlir::scf::continuousTileUsingSCF(RewriterBase &rewriter, TilingInterface op,
+                                  const scf::SCFTilingOptions &options) {
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPointAfter(op);
+
+  std::map<int, OpFoldResult> sizesMap;
+
+  if (!options.tileSizeComputationFunction) {
+    return rewriter.notifyMatchFailure(
+        op, "missing tile size computation function");
----------------
ftynse wrote:

These error messages are invisible unless the function is called from a pattern rewriting engine, which I suspect is not the case here. If you want to provide a better error description, you can either use `DiagnosedSilenceableFailure` from the Transform dialect or ask the caller to give an `llvm::raw_ostream`, defaulted to `llvm::nulls` that will accept the error message.

https://github.com/llvm/llvm-project/pull/82792