[MLIR] Add continuous tiling to TileUsingForOp

This patch adds continuous tiling options to TileUsingForOp.
In cases where the tensor dimensions are not exactly divisible
by the tile size, we are left with leftover tensor chunks that
are irregularly tiled. This approach attempts to tile the leftover
chunk with a smaller tile size and repeats this process recursively
using exponentially diminishing tile sizes. The transform eventually
generates a chain of loops that apply tiling using diminishing tile
sizes. The transform lowers from the linalg dialect to scf dialect.

Author: muneebkhan85

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

@@ -1833,7 +1833,9 @@ def TileUsingForOp : Op<Transform_Dialect, "structured.tile_using_for",
     be as many handles as `ShapedType::kDynamic` values in the
     `static_sizes` attribute. A static size of `0` indicates that the dimension
     should not be tiled. No loop will be generated for such dimensions. If all
-    tile sizes are `0`, this transform is effectively a no-op.
+    tile sizes are `0`, this transform is effectively a no-op. To apply
+    continuous tiling `continuous_tiles` needs to be supplied with as many
+    boolean values as there are nested loops.
 
     This op returns handles to the tiled op (in the generated loop nest) and the
     generated loops. The number of loops is the number of tile sizes that are
@@ -1859,6 +1861,7 @@ def TileUsingForOp : Op<Transform_Dialect, "structured.tile_using_for",
   let arguments = (ins TransformHandleTypeInterface:$target,
                    Variadic<TransformAnyParamTypeOrAnyHandle>:$dynamic_sizes,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_sizes,
+                   DefaultValuedOptionalAttr<DenseBoolArrayAttr, "{}">:$continuous_tiles,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$interchange,
                    DefaultValuedOptionalAttr<DenseBoolArrayAttr, "{}">:$scalable_sizes);
   let results = (outs TransformHandleTypeInterface:$tiled_linalg_op,
@@ -1867,22 +1870,26 @@ def TileUsingForOp : Op<Transform_Dialect, "structured.tile_using_for",
     OpBuilder<(ins "TypeRange":$loopTypes,
                    "Value":$target,
                    "ArrayRef<int64_t>":$staticTileSizes,
+                   CArg<"ArrayRef<bool>", "{}">:$continuousTiles,
                    CArg<"ArrayRef<int64_t>", "{}">:$interchange,
                    CArg<"std::optional<ArrayRef<bool>>", "std::nullopt">:
                       $scalableSizes)>,
     OpBuilder<(ins "TypeRange":$loopTypes,
                    "Value":$target,
                    "ArrayRef<OpFoldResult>":$mixedTileSizes,
+                   CArg<"ArrayRef<bool>", "{}">:$continuousTiles,
                    CArg<"ArrayRef<int64_t>", "{}">:$interchange,
                    CArg<"std::optional<ArrayRef<bool>>", "std::nullopt">:
                       $scalableSizes)>,
     OpBuilder<(ins "Value":$target,
                    "ArrayRef<int64_t>":$staticTileSizes,
+                   CArg<"ArrayRef<bool>", "{}">:$continuousTiles,
                    CArg<"ArrayRef<int64_t>", "{}">:$interchange,
                    CArg<"std::optional<ArrayRef<bool>>", "std::nullopt">:
                       $scalableSizes)>,
     OpBuilder<(ins "Value":$target,
                    "ArrayRef<OpFoldResult>":$mixedTileSizes,
+                   CArg<"ArrayRef<bool>", "{}">:$continuousTiles,
                    CArg<"ArrayRef<int64_t>", "{}">:$interchange,
                    CArg<"std::optional<ArrayRef<bool>>", "std::nullopt">:
                       $scalableSizes)>,

mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h

@@ -71,6 +71,14 @@ struct SCFTilingOptions {
         mapping, [](auto attr) -> Attribute { return attr; });
     return *this;
   }
+
+  /// Specify which loops in the loop nest are to be continuously tiled.
+  SmallVector<bool> continuousTileMappingVector = {};
+  SCFTilingOptions &setCTileMapping(ArrayRef<bool> ctile) {
+    continuousTileMappingVector =
+        llvm::map_to_vector(ctile, [](auto attr) -> bool { return attr; });
+    return *this;
+  }
 };
 
 /// Transformation information returned after tiling.
@@ -92,6 +100,12 @@ FailureOr<SCFTilingResult> tileUsingSCF(RewriterBase &rewriter,
                                         TilingInterface op,
                                         const SCFTilingOptions &options);
 
