[Mlir-commits] [mlir] [mlir][vector-to-gpu] Do not prepare masked contractions for MMA (PR #226732)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 26 08:11:21 PDT 2026
https://github.com/Vaisman updated https://github.com/llvm/llvm-project/pull/226732
>From bc0fe0847ec9ea19048c23790e4fe2b9c1094f51 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 fe6810e01addc2..4533adeda4100b 100644
--- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
+++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
@@ -439,6 +439,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 0d32a7600cf8ed..1dd9a11f8f7bcf 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 6ab27f0db89fd1..92dd080c924be3 100644
--- a/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir
+++ b/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir
@@ -752,3 +752,40 @@ func.func @matmul_masked_write(%arg0: memref<16x16xf16>, %arg1: memref<16x16xf16
vector.mask %mask { vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<16x16xf16> } : vector<16x16xi1>
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