[Mlir-commits] [mlir] [mlir][vector] Support multi-dimensional vectors in VectorFromElementsLowering (PR #151175)

[Kunwar Grover]

Groverkss wrote:

> I don’t think we can make a call on unrolling vs linearization. Unrolling will bloat the code size when unrolling a large dimension whereas linearization will generate fewer ops (best case a single op). Vector shuffles will be generated anyways by LLVM regardless of what we do at MLIR level. The right call is probably project dependent.
> 
> >

I agree the right call is project dependent, but I don't agree with "vector shuffles will be generated anyways by LLVM...". Take for example:

a = load : vector<2x8xf32>
b = elementwise a : vector<2x8xf32>
store b : vector<2x8xf32>

Unrolling:

a_0 = load: vector<8xf32>
a_1 = load: vector<8xf32>
b_0 = elementwise a_0 : vector<8xf32>
b_1 = elemenetwise a_1 : vector<8xf32>
store b_0 : vector<8xf32>
store b_1 : vector<8xf32>

Flattening:

// Loads cannot be flattened, they have to be unrolled and then shuffled into a single vector
a_0 = load: vector<8xf32>
a_1 = load: vector<8xf32>
a = shuffle a_0, a_1 : vector<8xf32>, vector<8xf32> -> vector<16xf32>
b = elementwise a : vector<8xf32>
b_0 = shufflevector b, poison : vector<8xf32>, vector<8xf32> -> vector<4xf32>
b_1 = shufflevector b, poison : vector<8xf32>, vector<8xf32> -> vector<4xf32>
// Stores cannot be flattened, they have to be sliced using shuffles
store b_0 : vector<8xf32>
store b_1: vector<8xf32>
``

In the unrolling IR, we have no shuffles, while in the flattening IR we do have shuffles.

This is what I meant by my comment, that locally, flattening looks nice, because you get a single operation on a wide vector, but it has implications across the entire IR which need to be considered properly.

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