+/// Method to continuously tile an op that implements the `TilingInterface`
+/// using `scf.for` for iterating over the tiles.
+FailureOr<SCFTilingResult>
+continuousTileUsingSCF(RewriterBase &rewriter, TilingInterface op,
+                       const SCFTilingOptions &options);
+
 /// Options used to control tile + fuse.
 struct SCFTileAndFuseOptions {
   /// The tiling options used to control the tiling of the consumer.

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

@@ -2476,39 +2476,42 @@ DiagnosedSilenceableFailure transform::TileReductionUsingForallOp::applyToOne(
 void transform::TileUsingForOp::build(
     OpBuilder &builder, OperationState &result, TypeRange loopTypes,
     Value target, ArrayRef<int64_t> staticTileSizes,
-    ArrayRef<int64_t> interchange,
+    ArrayRef<bool> continuousTiles, ArrayRef<int64_t> interchange,
     std::optional<ArrayRef<bool>> scalableSizes) {
   return build(builder, result, loopTypes,
                /*target=*/target,
                /*mixedTileSizes=*/
                getAsOpFoldResult(builder.getI64ArrayAttr(staticTileSizes)),
-               interchange, scalableSizes);
+               continuousTiles, interchange, scalableSizes);
 }
 
 void transform::TileUsingForOp::build(
     OpBuilder &builder, OperationState &result, Value target,
-    ArrayRef<int64_t> staticTileSizes, ArrayRef<int64_t> interchange,
+    ArrayRef<int64_t> staticTileSizes, ArrayRef<bool> continuousTiles,
+    ArrayRef<int64_t> interchange,
     std::optional<ArrayRef<bool>> scalableSizes) {
   build(builder, result, target,
         getAsOpFoldResult(builder.getI64ArrayAttr(staticTileSizes)),
-        interchange, scalableSizes);
+        builder.getDenseBoolArrayAttr(continuousTiles), interchange,
+        scalableSizes);
 }
 
 void transform::TileUsingForOp::build(
     OpBuilder &builder, OperationState &result, Value target,
-    ArrayRef<OpFoldResult> mixedTileSizes, ArrayRef<int64_t> interchange,
+    ArrayRef<OpFoldResult> mixedTileSizes, ArrayRef<bool> continuousTiles,
+    ArrayRef<int64_t> interchange,
     std::optional<ArrayRef<bool>> scalableSizes) {
   // Loop types are automaticaly splat by the callee, setting up one is
   // enough.
   SmallVector<Type> loopTypes(1, builder.getType<transform::AnyOpType>());
-  build(builder, result, loopTypes, target, mixedTileSizes, interchange,
-        scalableSizes);
+  build(builder, result, loopTypes, target, mixedTileSizes, continuousTiles,
+        interchange, scalableSizes);
 }
 
 void transform::TileUsingForOp::build(
     OpBuilder &builder, OperationState &result, TypeRange loopTypes,
     Value target, ArrayRef<OpFoldResult> mixedTileSizes,
-    ArrayRef<int64_t> interchange,
+    ArrayRef<bool> continuousTiles, ArrayRef<int64_t> interchange,
     std::optional<ArrayRef<bool>> scalableSizes) {
   SmallVector<int64_t> staticTileSizes;
   SmallVector<Value> dynamicTileSizes;
@@ -2517,6 +2520,7 @@ void transform::TileUsingForOp::build(
   // attributes for multiple variadic operands. In the absence of this,
   // horrible bugs ensue.
   auto staticTileSizesAttr = builder.getDenseI64ArrayAttr(staticTileSizes);
+  auto continuousTilesAttr = builder.getDenseBoolArrayAttr(continuousTiles);
   unsigned numExpectedLoops =
       staticTileSizes.size() - llvm::count(staticTileSizes, 0);
   SmallVector<Type> resultTypes;
@@ -2535,6 +2539,7 @@ void transform::TileUsingForOp::build(
         /*target=*/target,
         /*dynamic_sizes=*/dynamicTileSizes,
         /*static_sizes=*/staticTileSizesAttr,
+        /*continuous_tiles=*/continuousTilesAttr,
         /*interchange=*/builder.getDenseI64ArrayAttr(interchange),
         /*scalable_sizes=*/expandedScalableSizes);
 }
@@ -2675,8 +2680,15 @@ transform::TileUsingForOp::apply(transform::TransformRewriter &rewriter,
     }
 
     tilingOptions.setInterchange(getInterchange());
-    FailureOr<scf::SCFTilingResult> maybeTilingResult =
-        tileUsingSCF(rewriter, tilingInterface, tilingOptions);
+    tilingOptions.setCTileMapping(getContinuousTiles());
+
+    FailureOr<scf::SCFTilingResult> maybeTilingResult;
+    if (tilingOptions.continuousTileMappingVector.empty())
+      maybeTilingResult =
+          tileUsingSCF(rewriter, tilingInterface, tilingOptions);
+    else
+      maybeTilingResult =
+          continuousTileUsingSCF(rewriter, tilingInterface, tilingOptions);
     if (failed(maybeTilingResult))
       return DiagnosedSilenceableFailure::definiteFailure();
 
@@ -2725,6 +2737,18 @@ ParseResult parseOptionalInterchange(OpAsmParser &parser,
   return success();
 }
 
+ParseResult parseOptionalContinuousTiles(OpAsmParser &parser,
+                                         OperationState &result) {
+  if (failed(parser.parseOptionalKeyword("continuous_tiles")))
+    return success();
+  if (failed(parser.parseEqual()))
+    return failure();
+  result.addAttribute(
+      transform::TileUsingForOp::getContinuousTilesAttrName(result.name),
+      DenseBoolArrayAttr::parse(parser, Type{}));
+  return success();
+}
+
 void printOptionalInterchange(OpAsmPrinter &p,
                               ArrayRef<int64_t> interchangeVals) {
   if (!interchangeVals.empty()) {
@@ -2747,6 +2771,7 @@ ParseResult transform::TileUsingForOp::parse(OpAsmParser &parser,
   if (parser.parseOperand(target) || parser.getCurrentLocation(&operandLoc) ||
       parseDynamicIndexList(parser, dynamicSizes, staticSizes, scalableVals) ||
       parseOptionalInterchange(parser, result) ||
+      parseOptionalContinuousTiles(parser, result) ||
       parser.parseOptionalAttrDict(result.attributes) ||
       parser.parseColonType(functionalType))
     return ParseResult::failure();