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

[Benjamin Maxwell]

https://github.com/MacDue created https://github.com/llvm/llvm-project/pull/81160

This tests both #80148 and #80170 work together to allow unrolling the reduction dimension of a matmul.

>From 74a7a4fe75b5409a181f33942eab842a06b8871b Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Thu, 8 Feb 2024 16:52:41 +0000
Subject: [PATCH] [mlir][ArmSME][test] Unroll reduction dimension in
 multi-tile-matmul.mlir

This tests both #80148 and #80170 work together to allow unrolling the
reduction dimension of a matmul.
---
 .../Dialect/Linalg/CPU/ArmSME/multi-tile-matmul.mlir   | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul.mlir
index 327f237ba8948..d5c35068ccb32 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul.mlir
@@ -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