iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.
http://iree.dev/
Apache License 2.0
2.56k stars 568 forks source link

Adding support for scalable vectors to the tiling and vectorization constructs in the Transform dialect #14275

Open banach-space opened 1 year ago

banach-space commented 1 year ago

[Implements #13334 ]

As per https://discourse.llvm.org/t/rfc-scalable-vectorisation-in-linalg/, add support for scalable tile sizes:

transform.structured.tile %0 [2, [4]]]

and scalable vector sizes:

transform.structured.masked_vectorize %0 vector_sizes [[4]]]

Ultimately, we should support the following (combined tiling and vectorisation):

%vscale_base = transform.param.const 4 : i64 -> index
%tile_size, %split_point = 
  transform.vector.compute_scalable_tile_size_and_split_point
  base(%vscale_base) (%op)
%main, %remainder = transform.structured.split %op after %split_point
%main_t = transform.structured.tile %main [%tile_size]
transform.structured.vectorize %main_t
transform.structured.masked_vectorize %remainder [[%vscale_base]]

CC @dcaballe

banach-space commented 1 year ago

The patches so far:

This is sufficient for scalable vectorisation and tiling, but only allows for the trailing vector/tile size to be scalable.