[Mlir-commits] [mlir] [mlir][linalg] Vectorize directly to a named contraction (PR #147296)

[Adam Siemieniuk]

adam-smnk wrote:

I think we're talking about two slightly different things. Let me try to clarify.

`linalg.contract` has unique maps as in the operation requires to have its `indexing_maps` to be explicitly defined. As in your above example, `contract` already has maps and generalization only has to generate correct body. It's basically a 1-to-1 conversion.

For the earlier case I took your statement:
> When we broadcast like this, the maps should be the default indexing maps for linalg.matmul

as in resolving broadcasting for this matmul:
```mlir
%0 = linalg.matmul
    indexing_maps = [affine_map<(m, n, k) -> (k)>,  // broadcast LHS
                      affine_map<(m, n, k) -> (k)>, // broadcast RHS
                      affine_map<(m, n, k) -> (m, n)>]
```
should yield
```mlir
    indexing_maps = [affine_map<(m, n, k) -> (m, k)>,  // broadcast LHS
                      affine_map<(m, n, k) -> (k, n)>, // broadcast RHS
                      affine_map<(m, n, k) -> (m, n)>]
```
as these can be derived from the default indexing maps for `linalg.matmul.

I agree that's a valid strategy. However, the same `linalg.contract` with implicit (broadcasted) parallel dimensions on LHS and RHS can be realized  in any possible permutation of these maps as there's no default definition here. For example:
```mlir
indexing_maps = [affine_map<(m, n, k) -> (m, n, k)>,
                     affine_map<(m, n, k) -> (k)>,
                     affine_map<(m, n, k) -> (m, n)>]
```
is a valid interpretation for broadcast decomposition.

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