[Mlir-commits] [mlir] [mlir][sme] Include hoisting in the mult-tile e2e test (PR #202979)

[Andrzej Warzyński]

https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/202979

>From 03c84fa1719ec8c19fbd63e88105ebb6face1107 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Wed, 10 Jun 2026 13:57:01 +0000
Subject: [PATCH] [mlir][sme] Update the multi-tile e2e example

These changes enable hoisting of the accumulator load/store operations
out of the K loop.
---
 .../Linalg/CPU/ArmSME/multi-tile-matmul.mlir  | 28 ++++++++-----------
 1 file changed, 12 insertions(+), 16 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 007189ad9d578..66d3ed2fd7c1b 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
@@ -69,6 +69,8 @@ module attributes {transform.with_named_sequence} {
   transform.named_sequence @__transform_main(%module : !transform.any_op {transform.consumed}) {
     %matmul = transform.structured.match ops{["linalg.matmul"]} in %module
       : (!transform.any_op) -> !transform.any_op
+    %func = transform.structured.match ops{["func.func"]} in %module
+      : (!transform.any_op) -> !transform.any_op
 
     // 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
@@ -80,28 +82,22 @@ module attributes {transform.with_named_sequence} {
     transform.structured.vectorize %tiled_linalg_op vector_sizes [[8], [8], 4]
       : !transform.any_op
 
-    // Step 3: Bufferize ahead of TransferReadDropUnitDimsPattern, which
-    // currently only supports memrefs.
-    %bufferize = transform.bufferization.one_shot_bufferize %module
-      {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
+    // Step 3: Lower vector.mask %mask { vector.transfer_* } to vector.transfer_* %mask
+    // (Unlocks hoisting below)
+    transform.apply_patterns to %loop_k {
+      transform.apply_patterns.vector.lower_masked_transfers
+    } : !transform.op<"scf.for">
 
-    %func = transform.structured.match ops{["func.func"]} in %bufferize
-      : (!transform.any_op) -> !transform.any_op
+    // Step 4: Hoist the C accumulator load/store out of the k-loop while still
+    // in tensor form, so transfer_write has a result value the loop can yield.
+    transform.apply_licm to %loop_k : !transform.op<"scf.for">
+    transform.loop.hoist_loop_invariant_subsets %loop_k : !transform.op<"scf.for">
 
-    // Step 4: Lower vector.multi_reduction to vector.contract (+ some helpful patterns).
+    // Step 5: Lower to vector.outerproduct.
     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 5: Lower vector.contract to vector.outerproduct. Also drop unit
-    // dims, specifically to prevent vector.transfer_read of vector<[8]x1xf32>,
-    // which can't be lowered in generic path.
-    transform.apply_patterns to %func {
       transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
-      transform.apply_patterns.vector.lower_masks
-      transform.apply_patterns.vector.rank_reducing_subview_patterns
     } : !transform.any_op
 
     transform.yield