[Mlir-commits] [mlir] [mlir][vector] Fix crash on untraceable masks in getCompressedMaskOp (PR #207299)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 8 04:21:12 PDT 2026
https://github.com/marquisburg updated https://github.com/llvm/llvm-project/pull/207299
>From b1737aa3922d37bd33bbe450da525b46111a5bf4 Mon Sep 17 00:00:00 2001
From: GhostDragonWasTook <marquis at marquiseco.co.za>
Date: Fri, 3 Jul 2026 02:40:25 +0200
Subject: [PATCH 1/2] [mlir][vector] Fix crash on untraceable masks in
getCompressedMaskOp
The loop tracing a mask back to its creation op assumed the chain
always ends at a create_mask/constant_mask/arith.constant, a
block-argument mask made it fall through to isa<> on a null op
(segfault), and any other unhandled defining op spun it forever.
Fixes #206928
---
.../Transforms/VectorEmulateNarrowType.cpp | 8 +++++-
.../emulate-narrow-type-unsupported.mlir | 28 +++++++++++++++++++
2 files changed, 35 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
index 9faaebdcf8f35..e37296e50fb29 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
@@ -98,10 +98,16 @@ static FailureOr<Operation *> getCompressedMaskOp(OpBuilder &rewriter,
if (auto extractOp = dyn_cast<vector::ExtractOp>(maskOp)) {
maskOp = extractOp.getSource().getDefiningOp();
extractOps.push_back(extractOp);
+ } else {
+ // Unsupported mask-defining op (e.g. a block argument, which has no
+ // defining op, or an op we cannot trace through). Bail out rather than
+ // looping forever or dereferencing a null op below.
+ break;
}
}
- if (!isa<arith::ConstantOp, vector::CreateMaskOp, vector::ConstantMaskOp>(
+ if (!maskOp ||
+ !isa<arith::ConstantOp, vector::CreateMaskOp, vector::ConstantMaskOp>(
maskOp))
return failure();
diff --git a/mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir b/mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir
index a5a6fc4acfe10..74f72ed5a5d9c 100644
--- a/mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir
+++ b/mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir
@@ -109,3 +109,31 @@ func.func @vector_maskedstore_2d_i8_negative(%arg0: index, %arg1: index, %arg2:
// CHECK-LABEL: func @vector_maskedstore_2d_i8_negative
// CHECK: memref.alloc() : memref<3x8xi8>
// CHECK-NOT: i32
+
+// -----
+
+// The mask is a block argument, so its defining op cannot be traced back to a
+// supported mask-creation op. This must gracefully bail out instead of crashing
+// (see llvm/llvm-project#206928).
+
+func.func @vector_maskedstore_i8_block_arg_mask_negative(%arg0: memref<?xi8>, %mask: vector<8xi1>, %value: vector<8xi8>) {
+ %c0 = arith.constant 0 : index
+ vector.maskedstore %arg0[%c0], %mask, %value : memref<?xi8>, vector<8xi1>, vector<8xi8>
+ return
+}
+
+// CHECK-LABEL: func @vector_maskedstore_i8_block_arg_mask_negative
+// CHECK: vector.maskedstore
+
+// -----
+
+// As above, but for `vector.maskedload`.
+
+func.func @vector_maskedload_i8_block_arg_mask_negative(%arg0: memref<?xi8>, %mask: vector<8xi1>, %passthru: vector<8xi8>) -> vector<8xi8> {
+ %c0 = arith.constant 0 : index
+ %0 = vector.maskedload %arg0[%c0], %mask, %passthru : memref<?xi8>, vector<8xi1>, vector<8xi8> into vector<8xi8>
+ return %0 : vector<8xi8>
+}
+
+// CHECK-LABEL: func @vector_maskedload_i8_block_arg_mask_negative
+// CHECK: vector.maskedload
>From 561f1ab1a0472ae1416b66b56a180113653473f9 Mon Sep 17 00:00:00 2001
From: GhostDragonWasTook <marquis at marquiseco.co.za>
Date: Wed, 8 Jul 2026 13:20:02 +0200
Subject: [PATCH 2/2] [mlir][vector][NFC] Tidy getCompressedMaskOp bail-out and
negative tests
Replace the break-then-recheck bail-out in the mask trace loop with an
early return when a defining op is not a vector.extract, and use
isa_and_present for the post-loop check so the block-argument (null
maskOp) case is handled without a separate guard. NFC.
Rename the negative tests in emulate-narrow-type-unsupported.mlir to use
the negative_ prefix per the MLIR testing guide, and move the
block-argument maskedload test into the vector.maskedload section.
---
.../Transforms/VectorEmulateNarrowType.cpp | 28 ++++-----
.../emulate-narrow-type-unsupported.mlir | 60 +++++++++----------
2 files changed, 44 insertions(+), 44 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
index e37296e50fb29..e237ae0568ae4 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
@@ -89,26 +89,26 @@ static FailureOr<Operation *> getCompressedMaskOp(OpBuilder &rewriter,
numSrcElemsPerDest;
Operation *maskOp = mask.getDefiningOp();
- SmallVector<vector::ExtractOp, 2> extractOps;
+ // Chain of `vector.extract` ops that lead to the op that created the mask.
// TODO: add support to `vector.broadcast`.
- // Finding the mask creation operation.
+ SmallVector<vector::ExtractOp, 2> extractOps;
+
+ // Trace the mask back to its creation op, looking through `vector.extract`
+ // ops. Any other defining op is unsupported.
while (maskOp &&
!isa<arith::ConstantOp, vector::CreateMaskOp, vector::ConstantMaskOp>(
maskOp)) {
- if (auto extractOp = dyn_cast<vector::ExtractOp>(maskOp)) {
- maskOp = extractOp.getSource().getDefiningOp();
- extractOps.push_back(extractOp);
- } else {
- // Unsupported mask-defining op (e.g. a block argument, which has no
- // defining op, or an op we cannot trace through). Bail out rather than
- // looping forever or dereferencing a null op below.
- break;
- }
+ auto extractOp = dyn_cast<vector::ExtractOp>(maskOp);
+ if (!extractOp)
+ return failure();
+ maskOp = extractOp.getSource().getDefiningOp();
+ extractOps.push_back(extractOp);
}
- if (!maskOp ||
- !isa<arith::ConstantOp, vector::CreateMaskOp, vector::ConstantMaskOp>(
- maskOp))
+ // `maskOp` is null when the mask is a block argument, which has no defining
+ // op.
+ if (!isa_and_present<arith::ConstantOp, vector::CreateMaskOp,
+ vector::ConstantMaskOp>(maskOp))
return failure();
// Computing the "compressed" mask. All the emulation logic (i.e. computing
diff --git a/mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir b/mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir
index 74f72ed5a5d9c..d3abccf641914 100644
--- a/mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir
+++ b/mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir
@@ -7,14 +7,14 @@
/// vector.load
///----------------------------------------------------------------------------------------
-func.func @vector_load_2d_i8_negative(%arg1: index, %arg2: index) -> vector<2x4xi8> {
+func.func @negative_vector_load_2d_i8(%arg1: index, %arg2: index) -> vector<2x4xi8> {
%0 = memref.alloc() : memref<3x4xi8>
%1 = vector.load %0[%arg1, %arg2] : memref<3x4xi8>, vector<2x4xi8>
return %1 : vector<2x4xi8>
}
// No support for loading 2D vectors - expect no conversions
-// CHECK-LABEL: func @vector_load_2d_i8_negative
+// CHECK-LABEL: func @negative_vector_load_2d_i8
// CHECK: memref.alloc() : memref<3x4xi8>
// CHECK-NOT: i32
@@ -24,14 +24,14 @@ func.func @vector_load_2d_i8_negative(%arg1: index, %arg2: index) -> vector<2x4x
/// vector.transfer_read
///----------------------------------------------------------------------------------------
-func.func @vector_transfer_read_2d_i4_negative(%arg1: index, %arg2: index) -> vector<2x8xi4> {
+func.func @negative_vector_transfer_read_2d_i4(%arg1: index, %arg2: index) -> vector<2x8xi4> {
%c0 = arith.constant 0 : i4
%0 = memref.alloc() : memref<3x8xi4>
%1 = vector.transfer_read %0[%arg1, %arg2], %c0 {in_bounds = [true, true]} :
memref<3x8xi4>, vector<2x8xi4>
return %1 : vector<2x8xi4>
}
-// CHECK-LABEL: func @vector_transfer_read_2d_i4_negative
+// CHECK-LABEL: func @negative_vector_transfer_read_2d_i4
// CHECK: memref.alloc() : memref<3x8xi4>
// CHECK-NOT: i32
@@ -41,7 +41,7 @@ func.func @vector_transfer_read_2d_i4_negative(%arg1: index, %arg2: index) -> ve
/// vector.maskedload
///----------------------------------------------------------------------------------------
-func.func @vector_maskedload_2d_i8_negative(%arg1: index, %arg2: index, %arg3: index, %passthru: vector<2x4xi8>) -> vector<2x4xi8> {
+func.func @negative_vector_maskedload_2d_i8(%arg1: index, %arg2: index, %arg3: index, %passthru: vector<2x4xi8>) -> vector<2x4xi8> {
%0 = memref.alloc() : memref<3x4xi8>
%mask = vector.create_mask %arg3, %arg3 : vector<2x4xi1>
%1 = vector.maskedload %0[%arg1, %arg2], %mask, %passthru :
@@ -49,17 +49,31 @@ func.func @vector_maskedload_2d_i8_negative(%arg1: index, %arg2: index, %arg3: i
return %1 : vector<2x4xi8>
}
-// CHECK-LABEL: func @vector_maskedload_2d_i8_negative
+// CHECK-LABEL: func @negative_vector_maskedload_2d_i8
// CHECK: memref.alloc() : memref<3x4xi8>
// CHECK-NOT: i32
// -----
+// The mask is a block argument and cannot be traced back to a supported
+// mask-creation op, so the load is left unconverted.
+
+func.func @negative_vector_maskedload_i8_block_arg_mask(%arg0: memref<?xi8>, %mask: vector<8xi1>, %passthru: vector<8xi8>) -> vector<8xi8> {
+ %c0 = arith.constant 0 : index
+ %0 = vector.maskedload %arg0[%c0], %mask, %passthru : memref<?xi8>, vector<8xi1>, vector<8xi8> into vector<8xi8>
+ return %0 : vector<8xi8>
+}
+
+// CHECK-LABEL: func @negative_vector_maskedload_i8_block_arg_mask
+// CHECK: vector.maskedload
+
+// -----
+
///----------------------------------------------------------------------------------------
/// vector.extract -> vector.masked_load
///----------------------------------------------------------------------------------------
-func.func @vector_extract_maskedload_2d_i4_negative(%arg1: index) -> vector<8x8x16xi4> {
+func.func @negative_vector_extract_maskedload_2d_i4(%arg1: index) -> vector<8x8x16xi4> {
%0 = memref.alloc() : memref<8x8x16xi4>
%c0 = arith.constant 0 : index
%c16 = arith.constant 16 : index
@@ -73,7 +87,7 @@ func.func @vector_extract_maskedload_2d_i4_negative(%arg1: index) -> vector<8x8x
return %63 : vector<8x8x16xi4>
}
-// CHECK-LABEL: func @vector_extract_maskedload_2d_i4_negative
+// CHECK-LABEL: func @negative_vector_extract_maskedload_2d_i4
// CHECK: memref.alloc() : memref<8x8x16xi4>
// CHECK-NOT: i32
@@ -83,13 +97,13 @@ func.func @vector_extract_maskedload_2d_i4_negative(%arg1: index) -> vector<8x8x
/// vector.store
///----------------------------------------------------------------------------------------
-func.func @vector_store_2d_i8_negative(%arg0: vector<2x8xi8>, %arg1: index, %arg2: index) {
+func.func @negative_vector_store_2d_i8(%arg0: vector<2x8xi8>, %arg1: index, %arg2: index) {
%0 = memref.alloc() : memref<4x8xi8>
vector.store %arg0, %0[%arg1, %arg2] :memref<4x8xi8>, vector<2x8xi8>
return
}
-// CHECK-LABEL: func @vector_store_2d_i8_negative
+// CHECK-LABEL: func @negative_vector_store_2d_i8
// CHECK: memref.alloc() : memref<4x8xi8>
// CHECK-NOT: i32
@@ -99,41 +113,27 @@ func.func @vector_store_2d_i8_negative(%arg0: vector<2x8xi8>, %arg1: index, %arg
/// vector.maskedstore
///----------------------------------------------------------------------------------------
-func.func @vector_maskedstore_2d_i8_negative(%arg0: index, %arg1: index, %arg2: index, %value: vector<2x8xi8>) {
+func.func @negative_vector_maskedstore_2d_i8(%arg0: index, %arg1: index, %arg2: index, %value: vector<2x8xi8>) {
%0 = memref.alloc() : memref<3x8xi8>
%mask = vector.create_mask %arg2, %arg2 : vector<2x8xi1>
vector.maskedstore %0[%arg0, %arg1], %mask, %value : memref<3x8xi8>, vector<2x8xi1>, vector<2x8xi8>
return
}
-// CHECK-LABEL: func @vector_maskedstore_2d_i8_negative
+// CHECK-LABEL: func @negative_vector_maskedstore_2d_i8
// CHECK: memref.alloc() : memref<3x8xi8>
// CHECK-NOT: i32
// -----
-// The mask is a block argument, so its defining op cannot be traced back to a
-// supported mask-creation op. This must gracefully bail out instead of crashing
-// (see llvm/llvm-project#206928).
+// The mask is a block argument and cannot be traced back to a supported
+// mask-creation op, so the store is left unconverted.
-func.func @vector_maskedstore_i8_block_arg_mask_negative(%arg0: memref<?xi8>, %mask: vector<8xi1>, %value: vector<8xi8>) {
+func.func @negative_vector_maskedstore_i8_block_arg_mask(%arg0: memref<?xi8>, %mask: vector<8xi1>, %value: vector<8xi8>) {
%c0 = arith.constant 0 : index
vector.maskedstore %arg0[%c0], %mask, %value : memref<?xi8>, vector<8xi1>, vector<8xi8>
return
}
-// CHECK-LABEL: func @vector_maskedstore_i8_block_arg_mask_negative
+// CHECK-LABEL: func @negative_vector_maskedstore_i8_block_arg_mask
// CHECK: vector.maskedstore
-
-// -----
-
-// As above, but for `vector.maskedload`.
-
-func.func @vector_maskedload_i8_block_arg_mask_negative(%arg0: memref<?xi8>, %mask: vector<8xi1>, %passthru: vector<8xi8>) -> vector<8xi8> {
- %c0 = arith.constant 0 : index
- %0 = vector.maskedload %arg0[%c0], %mask, %passthru : memref<?xi8>, vector<8xi1>, vector<8xi8> into vector<8xi8>
- return %0 : vector<8xi8>
-}
-
-// CHECK-LABEL: func @vector_maskedload_i8_block_arg_mask_negative
-// CHECK: vector.maskedload
More information about the Mlir-commits
mailing list