[Mlir-commits] [mlir] [mlir][linalg] Add an e2e test for linalg.matmul_transpose_a to ArmSME (PR #71644)
Cullen Rhodes
llvmlistbot at llvm.org
Wed Nov 8 08:03:28 PST 2023
https://github.com/c-rhodes updated https://github.com/llvm/llvm-project/pull/71644
>From 1d828fd981720d790e14f6896d8d6ce42340d7db Mon Sep 17 00:00:00 2001
From: Cullen Rhodes <cullen.rhodes at arm.com>
Date: Tue, 31 Oct 2023 11:07:59 +0000
Subject: [PATCH 1/3] [mlir][linalg] Add an e2e test for
linalg.matmul_transpose_a to ArmSME
This patch adds an integration test demonstrating the first e2e example
lowering a linalg.matmul to SME via vector.outerproduct.
The test uses a 'linalg.matmul_transpose_a' rather than 'linalg.matmul'
since the latter emits a 'vector.transfer_read' with a vector type of
'vector<[4]x1xf32>' that can't be currently lowered via generic (SVE)
path, since it has leading scalable dim.
---
.../Linalg/CPU/ArmSME/matmul-transpose-a.mlir | 77 +++++++++++++++++++
1 file changed, 77 insertions(+)
create mode 100644 mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir
new file mode 100644
index 000000000000000..bf445e05fb70e8d
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir
@@ -0,0 +1,77 @@
+// RUN: mlir-opt %s \
+// RUN: -transform-interpreter -test-transform-dialect-erase-schedule \
+// RUN: -one-shot-bufferize="bufferize-function-boundaries" -canonicalize \
+// RUN: -enable-arm-streaming="mode=locally enable-za" \
+// RUN: -convert-vector-to-arm-sme -convert-arm-sme-to-scf \
+// RUN: -convert-vector-to-scf -cse -arm-sve-legalize-vector-storage \
+// RUN: -convert-vector-to-llvm=enable-arm-sme \
+// RUN: -convert-vector-to-llvm=enable-arm-sve \
+// RUN: -cse -canonicalize -allocate-arm-sme-tiles -test-lower-to-llvm | \
+// RUN: %mcr_aarch64_cmd \
+// RUN: -e=main -entry-point-result=void \
+// RUN: -march=aarch64 -mattr="+sve,+sme" \
+// RUN: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils | \
+// RUN: FileCheck %s
+
+func.func @matmul_transpose_a(%A : tensor<?x?xf32>, %B : tensor<?x?xf32>, %C : tensor<?x?xf32>) {
+ %res = linalg.matmul_transpose_a ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32>
+ %xf = tensor.cast %res : tensor<?x?xf32> to tensor<*xf32>
+ call @printMemrefF32(%xf) : (tensor<*xf32>) -> ()
+ return
+}
+
+func.func @main() {
+ %c0 = arith.constant 0 : i32
+ %c4 = arith.constant 4 : index
+
+ %A = arith.constant dense<[
+ [ 1.0, 2.0, 3.0, 4.0 ],
+ [ 5.0, 6.0, 7.0, 8.0 ],
+ [ 9.0, 10.0, 11.0, 12.0 ],
+ [ 13.0, 14.0, 15.0, 16.0 ]
+ ]> : tensor<4x4xf32>
+
+ %A_dyn = tensor.cast %A : tensor<4x4xf32> to tensor<?x?xf32>
+
+ %C_init = bufferization.alloc_tensor(%c4, %c4) : tensor<?x?xf32>
+ %C = linalg.fill ins(%c0 : i32) outs(%C_init : tensor<?x?xf32>) -> tensor<?x?xf32>
+
+ // CHECK: Unranked Memref {{.*}} rank = 2 offset = 0 sizes = [4, 4] strides = [4, 1] data =
+ // CHECK: [276, 304, 332, 360]
+ // CHECK: [304, 336, 368, 400]
+ // CHECK: [332, 368, 404, 440]
+ // CHECK: [360, 400, 440, 480]
+ call @matmul_transpose_a(%A_dyn, %A_dyn, %C) : (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) -> ()
+
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%module : !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.matmul_transpose_a"]} in %module
+ : (!transform.any_op) -> !transform.any_op
+ %tiled_linalg_op, %loops:3 = transform.structured.tile_using_for %0[[4], [4], 1]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ transform.structured.vectorize %tiled_linalg_op vector_sizes [[4], [4], 1]
+ : !transform.any_op
+
+ %func = transform.structured.match ops{["func.func"]} in %module
+ : (!transform.any_op) -> !transform.any_op
+
+ transform.apply_patterns to %func {
+ transform.apply_patterns.vector.lower_masked_transfers
+ transform.apply_patterns.vector.transfer_permutation_patterns
+ transform.apply_patterns.vector.reduction_to_contract
+ } : !transform.any_op
+
+ transform.apply_patterns to %func {
+ transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
+ transform.apply_patterns.vector.lower_masks
+ } : !transform.any_op
+
+ transform.yield
+ }
+}
+
+func.func private @printMemrefF32(%ptr : tensor<*xf32>)
>From 6cb5e9a2a28006209b0dbf8fa5df3c4699cb04dd Mon Sep 17 00:00:00 2001
From: Cullen Rhodes <cullen.rhodes at arm.com>
Date: Wed, 8 Nov 2023 14:48:38 +0000
Subject: [PATCH 2/3] address comments
---
.../Linalg/CPU/ArmSME/matmul-transpose-a.mlir | 49 +++++++++++++------
1 file changed, 34 insertions(+), 15 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir
index bf445e05fb70e8d..dda8c93380768c9 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir
@@ -23,25 +23,37 @@ func.func @matmul_transpose_a(%A : tensor<?x?xf32>, %B : tensor<?x?xf32>, %C : t
func.func @main() {
%c0 = arith.constant 0 : i32
- %c4 = arith.constant 4 : index
+ %c7 = arith.constant 7 : index
%A = arith.constant dense<[
- [ 1.0, 2.0, 3.0, 4.0 ],
- [ 5.0, 6.0, 7.0, 8.0 ],
- [ 9.0, 10.0, 11.0, 12.0 ],
- [ 13.0, 14.0, 15.0, 16.0 ]
- ]> : tensor<4x4xf32>
+ [ 1., 2., 3., 4., 5., 6., 7.],
+ [ 8., 9., 10., 11., 12., 13., 14.],
+ [15., 16., 17., 18., 19., 20., 21.],
+ [22., 23., 24., 25., 26., 27., 28.],
+ [29., 30., 31., 32., 33., 34., 35.],
+ [36., 37., 38., 39., 40., 41., 42.],
+ [43., 44., 45., 46., 47., 48., 49.],
+ [50., 51., 52., 53., 54., 55., 56.],
+ [57., 58., 59., 60., 61., 62., 63.],
+ [64., 65., 66., 67., 68., 69., 70.],
+ [71., 72., 73., 74., 75., 76., 77.],
+ [78., 79., 80., 81., 82., 83., 84.],
+ [85., 86., 87., 88., 89., 90., 91.]
+ ]> : tensor<13x7xf32>
- %A_dyn = tensor.cast %A : tensor<4x4xf32> to tensor<?x?xf32>
+ %A_dyn = tensor.cast %A : tensor<13x7xf32> to tensor<?x?xf32>
- %C_init = bufferization.alloc_tensor(%c4, %c4) : tensor<?x?xf32>
+ %C_init = bufferization.alloc_tensor(%c7, %c7) : tensor<?x?xf32>
%C = linalg.fill ins(%c0 : i32) outs(%C_init : tensor<?x?xf32>) -> tensor<?x?xf32>
- // CHECK: Unranked Memref {{.*}} rank = 2 offset = 0 sizes = [4, 4] strides = [4, 1] data =
- // CHECK: [276, 304, 332, 360]
- // CHECK: [304, 336, 368, 400]
- // CHECK: [332, 368, 404, 440]
- // CHECK: [360, 400, 440, 480]
+ // CHECK: Unranked Memref {{.*}} rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data =
+ // CHECK: [32955, 33514, 34073, 34632, 35191, 35750, 36309]
+ // CHECK: [33514, 34086, 34658, 35230, 35802, 36374, 36946]
+ // CHECK: [34073, 34658, 35243, 35828, 36413, 36998, 37583]
+ // CHECK: [34632, 35230, 35828, 36426, 37024, 37622, 38220]
+ // CHECK: [35191, 35802, 36413, 37024, 37635, 38246, 38857]
+ // CHECK: [35750, 36374, 36998, 37622, 38246, 38870, 39494]
+ // CHECK: [36309, 36946, 37583, 38220, 38857, 39494, 40131]
call @matmul_transpose_a(%A_dyn, %A_dyn, %C) : (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) -> ()
return
@@ -49,22 +61,29 @@ func.func @main() {
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%module : !transform.any_op {transform.readonly}) {
- %0 = transform.structured.match ops{["linalg.matmul_transpose_a"]} in %module
+ %matmul_transpose_a = transform.structured.match ops{["linalg.matmul_transpose_a"]} in %module
: (!transform.any_op) -> !transform.any_op
- %tiled_linalg_op, %loops:3 = transform.structured.tile_using_for %0[[4], [4], 1]
+
+ // Step 1: Tile for size [4] x [4], which corresponds to SVLs x SVLs, where
+ // SVLs is the number of 32-bit elements in a vector of SVL bits.
+ %tiled_linalg_op, %loops:3 = transform.structured.tile_using_for %matmul_transpose_a[[4], [4], 1]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+
+ // Step 2: Vectorize
transform.structured.vectorize %tiled_linalg_op vector_sizes [[4], [4], 1]
: !transform.any_op
%func = transform.structured.match ops{["func.func"]} in %module
: (!transform.any_op) -> !transform.any_op
+ // Step 3: Lower vector.multi_reduction to vector.contract (+ some helpful patterns)
transform.apply_patterns to %func {
transform.apply_patterns.vector.lower_masked_transfers
transform.apply_patterns.vector.transfer_permutation_patterns
transform.apply_patterns.vector.reduction_to_contract
} : !transform.any_op
+ // Step 4: Lower vector.contract to vector.outerproduct
transform.apply_patterns to %func {
transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
transform.apply_patterns.vector.lower_masks
>From 9a55898dbe3c154372630e4d15b11ba3ceb1b3ac Mon Sep 17 00:00:00 2001
From: Cullen Rhodes <cullen.rhodes at arm.com>
Date: Wed, 8 Nov 2023 16:03:07 +0000
Subject: [PATCH 3/3] add fullstops
---
.../Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir
index dda8c93380768c9..28179fed31eca4b 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir
@@ -69,21 +69,21 @@ module attributes {transform.with_named_sequence} {
%tiled_linalg_op, %loops:3 = transform.structured.tile_using_for %matmul_transpose_a[[4], [4], 1]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
- // Step 2: Vectorize
+ // Step 2: Vectorize.
transform.structured.vectorize %tiled_linalg_op vector_sizes [[4], [4], 1]
: !transform.any_op
%func = transform.structured.match ops{["func.func"]} in %module
: (!transform.any_op) -> !transform.any_op
- // Step 3: Lower vector.multi_reduction to vector.contract (+ some helpful patterns)
+ // Step 3: Lower vector.multi_reduction to vector.contract (+ some helpful patterns).
transform.apply_patterns to %func {
transform.apply_patterns.vector.lower_masked_transfers
transform.apply_patterns.vector.transfer_permutation_patterns
transform.apply_patterns.vector.reduction_to_contract
} : !transform.any_op
- // Step 4: Lower vector.contract to vector.outerproduct
+ // Step 4: Lower vector.contract to vector.outerproduct.
transform.apply_patterns to %func {
transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
transform.apply_patterns.vector.lower_masks
More information about the Mlir-commits
mailing list