[all-commits] [llvm/llvm-project] c91d3b: [mlir][vector] Constrain patterns: vector.contract...

[Andrzej Warzyński via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: c91d3b0b08ee133573b4c24f239306a9f1741174
      https://github.com/llvm/llvm-project/commit/c91d3b0b08ee133573b4c24f239306a9f1741174
  Author: Andrzej Warzynski <andrzej.warzynski at arm.com>
  Date:   2023-10-06 (Fri, 06 Oct 2023)

  Changed paths:
    M mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
    A mlir/test/Dialect/Vector/vector-contract-to-outerproduct-transforms-unsupported.mlir
    M mlir/test/Dialect/Vector/vector-contract-to-outerproduct-transforms.mlir

  Log Message:
  -----------
  [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