[Mlir-commits] [mlir] [mlir][Vector] Add `vector.shuffle` tree transformation (PR #145740)

[Diego Caballero]

dcaballe wrote:

> My understanding is that the total number of shuffles will still be O(T)

The number of shuffles is linear to the number of `vector.to_elements`, regardless of the combination pattern. This is guaranteed by the following properties (from the doc):

```
///   1. Vectors are shuffled in pairs by order of appearance in
///      the `vector.from_elements` operand list.
///   2. Each input vector to each level is used only once.
```

More "scattered" patterns will lead to wider `vector.shuffle` ops but not to more or less `vector.shuffle` ops.

In terms of code size, the actual "gain" comes when we lower the `vector.from/to_elements` vs `vector.shuffle` ops, let's say, to LLVM. The former will lower to an explosion of vector insertelement/extractelement operations while the latter will directly lower to the shufflevector counterpart operations.

> Only partially related to this PR, but I also wonder if there is something like this is LLVMIR, where a sequence of shufflevector instructions is turned into a tree to reduce the average mask size.

LLVM's InstCombine implements something similar but at a much higher cost. It goes through the large sequences of extractelement/insertelement instructions to reconstruct the information that `vector.from/to_elements` provide "for free" (e.g., group extractelement/insertelement instructions extracting/inserting from/into the same vector, determine elements that are actually used, replaced those unused with poison values...). To contextualize "higher cost", we have a case where this InstCombine transformation accounts for >90% of the E2E compilation time. Enabling this transformation makes InstCombine compilation time negligible... same as the compilation time of this transformation itself :).

InstCombine also applies some canonicalization to shufflevector instructions but I'm not aware of any major transformation happening to the actual tree shuffle sequence generated by this transformation.

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