[Mlir-commits] [mlir] [mlir][ArmSME] Reject non-unit-stride `tile_load/tile_store` memrefs and pass `layout{IdentityLayoutMap}` to matmul tests (PR #214959)

[Federico Bruzzone]

https://github.com/FedericoBruzzone created https://github.com/llvm/llvm-project/pull/214959

#210952 added a stride-unit verifier to `vector.maskedload/maskedstore/expandload/compressstore`, which correctly caught a pre-existing bug: `linalg.matmul` lowered through ArmSME (i.e., `multi-tile-matmul-mixed-types.mlir`,
`matmul.mlir`) produces `vector.maskedload` on `memrefs` whose type doesn't guarantee contiguity, [breaking the buildbot](https://lab.llvm.org/buildbot/#/builders/121/builds/2588).

Traced the `maskedload` back to `arm_sme.tile_load`/`tile_store`, whose docs already state: "the slice of memory must be contiguous", but op had a verifier enforcing it. **Adds one, matching the documented contract**.

With the check in place, the two integration tests fail earlier and more clearly, at `arm_sme.tile_load` instead of deep in the SCF lowering.
Essentially, `bufferize-function-boundaries` defaults to `infer-layout-map`, which always gives function *arguments* a fully
dynamic layout, even though these tests always pass contiguous data. This PR fixes this by requiring `layout{IdentityLayoutMap}` explicitly in both tests

Since on macOS M4 Pro the SME tests are disabled (`MLIR_RUN_ARM_SME_TESTS` requires Linux hwcap), I manually verified both tests by manually running `mlir-opt`.


>From c8e82010f174e953bc5992d13cc9210756a058b7 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Sat, 8 Aug 2026 12:19:57 +0200
Subject: [PATCH] [mlir][ArmSME] Reject non-unit-stride tile_load/tile_store
 memrefs

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 .../mlir/Dialect/ArmSME/IR/ArmSMEOps.td       |  2 ++
 mlir/lib/Dialect/ArmSME/IR/ArmSME.cpp         | 19 +++++++++++++++++++
 .../Dialect/Linalg/CPU/ArmSME/matmul.mlir     |  8 ++++++--
 .../ArmSME/multi-tile-matmul-mixed-types.mlir |  8 ++++++--
 4 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEOps.td b/mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEOps.td
index 2f083b55d4904..264c3969a1152 100644
--- a/mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEOps.td
+++ b/mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEOps.td
@@ -329,6 +329,7 @@ def TileLoadOp : ArmSME_Op<"tile_load", [
     CPred<"bool(getPadding()) == bool(getMask())">
   >,
 ]> {
+  let hasVerifier = 1;
   let summary = "Tile load operation";
   let description = [{
     Loads a 2D SME "virtual tile" from memory defined by a base and indices,
@@ -411,6 +412,7 @@ def TileStoreOp : ArmSME_Op<"tile_store", [
   AllElementTypesMatch<["valueToStore", "base"]>,
   HasMatchingMaskTypeConstraint<"valueToStore", "mask">,
 ]> {
+  let hasVerifier = 1;
   let summary = "Tile store operation";
   let description = [{
     Stores a 2D SME "virtual tile" to memory defined by a base and indices,
diff --git a/mlir/lib/Dialect/ArmSME/IR/ArmSME.cpp b/mlir/lib/Dialect/ArmSME/IR/ArmSME.cpp
index bc505f936d6c5..de0a868415556 100644
--- a/mlir/lib/Dialect/ArmSME/IR/ArmSME.cpp
+++ b/mlir/lib/Dialect/ArmSME/IR/ArmSME.cpp
@@ -46,6 +46,25 @@ LogicalResult verifyArmSMETileOpInterface(Operation *op) {
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/ArmSME/IR/ArmSMEAttrDefs.cpp.inc"
 
+//===----------------------------------------------------------------------===//
+// TileLoadOp / TileStoreOp
+//===----------------------------------------------------------------------===//
+
+// The slice of memory loaded/stored by tile_load/tile_store is read/written
+// one contiguous row (tile slice) at a time, along the memref's most minor
+// dimension, so that dimension must have unit stride.
+LogicalResult TileLoadOp::verify() {
+  if (!getMemRefType().isLastDimUnitStride())
+    return emitOpError("most minor dimension of memref must have unit stride");
+  return success();
+}
+
+LogicalResult TileStoreOp::verify() {
+  if (!getMemRefType().isLastDimUnitStride())
+    return emitOpError("most minor dimension of memref must have unit stride");
+  return success();
+}
+
 void ArmSMEDialect::initialize() {
   addAttributes<
 #define GET_ATTRDEF_LIST
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul.mlir
index e2c0f1d22fea1..147ab7401c5cf 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul.mlir
@@ -67,8 +67,12 @@ module attributes {transform.with_named_sequence} {
       : !transform.any_op
 
     // Step 3: Bufferize ahead of TransferReadDropUnitDimsPattern, which
-    // currently only supports memrefs.
-    %bufferize = transform.bufferization.one_shot_bufferize %module
+    // currently only supports memrefs. Force an identity (contiguous) layout
+    // map at function boundaries: the default inferred layout is fully
+    // dynamic for function arguments, which later fails vector-to-ArmSME
+    // lowering's requirement that the tile memref have unit stride on its
+    // most minor dimension.
+    %bufferize = transform.bufferization.one_shot_bufferize layout{IdentityLayoutMap} %module
       {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
 
     %func = transform.structured.match ops{["func.func"]} in %bufferize
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul-mixed-types.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul-mixed-types.mlir
index 71798a6affbbc..caa4d318f5f06 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul-mixed-types.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul-mixed-types.mlir
@@ -82,8 +82,12 @@ module attributes {transform.with_named_sequence} {
       : !transform.any_op
 
     // Step 3: Bufferize ahead of TransferReadDropUnitDimsPattern, which
-    // currently only supports memrefs.
-    %bufferize = transform.bufferization.one_shot_bufferize %module
+    // currently only supports memrefs. Force an identity (contiguous) layout
+    // map at function boundaries: the default inferred layout is fully
+    // dynamic for function arguments, which later fails vector-to-ArmSME
+    // lowering's requirement that the tile memref have unit stride on its
+    // most minor dimension.
+    %bufferize = transform.bufferization.one_shot_bufferize layout{IdentityLayoutMap} %module
       {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
 
     %func = transform.structured.match ops{["func.func"]} in %bufferize