[Mlir-commits] [mlir] [mlir][ArmSME] Reject non-unit-stride `tile_load/tile_store` memrefs and pass `layout{IdentityLayoutMap}` to matmul tests (PR #214959)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Aug 8 04:55:08 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-sme
@llvm/pr-subscribers-mlir-linalg
Author: Federico Bruzzone (FedericoBruzzone)
<details>
<summary>Changes</summary>
#<!-- -->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`.
---
Full diff: https://github.com/llvm/llvm-project/pull/214959.diff
4 Files Affected:
- (modified) mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEOps.td (+2)
- (modified) mlir/lib/Dialect/ArmSME/IR/ArmSME.cpp (+19)
- (modified) mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul.mlir (+6-2)
- (modified) mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul-mixed-types.mlir (+6-2)
``````````diff
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
``````````
</details>
https://github.com/llvm/llvm-project/pull/214959
More information about the Mlir-commits
mailing list