This is an archive of the discontinued LLVM Phabricator instance.

[mlir][linalg] Add a test for linalg.matmul --> vector.outerproduct
ClosedPublic

Authored by awarzynski on May 12 2023, 9:42 AM.

Details

Summary

Representing matmuls as a sum of outer products is central to various
matrix extensions (e.g. Arm's SME). This test demonstrates how to use
Linalg's vectoriser and Vector's lowerings to represent linalg.matmul
as a chain of vector.outerproduct Ops.

Diff Detail

Event Timeline

awarzynski created this revision.May 12 2023, 9:42 AM
Herald added a project: Restricted Project. · View Herald Transcript
awarzynski requested review of this revision.May 12 2023, 9:42 AM
springerm accepted this revision.May 12 2023, 9:44 AM
This revision is now accepted and ready to land.May 12 2023, 9:44 AM
thopre added inline comments.
mlir/test/Dialect/Linalg/vectorization.mlir
85 ↗(On Diff #521690)

I don't think CHECK-SAME is well defined after a CHECK-LABEL. It might work now but conceptually it feels wrong to me. It should be after a CHECK: I'd rather duplicate the CHECK func.func @outerproduct_matmul

@jdenny Is CHECK-SAME after CHECK-LABEL officially supported? I couldn't find anything obvious in the documentation.

thopre added inline comments.May 12 2023, 2:18 PM
mlir/test/Dialect/Linalg/vectorization.mlir
85 ↗(On Diff #521690)

That said, a lot of tests already follow this pattern so maybe it should just be documented as expected to work. In other words, no objection on this patch.

Thanks @awarzynski! LGTM, just one comment.

mlir/test/Dialect/Linalg/vectorization.mlir
110 ↗(On Diff #521690)

This looks like a small integration test. Should we create two tests, one for vectorize and one for lower_contraction? If not, I think we should move it to a different .mlir file as here we are only testing vectorize.

awarzynski added inline comments.May 13 2023, 3:07 AM
mlir/test/Dialect/Linalg/vectorization.mlir
110 ↗(On Diff #521690)

I agree and am glad that you have pointed this out :)

  1. My main goal is to have an in-tree example/test that demonstrates how to get from linalg.matmul to vector.outerproduct. So, I prefer 1 integration test rather than 2 regression tests.
  2. All "proper" Integration tests in MLIR (in mlir/test/Integration) seem to be end-to-end (i.e. "compiler" + "run" + "check the runtime output"). This test is much simpler than that, so I'd keep it in this directory.

An updated version is on its way :)

Move the test to a dedicated file

Btw, I've just learnt that matmul_step_04_tile_and_vector_mask.mlir in IREE demonstrates similar lowering path (albeit much more elaborate). @qcolombet, thank you for sharing!

Cool! Then perhaps moving it to a separate file would be enough.

dcaballe accepted this revision.May 16 2023, 12:09 AM

Oops, I didn't see the update :)

mlir/test/Dialect/Linalg/vectorization.mlir
110 ↗(On Diff #521690)

@dcaballe note that the separate unit tests you mention exist already.

xgupta added a subscriber: xgupta.May 16 2023, 2:07 AM
xgupta added inline comments.
mlir/test/Dialect/Linalg/transform-op-matmul-to-outerproduct.mlir
1

Just noticed, does it requires -split-input-file option here? I think it work without it also.

awarzynski added inline comments.May 16 2023, 2:54 AM
mlir/test/Dialect/Linalg/transform-op-matmul-to-outerproduct.mlir
1
xgupta added inline comments.May 16 2023, 3:18 AM
mlir/test/Dialect/Linalg/transform-op-matmul-to-outerproduct.mlir
1

Thanks for addressing the comment!

Matt added a subscriber: Matt.May 19 2023, 2:21 PM