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

[awarzynski]

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.

[thopre]

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]

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.

[dcaballe]

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]

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 :)

[awarzynski]

Move the test to a dedicated file

[awarzynski]

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!

[dcaballe]

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

[dcaballe]

Oops, I didn't see the update :)

[nicolasvasilache]

mlir/test/Dialect/Linalg/vectorization.mlir

110 ↗

(On Diff #521690)

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

[xgupta]

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]

mlir/test/Dialect/Linalg/transform-op-matmul-to-outerproduct.mlir

1

Good catch and sorry that I missed that! Fixed here: https://github.com/llvm/llvm-project/commit/c2d55ef563a96933db8c171a7b624a75b67c2a0f

[xgupta]

mlir/test/Dialect/Linalg/transform-op-matmul-to-outerproduct.mlir

1

Thanks for addressing the comment!

mlir/test/Dialect/Linalg/transform-op-matmul-to-outerproduct.mlir

+// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file | FileCheck %s
+
+func.func @outerproduct_matmul(%A: memref<3x3xf32>, %B: memref<3x3xf32>, %C: memref<3x3xf32>) {
+  linalg.matmul ins(%A, %B: memref<3x3xf32>, memref<3x3xf32>)
+            outs(%C: memref<3x3xf32>)
+  return
+}
+
+// CHECK-LABEL:   func.func @outerproduct_matmul(
+// CHECK-SAME:      %[[VAL_0:.*]]: memref<3x3xf32>, %[[VAL_1:.*]]: memref<3x3xf32>, %[[VAL_2:.*]]: memref<3x3xf32>) {
+// CHECK:           %[[VAL_3:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_4:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK:           %[[VAL_5:.*]] = vector.transfer_read %[[VAL_0]]{{\[}}%[[VAL_3]], %[[VAL_3]]], %[[VAL_4]] {in_bounds = [true, true]} : memref<3x3xf32>, vector<3x3xf32>
+// CHECK:           %[[VAL_6:.*]] = vector.transfer_read %[[VAL_1]]{{\[}}%[[VAL_3]], %[[VAL_3]]], %[[VAL_4]] {in_bounds = [true, true]} : memref<3x3xf32>, vector<3x3xf32>
+// CHECK:           %[[VAL_7:.*]] = vector.transfer_read %[[VAL_2]]{{\[}}%[[VAL_3]], %[[VAL_3]]], %[[VAL_4]] {in_bounds = [true, true]} : memref<3x3xf32>, vector<3x3xf32>
+// CHECK:           %[[VAL_8:.*]] = vector.transpose %[[VAL_5]], [1, 0] : vector<3x3xf32> to vector<3x3xf32>
+// CHECK:           %[[VAL_9:.*]] = vector.extract %[[VAL_8]][0] : vector<3x3xf32>
+// CHECK:           %[[VAL_10:.*]] = vector.extract %[[VAL_6]][0] : vector<3x3xf32>
+// CHECK:           %[[VAL_11:.*]] = vector.outerproduct %[[VAL_9]], %[[VAL_10]], %[[VAL_7]] {kind = #vector.kind<add>} : vector<3xf32>, vector<3xf32>
+// CHECK:           %[[VAL_12:.*]] = vector.extract %[[VAL_8]][1] : vector<3x3xf32>
+// CHECK:           %[[VAL_13:.*]] = vector.extract %[[VAL_6]][1] : vector<3x3xf32>
+// CHECK:           %[[VAL_14:.*]] = vector.outerproduct %[[VAL_12]], %[[VAL_13]], %[[VAL_11]] {kind = #vector.kind<add>} : vector<3xf32>, vector<3xf32>
+// CHECK:           %[[VAL_15:.*]] = vector.extract %[[VAL_8]][2] : vector<3x3xf32>
+// CHECK:           %[[VAL_16:.*]] = vector.extract %[[VAL_6]][2] : vector<3x3xf32>
+// CHECK:           %[[VAL_17:.*]] = vector.outerproduct %[[VAL_15]], %[[VAL_16]], %[[VAL_14]] {kind = #vector.kind<add>} : vector<3xf32>, vector<3xf32>
+// CHECK:           vector.transfer_write %[[VAL_17]], %[[VAL_2]]{{\[}}%[[VAL_3]], %[[VAL_3]]] {in_bounds = [true, true]} : vector<3x3xf32>, memref<3x3xf32>
+// CHECK:           return
+// CHECK:         }
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
+  %2 = transform.structured.vectorize %1
+  transform.vector.lower_contraction %2 lowering_strategy = "outerproduct" : (!pdl.operation) -> !pdl.operation
+}