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

[Rolf Morel]

rolfmorel wrote:

Here's a suggestion for a test case. It requires the pending change to `transform.foreach` to land first though (#93705).

```mlir
// RUN: mlir-opt --transform-interpreter --split-input-file %s | FileCheck %s

PAYLOAD_CONTAINING_ANY_LINALG_HERE

module attributes {transform.with_named_sequence} {
  transform.named_sequence @tile_dim0_continously(%linalg: !transform.any_op {transform.readonly}, %target_size: !transform.param<i64>) -> !transform.any_op, !transform.any_op {
    %tile_sizes, %split_points = transform.structured.continuous_tile_sizes %linalg descending_from %target_size { dimension = 0 } : (!transform.any_op, !transform.param<i64>) -> (!transform.param<i64>, !transform.param<i64>)
    %linalg_splits, %empty = transform.structured.split %linalg after %split_points { dimension  = 0, multiway } : ...
    %tiled_splits, %loops = transform.foreach %linalg_splits, %tile_sizes {
      ^bb1(%linalg_split: !transform.any_op, %tile_size: !transform.param<i64>):
      %tiled_linalg_split, %dim0_loop = transform.structured.tile_using_for %linalg_split [%tile_size] : ...
      transform.yield %tiled_linalg_split, %dim0_loop : ...
    }
    transform.yield %tiled_splits, %loops : ...
  }
  transform.named_sequence @__transform_main(%payload: !transform.any_op) {
    %linalg = transform.structured.match ops{["linalg.generic"]} in %payload 
    %target_size = transform.param.constant 9 : !transform.param<index>
    // NB: we would need a foreach here/in the named sequence if %linalg would associate multiple linalg.generics
    %4 = transform.include @tile_dim0_continously failures(propagate) (%linalg) : (!transform.any_op) -> !transform.any_op
    transform.yield 
  }
}
```

Note the above assumes that `transform.structured.continuous_tile_sizes` is able to take `target_size` as a param instead of passing it as an attribute. (Note as well that the above is not quite syntactically correct: e.g., some types are left as homework.)

## Demonstrating abstraction

The main point of interest regarding the test case is that it demonstrates how to do abstraction with the Transform dialect. That is, it shows how a higher-level transform can be implemented with the Transform dialect while the details of this implementation can be hidden. In this way, the higher-level transform's functionality is exposed basically* as if it were available like any other Transform op.

(*: Note that currently a named sequence cannot be parametric w.r.t. attribute arguments that its contained ops take as attributes and not as params. So each op that we would like to use in named sequences - and whose attribute arguments we like expose as arguments of the sequence - needs to be modified to take that attribute as a `transform.param` argument as well. This is not ideal.)

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