[Mlir-commits] [mlir] [vector][mlir] Canonicalize to shape_cast where possible (PR #140583)

[Kunwar Grover]

Groverkss wrote:

The main issue with this canonicalization comes down to the fact that it drops
structural information of operations, which goes against "structured code
generation".

The crux of the issue raised in the previous prs, reverts and discourse
discussion is that shape_cast is not a "structured" N-D transformation, while
transpose, broadcast and extract are.

This is visible when describing each of these operations based on how their
result dimensions map to their input dimensions:

```
shape_cast -> output_dims = delinearize(linearize(input_dims, basis=input_shape), basis=output_shape)
transpose -> output_dims = permute(input_dims)
broadcast -> output_dims = concat(new_dims, input_dims)
extract -> output_dims = drop_front(input_dims, num_dropped=old_rank-new_rank)
```

Note that transpose, broadcast and extract represent a projected permutation
transformation while shape_cast represents an arbritary reshape transformation
which is less restricted than a projected permutation.

When coming from tensor land via vectorization, the vector code produced
preserves this structural information which is very useful. There are no
arbritary reshapes with linearize/delinearize semantics, and dimensional
structure is preserved.

This PR tries to canonicalize different ops which can do different things but have
a single point where they are similar, to a single, op which has the least
structure preserving semantics of them all.

The reasonings I'm seeing in this PR are:

> The point of having a canonical representation is to choose a unique
> representation for a computation that may have multiple alternative
> representations. That form doesn't have to be the "best" representation
> (whatever that means: it doesn’t have to be the most efficient, the most
> compact, the most aesthetically pleasing, the most restrictive…) It has to be
> unique and hold all the information that the original form held. 

> The original Ops are mechanically recoverable in every case, so I don't see
> any loss of information here.

I don't see canonicalizing towards a less restricted op to be a goal of
canonicalization at all:
https://mlir.llvm.org/docs/Canonicalization/#general-design

Infact, that is only just the initial bar for something to be considered as a
canonicalization:

> Canonicalize shouldn’t lose the semantic of original operation: the original
> information should always be recoverable from the transformed IR.

It's not the bar for something to be a canonicalization.

The goal that is relevant here, that I do see is:

> The goal of canonicalization is to make subsequent analyses and optimizations
> more effective. Therefore, performance improvements are not necessary for
> canonicalization.

There is no justification in this PR that canonicalizing towards a single operation
makes subsequent analysis and optimizations more effective. Infact, it requires
transformations to now add an edge case for special `shape_cast` that
could be the special form it was looking for. For example:

Let's say I'm writing a transformation (a commonly written one), which folds
a transpose-like op into a permutation map.

```
transpose 1x2x2x8 -> 2x2x8x1
contract { indexing_maps = [#map] }

transform to

contract { indexing_maps = [#transposed_map] }
```

It is easy to write, with just matching operations that have transpose like
semantics.

```
shape_cast 8x8 -> 2x2x8x1
contract { indexing_maps = [#map] }
```

Not only does my transformation has to look for transpose-like ops, it has to
look for shape_cast's that are special cased to be transpose-like.

This is making subsequent transformations harder to apply, this is actively going
against the goal of canonicalizations.

In practice, we have number of patterns in SPIRV backend in our downstream
compiler (IREE), which do transformations that look for
transpose/broadcast/extract/insert ops and use their semantics to analyze
before doing transformations (They don't rely on them for correctness). These
patterns will not run anymore, because they all need to handle a special case
of shape_cast acting like other ops, even though the pipeline never introduced
it.

Some examples (from @kuhar):

https://github.com/iree-org/iree/blob/main/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVInitialVectorLowering.cpp#L148
https://github.com/iree-org/iree/blob/main/compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp#L509

I'm sure I could find more examples from upstream, but these are just from
the top of my head.

Keeping this in mind, I don't agree this transformation is a canonicalization
for the entire vector dialect. It's okay to add this as a transformation pass,
pattern methods, or anything, as long as it's opt-in and doesn't force every
user of vector dialect to use this.

If you think otherwise, please send a RFC and we can discuss further there. This
is a very wide reaching change and I would like you to address why it's okay to
throw semantic information like projected permutation semantics and ask
transformations to analyze it back.

The reason I ask for an RFC is that there are justifications given in this PR
which I have no meaningful way to measure, and it falls on the author to prove
that this transformation is a canonicalization:

> We currently have 3+ ways to drop unit dimensions and perform vector reshapes
> that don’t imply data movement in the virtual vector world and we are paying
> the cost of maintaining all those redundant patterns. I think the big
> misalignment here is that we expect the canonical form to be the ideal input
> for all our use cases without doing any pre-processing and that is something
> we need to change for the shake of the overall project health.

A simple example of why this isn't a good arguement is we also have linalg.contraction, linalg.generic and linalg.matmul in the same dialect, but we don't choose to canonicalize them to each other. We have 3+ ways of writing a matmul in linalg. We should be consistent with this argument and apply it everywhere if so.

What I would like to see in the RFC is:

1. Show what "redundant" patterns we remove from having this canonicalization.
2. Measure the impact of upstream patterns where every pattern looking for
   a transpose/broadcast/extract now has to also account for a shape_cast,
   because one edge case falls through to that.
3. Show why this canonicalization aligns with the goals of canonicalizations.

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