[Mlir-commits] [mlir] f6ddf39 - [mlir][ArmNeon] Reject masked contractions in `I8MM/BFMMLA` lowering patterns (#213698)

[llvmlistbot at llvm.org]

Author: Federico Bruzzone
Date: 2026-08-06T12:26:41+02:00
New Revision: f6ddf3965910afb16285418fe7bbb9ca36a8b18e

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

LOG: [mlir][ArmNeon] Reject masked contractions in `I8MM/BFMMLA` lowering patterns (#213698)

`LowerContractionToNeonI8MMPattern` and
`LowerContractionToNeonBFMMLAPattern` rewrite a `vector.contract` into
several ops (multiple `arm_neon.intr.smmla`/`bfmmla` calls plus
`extract/insert/shape_cast` scaffolding).

When the matched `vector.contract` sits inside a `vector.mask` (e.g.,
from vectorizing a matmul with dynamically-shaped operands), this
in-place multi-op rewrite violates `vector.mask`'s invariant that its
region contain exactly **one** operation, and the verifier rejects the
result: _error: 'vector.mask' op expects only one operation to mask_.

Neither pattern checked for masking, unlike the generic
contraction-lowering
patterns. To fix the problem, we bail out via `notifyMatchFailure` when
the contraction is masked, so it falls through to a different lowering
path instead of producing invalid IR.

Found while adding neon `linalg.matmul` e2e tests in #212809.
AI was used to investigate the problem.

---------

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>

Added: 
    

Modified: 
    mlir/lib/Dialect/ArmNeon/Transforms/LowerContractToNeonPatterns.cpp
    mlir/test/Dialect/ArmNeon/lower-to-arm-neon.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/ArmNeon/Transforms/LowerContractToNeonPatterns.cpp b/mlir/lib/Dialect/ArmNeon/Transforms/LowerContractToNeonPatterns.cpp
index fd0a774dcf602..6b4285026ea04 100644
--- a/mlir/lib/Dialect/ArmNeon/Transforms/LowerContractToNeonPatterns.cpp
+++ b/mlir/lib/Dialect/ArmNeon/Transforms/LowerContractToNeonPatterns.cpp
@@ -457,6 +457,9 @@ class LowerContractionToNeonI8MMPattern
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(vector::ContractionOp op,
                                 PatternRewriter &rewriter) const override {
+    if (cast<vector::MaskableOpInterface>(op.getOperation()).isMasked())
+      return rewriter.notifyMatchFailure(
+          op, "masked contractions are not supported");
 
     VectorContractRewriterI8MM vcr;
     if (failed(vcr.matchAndInit(op, rewriter)))
@@ -473,6 +476,9 @@ class LowerContractionToNeonBFMMLAPattern
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(vector::ContractionOp op,
                                 PatternRewriter &rewriter) const override {
+    if (cast<vector::MaskableOpInterface>(op.getOperation()).isMasked())
+      return rewriter.notifyMatchFailure(
+          op, "masked contractions are not supported");
 
     VectorContractRewriterBFMMLA vcr;
     if (failed(vcr.matchAndInit(op, rewriter)))

diff  --git a/mlir/test/Dialect/ArmNeon/lower-to-arm-neon.mlir b/mlir/test/Dialect/ArmNeon/lower-to-arm-neon.mlir
index 5fc29c6442602..a8799028217df 100644
--- a/mlir/test/Dialect/ArmNeon/lower-to-arm-neon.mlir
+++ b/mlir/test/Dialect/ArmNeon/lower-to-arm-neon.mlir
@@ -515,6 +515,27 @@ func.func @vector_arm_neon_mnk_unroll(%lhs: vector<4x16xi8>, %rhs: vector<4x16xi
   return %res : vector<4x4xi32>
 }
 
+// -----
+
+// Masked contractions must be left alone: rewriting one in place would leave
+// more than one op inside the `vector.mask` region, which is invalid.
+
+// CHECK-LABEL: func @vector_arm_neon_masked_contract_not_lowered
+// CHECK:       %[[RES:.*]] = vector.mask %{{.*}} { vector.contract
+// CHECK-NOT:   arm_neon.intr.smmla
+// CHECK:       return %[[RES]]
+func.func @vector_arm_neon_masked_contract_not_lowered(
+    %lhs: vector<4x8xi8>, %rhs: vector<4x8xi8>, %acc: vector<4x4xi32>,
+    %m: index, %n: index, %k: index) -> vector<4x4xi32> {
+  %mask = vector.create_mask %m, %n, %k : vector<4x4x8xi1>
+  %lhs_extsi = arith.extsi %lhs : vector<4x8xi8> to vector<4x8xi32>
+  %rhs_extsi = arith.extsi %rhs : vector<4x8xi8> to vector<4x8xi32>
+  %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) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs_extsi, %rhs_extsi, %acc : vector<4x8xi32>, vector<4x8xi32> into vector<4x4xi32>
+  } : vector<4x4x8xi1> -> vector<4x4xi32>
+  return %res : vector<4x4xi32>
+}
+
 module attributes {transform.with_named_sequence} {
   transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {
     %func = transform.structured.match ops{["func.func"]} in %module : (!transform.any_op) -> !transform.op<"func.func">