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

[Rolf Morel]

rolfmorel wrote:

Here's an example of how, with the proposed interface, continuous tiling can be applied to two dimensions. This has the exponentially decreasing pattern of tile sizes for the inner dimension occurring within each tiling loop of the outer dimension (where these outer loops themselves have exponentially decreasing tile sizes):

```mlir
%linalg = transform.structured.match ops{["linalg.generic"]} in %payload // assume only one linalg matched
%tile_sizes_dim0, %split_points_dim0 = transform.structured.continuous_tile_sizes %linalg { dimension = 0, target_size = 9} : (!transform.any_op) -> (!transform.param<i64>, !transform.param<i64>)
%tile_sizes_dim1, %split_points_dim1 = transform.structured.continuous_tile_sizes %linalg { dimension = 1, target_size = 7} : (!transform.any_op) -> (!transform.param<i64>, !transform.param<i64>)
%linalg_dim0_splits = transform.structured.split %linalg after %split_points_dim0 { dimension  = 0, multiway }
transform.foreach %linalg_dim0_splits, %tile_sizes_dim0 {
  ^bb1(%linalg_dim0_split: !transform.any_op, %tile_size_dim0: !transform.any_param):
  %linalg_dim0_split_tiled, %tile_loop = transform.structured.tile_using_for %linalg_dim0_split [%tile_size_dim0]
  %linalg_dim1_splits = transform.structured.split %linalg_dim0_split_tiled after %split_points_dim1 { dimension = 1, multiway }
  transform.foreach %linalg_dim1_splits, %tile_sizes_dim1 {
    ^bb2(%linalg_dim1_split: !transform.any_op, %tile_size_dim1: !transform.any_param):
    transform.structured.tile_using_for %linalg_dim1_split [0, %tile_size_dim1]
    transform.yield
  }
  transform.yield
}
```

In the above, the `multiway` attribute tells `transform.structured.split` to interpret the vector of tile sizes associated with the `split_points` argument as a list of multiple points to split the target linalg along (rather than the handles for `target` and `split_points` being co-indexed with just one split point for each linalg).

This pattern naturally extends to higher dimensions.

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