[Mlir-commits] [mlir] [mlir][vector] Patterns to convert to shape_cast, where possible (PR #138777)

[Andrzej Warzyński]

banach-space wrote:

> I found https://github.com/llvm/llvm-project/commit/06dbb2856967a5497c6ddfad3d3fdfea20849f7e which removes patterns exactly like these, on the caoincalizers.

Thanks for finding and sharing that thread :) I wasn't aware of those earlier design discussions about restricting `vector.shape_cast`.

That’s a fairly old discussion (given MLIR’s age), and likely dates back to when there was an assumption we'd be using `llvm.intr.matrix` intrinsics. From the [vector.shape_cast documentation](https://mlir.llvm.org/docs/Dialects/Vector/#vectorshape_cast-vectorshapecastop):

> There is an exception to the folding expectation when targeting llvm.intr.matrix operations. We need a type conversion back and forth from a 2-D MLIR vector to a 1-D flattened LLVM vector.shape_cast lowering to LLVM is supported in that particular case, for now.

However, these matrix intrinsics are not widely used, and - as far as I know - the corresponding LLVM work has been on hold for quite some time (a few years?). From what I can tell, those requirements are now outdated.

In the last 2–3 years that I’ve been working in this space, the consistent takeaway has been:

> vector.shape_cast is the canonical form for reshapes that imply no data movement.

In general, the assumption is that `vector.shape_cast` encodes a reshape with no underlying value movement.

Regarding the example you mentioned:
```mlir
%0 = vector.shape_cast %arg0 : vector<2x3xf32> to vector<1x2x3xf32>
```
...and its lowering to:

```mlir
 %0 = ub.poison : vector<1x2x3xf32>
  %1 = vector.extract %arg0[0, 0] : f32 from vector<2x3xf32>
  %2 = vector.insert %1, %0 [0, 0, 0] : f32 into vector<1x2x3xf32>
  %3 = vector.extract %arg0[0, 1] : f32 from vector<2x3xf32>
  %4 = vector.insert %3, %2 [0, 0, 1] : f32 into vector<1x2x3xf32>
  %5 = vector.extract %arg0[0, 2] : f32 from vector<2x3xf32>
  %6 = vector.insert %5, %4 [0, 0, 2] : f32 into vector<1x2x3xf32>
  %7 = vector.extract %arg0[1, 0] : f32 from vector<2x3xf32>
  %8 = vector.insert %7, %6 [0, 1, 0] : f32 into vector<1x2x3xf32>
  %9 = vector.extract %arg0[1, 1] : f32 from vector<2x3xf32>
  %10 = vector.insert %9, %8 [0, 1, 1] : f32 into vector<1x2x3xf32>
  %11 = vector.extract %arg0[1, 2] : f32 from vector<2x3xf32>
  %12 = vector.insert %11, %10 [0, 1, 2] : f32 into vector<1x2x3xf32>
```

This lowering clearly involves data movement, and to me violates the design intent of `vector.shape_cast`. I wouldn’t use it as a guiding example - in fact, this may be a good time to revisit that lowering and consider disabling it or making it conditional.

So in principle: yes, I’m in favor of turning these patterns into canonicalizations. We should just be thoughtful about which constraints we want to enforce.

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