[mlir][vector] Constrain patterns: vector.contract -> vector.outerproduct
This patch constrains the patterns for converting `vector.contract` to
`vector.outerproduct` so that
* the reduction dimension is _not unrolled_ if the corresponding
dimension is scalable.
This is necessary as the current lowering is incorrect for scalable
dims. Indeed, the following unrolling for `vector.contract` would be
invalid if the corresponding dimension was scalable (K is the size of
the reduction dimension):
```
// K times. This is valid if K _is not_ scalable.
%lhs = vector.extract %LHS[0]
%rhs = vector.extract %RHS[0]
vector.outerproduct %lhs, %rhs
%lhs = vector.extract %LHS[1]
%rhs = vector.extract %RHS[1]
vector.outerproduct %lhs, %rhs
// ...
```
Instead, a `for` loop should be generated:
```
// This would be valid regardless of whether K is scalable or not
scf.for %k = 0 to K step 1
%lhs = vector.extract LHS[%k]
%rhs = vector.extract RHS[%k]
vector.outerproduct %lhs, %rhs
```
However, the lowering of:
* `vector.extract` of vector slices with dynamic indices
is incomplete and hence the implementation proposed above (with
`scf.for`) wouldn't work just yet, i.e. it wouldn't be possible to lower
it further. Instead, this patch disables unrolling in cases when the
reduction dimension is scalable, i.e. where the generated code would be
functionally incorrect.
In order to document unsupported cases, a dedicated test file is added:
* "vector-contract-to-outerproduct-transforms-unsupported.mlir"
This is the first patch in a series of patches that strives to update
these patterns (and to test them) for scalable vectors.
Resolves #68400
parent
092ef55a
Please register or sign in to comment