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

[llvmlistbot at llvm.org]

Author: Federico Bruzzone
Date: 2026-08-13T07:58:47+02:00
New Revision: 2a4cb1f1f80b49abe5b87aeb0155cee2e6bd3730

URL: https://github.com/llvm/llvm-project/commit/2a4cb1f1f80b49abe5b87aeb0155cee2e6bd3730
DIFF: https://github.com/llvm/llvm-project/commit/2a4cb1f1f80b49abe5b87aeb0155cee2e6bd3730.diff

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEOps.td
    mlir/lib/Dialect/ArmSME/IR/ArmSME.cpp
    mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul.mlir
    mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul-mixed-types.mlir

Removed: 
    


################################################################################
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