[Mlir-commits] [mlir] [MLIR][Vector] Add unroll pattern for vector.shape_cast (PR #164010)

[Andrzej Warzyński]

banach-space wrote:

> Hi @banach-space , my understanding is that ShapeCastOpRewritePattern works directly with the full transformation but Unrolling works in target-sized tiles (2x2), then assembles them. I also see there are similar patterns for other ops as well like BroadcastOpLowering (LowerVectorBroadcast.cpp) and UnrollBroadcast.

I’d like to re-visit this thread (apologies for not paying closer attention earlier).

To me, at a high level, this transformation fits more naturally under **lowering** rather than **unrolling**. Here’s why:

**UNROLLING**
```mlir
// BEFORE
vector.broadcast %v : vector<4xf32> to vector<4x4xf32>
// AFTER (vector.broadcast on lower rank vectors)
%v0 = vector.extract_strided_slice  : vector<4xf32> to vector<2xf32>
%v0_bc = vector.broadcast %v0 : vector<2xf32> to vector<2x2xf32>
vector.insert_strided_slice %v0_bc : vector<2x2xf32> into vector<4x4xf32>
(...)
```

**LOWERING**
```mlir
// BEFORE
vector.broadcast %v : vector<4xf32> to vector<4x4xf32>
// AFTER (no vector.broadcast)
%0 = ub.poison : vector<4x4xf32>
%1 = vector.insert %arg0, %0 [0] : vector<4xf32> into vector<4x4xf32>
%2 = vector.insert %arg0, %1 [1] : vector<4xf32> into vector<4x4xf32>
%3 = vector.insert %arg0, %2 [2] : vector<4xf32> into vector<4x4xf32>
%4 = vector.insert %arg0, %3 [3] : vector<4xf32> into vector<4x4xf32>
```

**In words:**
* Unrolling a `vector.<op_name>` that operates on rank-N vectors rewrites it as a sequence of `vector.<op_name>`s that operate on rank-K vectors, where K < N.
* Lowering a vector operation `vector.<op_name>` on rank-N vectors rewrites it as a sequence of other vector operations that we consider “lower-level”.

This patch implements the latter pattern.

So rather than adding it as an unrolling pattern, I’d suggest refactoring the existing lowering pattern to take the target vector shape as a parameter. From what I can tell, this aligns with @dcaballe 's [suggestion](https://github.com/llvm/llvm-project/pull/164010#issuecomment-3438225952) to re-use more code.

Thoughts?

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