[Mlir-commits] [mlir] b12c64e - [mlir][vector] Fix crash on untraceable masks in getCompressedMaskOp (#207299)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jul 10 07:51:24 PDT 2026
Author: Marquis
Date: 2026-07-10T15:51:19+01:00
New Revision: b12c64e0a26e9d0a45bc013bb1740e16100a700c
URL: https://github.com/llvm/llvm-project/commit/b12c64e0a26e9d0a45bc013bb1740e16100a700c
DIFF: https://github.com/llvm/llvm-project/commit/b12c64e0a26e9d0a45bc013bb1740e16100a700c.diff
LOG: [mlir][vector] Fix crash on untraceable masks in getCompressedMaskOp (#207299)
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
Claude Fable 5 did assist with this.
Added:
Modified:
mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
index 9faaebdcf8f35..e237ae0568ae4 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
@@ -89,20 +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);
- }
+ auto extractOp = dyn_cast<vector::ExtractOp>(maskOp);
+ if (!extractOp)
+ return failure();
+ maskOp = extractOp.getSource().getDefiningOp();
+ extractOps.push_back(extractOp);
}
- if (!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 a5a6fc4acfe10..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,13 +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 and cannot be traced back to a supported
+// mask-creation op, so the store is left unconverted.
+
+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 @negative_vector_maskedstore_i8_block_arg_mask
+// CHECK: vector.maskedstore
More information about the Mlir-commits
mailing list