[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:03:03 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-vector
Author: Vaisman
<details>
<summary>Changes</summary>
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.
---
Full diff: https://github.com/llvm/llvm-project/pull/226732.diff
4 Files Affected:
- (modified) mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp (+5)
- (modified) mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp (+4)
- (modified) mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir (+37)
- (modified) mlir/test/Dialect/Vector/vector-contract-matmul-transforms.mlir (+35)
``````````diff
diff --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
index 4ed868eca25c7..6cc28ad24b6b9 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 80b2b017404e3..20a145d9f92fa 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 3af86a6dedff7..e9dee270658d8 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 ac190b4bec3cd..b8095c775494d 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>
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/226732
More information about the Mlir-commits
mailing list