[Mlir-commits] [mlir] [mlir][vector-to-gpu] Do not prepare masked contractions for MMA (PR #226732)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Sep 26 14:02:24 PDT 2026
https://github.com/Vaisman created https://github.com/llvm/llvm-project/pull/226732
PrepareContractToGPUMMA and CanonicalizeContractMatmulToMMT ignore an enclosing vector.mask.
When they insert transposes, those end up inside the mask region and the IR no longer verifies.
When they only swap the operands, m and n trade places under the old mask. For square shapes this still verifies, but the bounds apply to the wrong dimensions.
Bail out on masked contractions in both patterns.
CanonicalizeContractMatmulToMMT lives in the Vector dialect and is also used outside VectorToGPU.
>From 848cc47c636e63a827cf6cd1cb2811183837f0c4 Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Fri, 25 Sep 2026 19:51:47 +0200
Subject: [PATCH] [mlir][vector-to-gpu] Do not prepare masked contractions for
MMA
PrepareContractToGPUMMA and CanonicalizeContractMatmulToMMT ignore an
enclosing vector.mask. When they insert transposes, those end up inside
the mask region and the IR no longer verifies. When they only swap the
operands, m and n trade places under the old mask. For square shapes
this still verifies, but the bounds apply to the wrong dimensions.
Bail out on masked contractions in both patterns.
CanonicalizeContractMatmulToMMT lives in the Vector dialect and is also
used outside VectorToGPU.
---
.../Conversion/VectorToGPU/VectorToGPU.cpp | 5 +++
.../Vector/Transforms/VectorTransforms.cpp | 4 ++
.../VectorToGPU/vector-to-mma-ops.mlir | 37 +++++++++++++++++++
.../vector-contract-matmul-transforms.mlir | 35 ++++++++++++++++++
4 files changed, 81 insertions(+)
diff --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
index 4ed868eca25c7b..6cc28ad24b6b95 100644
--- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
+++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
@@ -418,6 +418,11 @@ struct PrepareContractToGPUMMA
LogicalResult matchAndRewrite(vector::ContractionOp op,
PatternRewriter &rewriter) const override {
+ // The mask is not updated: the rewrite may create ops inside the
+ // vector.mask region or swap the m and n iteration dimensions.
+ if (op.isMasked())
+ return rewriter.notifyMatchFailure(op, "masked contraction");
+
Location loc = op.getLoc();
Value lhs = op.getLhs(), rhs = op.getRhs(), res = op.getAcc();
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 80b2b017404e3c..20a145d9f92fad 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1797,6 +1797,10 @@ struct CanonicalizeContractMatmulToMMT final
PatternRewriter &rewriter) const override {
if (failed(filter(op)))
return failure();
+ // The mask is not updated: the rewrite may create ops inside the
+ // vector.mask region or swap the m and n iteration dimensions.
+ if (op.isMasked())
+ return rewriter.notifyMatchFailure(op, "masked contraction");
Location loc = op.getLoc();
Value lhs = op.getLhs();
diff --git a/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir b/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir
index 3af86a6dedff7d..e9dee270658d81 100644
--- a/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir
+++ b/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir
@@ -727,3 +727,40 @@ func.func @read_transpose_with_broadcast_3d(%arg0: memref<2x2x2xf16>, %arg1: mem
vector.transfer_write %B, %arg1[%c0, %c0] {in_bounds = [true, true]} : vector<2x2xf16>, memref<2x2xf16>
return
}
+
+// -----
+
+// A masked contraction is not prepared for MMA conversion. With transposes
+// the rewrite would create ops inside the vector.mask region; with only an
+// operand swap it would swap the m and n dimensions under an unchanged mask.
+
+#map1 = affine_map<(d0, d1, d2) -> (d2, d0)>
+#map2 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: func @masked_matmul_km_kn_mn
+// CHECK-NOT: vector.transpose
+// CHECK: vector.mask %{{.*}} { vector.contract
+// CHECK-NOT: vector.transpose
+// CHECK: return
+func.func @masked_matmul_km_kn_mn(%arg0: vector<16x16xf16>, %arg1: vector<16x16xf16>, %arg2: vector<16x16xf16>, %m: index, %n: index, %k: index) -> vector<16x16xf16> {
+ %mask = vector.create_mask %m, %n, %k : vector<16x16x16xi1>
+ %D = vector.mask %mask { vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16> } : vector<16x16x16xi1> -> vector<16x16xf16>
+ return %D : vector<16x16xf16>
+}
+
+// -----
+
+#map1 = affine_map<(d0, d1, d2) -> (d2, d0)>
+#map4 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map5 = affine_map<(d0, d1, d2) -> (d1, d0)>
+
+// CHECK-LABEL: func @masked_matmul_km_nk_nm
+// CHECK-SAME: (%[[A:.+]]: vector<16x16xf16>, %[[B:.+]]: vector<16x16xf16>, %[[C:.+]]: vector<16x16xf16>
+// CHECK: vector.mask %{{.*}} { vector.contract {{.+}} %[[A]], %[[B]], %[[C]]
+// CHECK: return
+func.func @masked_matmul_km_nk_nm(%arg0: vector<16x16xf16>, %arg1: vector<16x16xf16>, %arg2: vector<16x16xf16>, %m: index, %n: index, %k: index) -> vector<16x16xf16> {
+ %mask = vector.create_mask %m, %n, %k : vector<16x16x16xi1>
+ %D = vector.mask %mask { vector.contract {indexing_maps = [#map1, #map4, #map5], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16> } : vector<16x16x16xi1> -> vector<16x16xf16>
+ return %D : vector<16x16xf16>
+}
diff --git a/mlir/test/Dialect/Vector/vector-contract-matmul-transforms.mlir b/mlir/test/Dialect/Vector/vector-contract-matmul-transforms.mlir
index ac190b4bec3cdb..b8095c775494d9 100644
--- a/mlir/test/Dialect/Vector/vector-contract-matmul-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-contract-matmul-transforms.mlir
@@ -196,3 +196,38 @@ func.func @matmul_km_nk_nm_4x4xi32(%arg0: vector<4x4xi32>, %arg1: vector<4x4xi32
kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<4x4xi32>, vector<4x4xi32> into vector<4x4xi32>
return %res : vector<4x4xi32>
}
+
+// Masked contractions are not rewritten: the mask would not be updated.
+// CHECK-LABEL: func.func @masked_matmul_mk_kn_mn_4x4xi32
+// CHECK-SAME: ([[ARG0:%.+]]: vector<4x4xi32>, [[ARG1:%.+]]: vector<4x4xi32>, [[ARG2:%.+]]: vector<4x4xi32>
+// CHECK-NEXT: [[MASK:%.+]] = vector.create_mask
+// CHECK-NEXT: [[RES:%.+]] = vector.mask [[MASK]] { vector.contract {{.+}} [[ARG0]], [[ARG1]], [[ARG2]]
+// CHECK-NEXT: return [[RES]]
+func.func @masked_matmul_mk_kn_mn_4x4xi32(%arg0: vector<4x4xi32>, %arg1: vector<4x4xi32>, %arg2: vector<4x4xi32>, %m: index, %n: index, %k: index) -> vector<4x4xi32> {
+ %mask = vector.create_mask %m, %n, %k : vector<4x4x4xi1>
+ %res = vector.mask %mask {
+ vector.contract {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
+ affine_map<(d0, d1, d2) -> (d2, d1)>,
+ affine_map<(d0, d1, d2) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<4x4xi32>, vector<4x4xi32> into vector<4x4xi32>
+ } : vector<4x4x4xi1> -> vector<4x4xi32>
+ return %res : vector<4x4xi32>
+}
+
+// CHECK-LABEL: func.func @masked_matmul_mk_nk_nm_4x4xi32
+// CHECK-SAME: ([[ARG0:%.+]]: vector<4x4xi32>, [[ARG1:%.+]]: vector<4x4xi32>, [[ARG2:%.+]]: vector<4x4xi32>
+// CHECK-NEXT: [[MASK:%.+]] = vector.create_mask
+// CHECK-NEXT: [[RES:%.+]] = vector.mask [[MASK]] { vector.contract {{.+}} [[ARG0]], [[ARG1]], [[ARG2]]
+// CHECK-NEXT: return [[RES]]
+func.func @masked_matmul_mk_nk_nm_4x4xi32(%arg0: vector<4x4xi32>, %arg1: vector<4x4xi32>, %arg2: vector<4x4xi32>, %m: index, %n: index, %k: index) -> vector<4x4xi32> {
+ %mask = vector.create_mask %m, %n, %k : vector<4x4x4xi1>
+ %res = vector.mask %mask {
+ vector.contract {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
+ affine_map<(d0, d1, d2) -> (d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d1, d0)>],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<4x4xi32>, vector<4x4xi32> into vector<4x4xi32>
+ } : vector<4x4x4xi1> -> vector<4x4xi32>
+ return %res : vector<4x4xi32>
+}
More information about the Mlir-commits
mailing list