[Mlir-commits] [mlir] [mlir][vector] Add vector.transpose with unit-dim to vector.shape_cast pattern (PR #72105)

[llvmlistbot at llvm.org]

MaheshRavishankar wrote:

> I'm OOO this week and afk but I wanted to drop some quick comments until I'm back next week. Hopefully they are helpful!
> 
> * We are using `vector.shape_cast` in a few places, including some lowerings of `vector.transpose` so this PR is not adding anything fundamentally new.

It is since this would not happen at any point in the stack and now it is. So it is changing the behavior of an entire compiler stack. Moving to canonicalization have a high impact and should be done very deliberately. We've over-index on using canonicalizations for ease (I have done it many times do, and I am sure I will do it again accidentally). There are many cases where we run certain "foldings" multiple times, but they are all opt-in and I think that is fine. Canonicalizations are too heavy a hammer with almost unbounded side-effect.

> * `vector.shape_cast` has effectively been used as a vector reshape for static shapes, which is a no-op in practice. We had some implementation gaps that lowered to vector shuffles but I haven't seen that lowering kicking in in quite some time. I think these lowerings were introduced for practical reasons (if I hit an implementation gap I just lower to inefficient code instead of crashing). Nicolas, correct me if I'm wrong. **Perhaps the solution to this problem should go along the lines of tightening `vector.shape_cast` to that actual use case in practice?**
> * `vector.transpose` is a heavy operation which implies data movements so preserving a no-op `vector.transpose` has two main downsides:
> * 1. We will lower them to redundant data movement ops, which hopefully the backend compiler will optimize away (guess what, LLVM doesn't always :)). Otherwise, we would have to add special lowering for "no-op" transposes, which would make the already complex transpose lowerings even more complex. I don't think the added complexity would be justified when we have simpler ways to model this.

I understand this. We should not be lowering `vector.transpose` as is, but some lowering flow lowers a `vector.transpose` to `vector.shape_cast` is fine (thats what happened in the first PR of this chain of 2). I dont have an issue with that, except maybe it should be structured so that it is opt-out if needed (or opt-in), as it is here for SPIR-V lowering. 

> * 2. From an analysis point of view, a `vector.transpose` is an expensive op, which would be misleading for this case. This case is similar to optimizing away 1D transposes or transpose patterns that are not transposing any dimensions. In this sense, we have a bunch of patterns to optimize transposes away so I don't see this case any different.
> * I don't think `vector.transpose` should be a canonical representation of `1xN -> Nx1`. There is no data movement here so using a `vector.transpose` op would be misleading. There are also different vector ops that could implement the same transformation, e.g., `vector.extract` + `vector.broadcast`, so I don't think canonicalizing all of them to something as involved as a transpose op would work.

I think we are mixing lowering with canonicalization. `1xN -> Nx1` is a transpose... it can also be "implemented" as a no-op `shape_cast`, but that is a lowering optimization, not a canonical representation.

> * Not sure I understand the SPRIV-V problem. If the lowering supports `1xN -> Nx1` transposes, how is that it can't support a simple static reshape? Is this just a missing rewrite or is something fundamental missing?

This is exactly the point I am trying to make. In SPIR-V to lower such a `vector.shape_cast`, we basically need to recognize this as a `vector.transpose` and lower it that way, which is effectively "undoing" a canonicalization. To me that is an indication that `vector.transpose` should be left as is, and on the lowering flow to LLVM the lowering  should be optimized to a no-op `vector.shape_cast`.

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