[Mlir-commits] [mlir] [Vector] Add canonicalization for select(pred, true, false) -> broadcast(pred) (PR #147934)

[Kunwar Grover]

Groverkss wrote:

> Thanks!
> 
> > Adds a canonicalization pattern for vector.broadcast to fold:
> 
> Hm, to me this is "rewriting" `arith.select` as `vector.broadcast` (as in, intuitively this isn't really folding). Why is `vector.broadcast` more desirable than `arith.select`?
> 
> Lets discuss this first - you can skip my other comments in the meantime.

arith.select evaluates the condition and selects between the true and false vector. broadcasts directly propagates the condition forward. It also unlocks multiple folding opportunities. For example:

```
%select = arith.select %pred, %true, %false : vector<4xi1>
%broadcasted = vector.broadcast %select : vector<2x4xi1>
%load = vector.maskedload .... /*mask=*/%broadcasted
```

here, in the current ir, it is quite hard to actually reason if the vector.maskedload mask is either true or false (and not partially true/false), which prevents folding opportunities. On the other hand:

```
%select = vector.broadcast %pred : vector<4xi1>
%broadcasted = vector.broadcast %select : vector<2x4xi1>
%load = vector.maskedload .... /*mask=*/%broadcasted
```

will fold into:

```
%select = vector.broadcast %pred : vector<2x4xi1>
%load = vector.maskedload .... /*mask=*/%broadcasted
```

and now it is easy to match that the mask is just a broadcasted conditional, so it must be either always full or always empty.

vector.broadcast has very clear semantics that the entire vector has the same value, on the other hand, for select, we actually have to traverse the true/false values to check if the entire vector has the same value.

I will also note that LLVM InstCombine does a similar canonicalization as noted above: https://github.com/llvm/llvm-project/blob/main/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp#L2351

The LLVM canonicalization converts the i1 pred into a broadcasted vector i1 pred, which further just folds away directly, completing the same fold we are doing here.

Personally, for me, this is a very clear canonicalization. vector.broadcast has much more restricted semantics vs arith.select and composes better with other vector operations (unless we plan to write folders for every vector operation that interacts with broadcast, to also work with (select pred, true , false) which is very counter intuitive).

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