shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

[mlir][ArmSME][test] Unroll reduction dimension in multi-tile-matmul.mlir (#81160)

Browse Source 
This tests both #80148 and #80170 work together to allow unrolling the
reduction dimension of a matmul.
This commit is contained in:
Benjamin Maxwell 2024-02-12 09:40:58 +00:00 committed by GitHub
parent f7f947e620
commit 8a2a65ffdf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul.mlir
View File

@ -73,14 +73,14 @@ module attributes {transform.with_named_sequence} {
%matmul = transform.structured.match ops{["linalg.matmul"]} in %module
: (!transform.any_op) -> !transform.any_op
// Step 1: Tile for size [8] x [8], which corresponds to (2 x SVLs) x (2 x SVLs),
// where SVLs is the number of 32-bit elements in a vector of SVL bits.
// This uses all four 32-bit SME virtual tiles.
%tiled_linalg_op, %loop_i, %loop_j, %loop_k = transform.structured.tile_using_for %matmul[[8], [8], 1]
// Step 1: Tile for size [8] x [8] (unrolled by 4), which corresponds to
// (2 x SVLs) x (2 x SVLs), where SVLs is the number of 32-bit elements in a
// vector of SVL bits. This uses all four 32-bit SME virtual tiles.
%tiled_linalg_op, %loop_i, %loop_j, %loop_k = transform.structured.tile_using_for %matmul[[8], [8], 4]
: (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">, !transform.op<"scf.for">, !transform.op<"scf.for">)
// Step 2: Vectorize.
transform.structured.vectorize %tiled_linalg_op vector_sizes [[8], [8], 1]
transform.structured.vectorize %tiled_linalg_op vector_sizes [[8], [8], 4]
: !transform.any_op
// Step 3: Bufferize ahead of TransferReadDropUnitDimsPattern, which