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

[llvmlistbot at llvm.org]

muneebkhan85 wrote:

@ftynse I have introduced new test cases under `mlir/test/Dialect/Linalg/continuous-tiling-full.mlir` to demonstrate the full idea of how continuous tiling is applied. 

While the examples here test for static dimension sizes, we have encountered an issue for the dynamic case where dimension sizes are not known, i.e, `?x?`. It turns out that in this case, multiway splitOp generates an "additional" trailing linalg op whose dimensions would evaluate to zero at runtime (in case of generating `chunk_sizes` using `continuous_tile_sizes`). This can be attributed to `linalg::splitOp` creating a valid head and tail on the last `chunkSize`. This causes the foreach  #93705 to fail in this case as `linalg_splits` and `tile_sizes` are not the same length. We could have a fully working dynamic example, if the last element in `linalg_splits` could be "dropped" in some way. One way to do this is to add `zip_shortest` to `foreach` when it takes more than one handles. If you have another suggestion or preference, please let us know. 

Here's an example for the dynamic dimensions case. `linalg_splits` and `tile_sizes` do not have the same payload size (6 vs 5), which makes the `foreach` fail.

```
module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
    %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
    %tile_sizes, %chunk_sizes = transform.structured.continuous_tile_sizes %0 { dimension = 0, target_size = 9 } : (!transform.any_op) -> !transform.any_op
    %linalg_splits, %empty = transform.structured.split %0 after %chunk_sizes { dimension = 0, multiway } : !transform.any_op, !transform.any_op
    transform.foreach %linalg_splits, %tile_sizes : !transform.any_op, !transform.any_op {
    ^bb1(%linalg_split: !transform.any_op, %tile_size: !transform.any_op):
      %tiled_linalg_split, %dim0_loop = transform.structured.tile_using_for %linalg_split tile_sizes [%tile_size] : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
      transform.yield
    }
    transform.yield
  }
}

func.func @continuous_tile_linalg_matmul(
  %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
    -> tensor<?x?xf32> {
  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
                     outs(%arg2: tensor<?x?xf32>)
    -> tensor<?x?xf32>

  return %0 : tensor<?x?xf32>
}
```

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