[Mlir-commits] [mlir] Fix narrow type mask crash (PR #208185)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 8 02:47:40 PDT 2026
https://github.com/LouisLu060211 created https://github.com/llvm/llvm-project/pull/208185
The issue exists in `getCompressedMaskOp`, when it tries to trace the mask of a
`vector.maskedstore` back to the op
This breaks when the mask is a block argument: `getDefiningOp()` returns
null. The `isa<>` check crashes on the null pointer. In addition
if the mask comes from some other unsupported op (e.g.
`vector.broadcast`), the traversal loop never advances and hangs forever.
Fix both by bailing out with `failure()`, so the conversion just reports
a normal legalization error instead of crashing. Added a regression test
with a block-argument mask.
Assisted by: Fable-5
>From fb9fd156080e0809b97d174b1a570783b23ed5c0 Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Tue, 23 Jun 2026 18:02:12 +0800
Subject: [PATCH 1/3] [mlir][memref] Avoid mem2reg crash on element count
overflow
---
mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp | 13 +++++++++----
mlir/test/Dialect/MemRef/mem2reg.mlir | 15 +++++++++++++++
2 files changed, 24 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
index 6748e2cf71804..9a8acbb61cfab 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
@@ -67,7 +67,9 @@ SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
if (!type.hasStaticShape())
return {};
// Make sure the memref contains only a single element.
- if (type.getNumElements() != 1)
+ std::optional<int64_t> numElements =
+ ShapedType::tryGetNumElements(type.getShape());
+ if (!numElements || *numElements != 1)
return {};
return {MemorySlot{getResult(), type.getElementType()}};
@@ -290,9 +292,12 @@ struct MemRefDestructurableTypeExternalModel
getSubelementIndexMap(Type type) const {
auto memrefType = llvm::cast<MemRefType>(type);
constexpr int64_t maxMemrefSizeForDestructuring = 16;
- if (!memrefType.hasStaticShape() ||
- memrefType.getNumElements() > maxMemrefSizeForDestructuring ||
- memrefType.getNumElements() == 1)
+ if (!memrefType.hasStaticShape())
+ return {};
+ std::optional<int64_t> numElements =
+ ShapedType::tryGetNumElements(memrefType.getShape());
+ if (!numElements || *numElements > maxMemrefSizeForDestructuring ||
+ *numElements == 0)
return {};
DenseMap<Attribute, Type> destructured;
diff --git a/mlir/test/Dialect/MemRef/mem2reg.mlir b/mlir/test/Dialect/MemRef/mem2reg.mlir
index 8f937c4efe75e..9327db24ff39c 100644
--- a/mlir/test/Dialect/MemRef/mem2reg.mlir
+++ b/mlir/test/Dialect/MemRef/mem2reg.mlir
@@ -309,3 +309,18 @@ func.func @two_consecutive_merge_points(%cond1: i1, %cond2: i1) -> i32 {
// CHECK: return %[[RESULT]] : i32
return %result : i32
}
+
+
+// -----
+
+// Make sure mem2reg does not crash on an alloca whose element count overflows
+// int64. https://github.com/llvm/llvm-project/issues/204297
+// CHECK-LABEL: func.func @alloca_element_count_overflow
+func.func @alloca_element_count_overflow() {
+ return
+^bb1(%0: index):
+ // The alloca must be left untouched (not promoted).
+ // CHECK: memref.alloca() : memref<9223372036854775807x3xi32>
+ %alloca = memref.alloca() : memref<9223372036854775807x3xi32>
+ return
+}
>From 10afc402a685de09e498bc61449459cf8fe230a9 Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Tue, 30 Jun 2026 16:56:00 +0800
Subject: [PATCH 2/3] Avoid crash in test alias analysis operand printing
---
.../Analysis/test-alias-analysis-modref.mlir | 14 ++++++
mlir/test/lib/Analysis/TestAliasAnalysis.cpp | 50 +++++++++++++++----
2 files changed, 54 insertions(+), 10 deletions(-)
diff --git a/mlir/test/Analysis/test-alias-analysis-modref.mlir b/mlir/test/Analysis/test-alias-analysis-modref.mlir
index 2a9c612317d53..9c1fa17c1249a 100644
--- a/mlir/test/Analysis/test-alias-analysis-modref.mlir
+++ b/mlir/test/Analysis/test-alias-analysis-modref.mlir
@@ -112,3 +112,17 @@ func.func @conditional_all_effects(%arg: memref<2xf32>) attributes {test.ptr = "
"test.conditional_side_effect_op"() {has_effects = true, test.ptr = "conditional_side_effect_op"} : () -> i32
return {test.ptr = "return"}
}
+
+// -----
+
+// CHECK-LABEL: Testing : "block_argument_with_unit_test_ptr"
+// CHECK: func.func -> func.func.region0#0: ModRef
+module {
+ func.func @block_argument_with_unit_test_ptr(%arg0: f32 {test._ptr}) attributes {test.ptr} {
+ %0 = "test.ptr"(%arg0) {test._ptr} : (f32) -> f32
+ %1 = "test.ptr"(%arg0) {test.inclusive} : (f32) -> f32
+ %2 = "test.ptr"(%arg0) {test.exclusive} : (f32) -> f32
+ %3 = "test.ptr"(%arg0) {test.c_ptr} : (f32) -> f32
+ return
+ }
+}
\ No newline at end of file
diff --git a/mlir/test/lib/Analysis/TestAliasAnalysis.cpp b/mlir/test/lib/Analysis/TestAliasAnalysis.cpp
index 0125e403272a8..8b33a08db4ab1 100644
--- a/mlir/test/lib/Analysis/TestAliasAnalysis.cpp
+++ b/mlir/test/lib/Analysis/TestAliasAnalysis.cpp
@@ -16,33 +16,63 @@
#include "mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Pass/Pass.h"
+#include "mlir/IR/Value.h"
using namespace mlir;
-/// Print a value that is used as an operand of an alias query.
+namespace mlir {
+namespace test {
+
+/// Print an operation that is used as an operand of an alias query.
static void printAliasOperand(Operation *op) {
- llvm::errs() << op->getAttrOfType<StringAttr>("test.ptr").getValue();
+ if (!op) {
+ llvm::errs() << "<<NULL OPERATION>>";
+ return;
+ }
+
+ if (StringAttr attr = op->getAttrOfType<StringAttr>("test.ptr")) {
+ llvm::errs() << attr.getValue();
+ return;
+ }
+
+ llvm::errs() << op->getName().getStringRef();
}
+
+/// Print a value that is used as an operand of an alias query.
static void printAliasOperand(Value value) {
if (BlockArgument arg = dyn_cast<BlockArgument>(value)) {
Region *region = arg.getParentRegion();
+ Operation *parentOp = region ? region->getParentOp() : nullptr;
+
+ if (parentOp) {
+ if (StringAttr attr = parentOp->getAttrOfType<StringAttr>("test.ptr"))
+ llvm::errs() << attr.getValue();
+ else
+ llvm::errs() << parentOp->getName().getStringRef();
+ } else {
+ llvm::errs() << "<<UNKNOWN PARENT>>";
+ }
+
+ if (region)
+ llvm::errs() << ".region" << region->getRegionNumber();
+
unsigned parentBlockNumber = arg.getOwner()->computeBlockNumber();
- llvm::errs() << region->getParentOp()
- ->getAttrOfType<StringAttr>("test.ptr")
- .getValue()
- << ".region" << region->getRegionNumber();
if (parentBlockNumber != 0)
llvm::errs() << ".block" << parentBlockNumber;
+
llvm::errs() << "#" << arg.getArgNumber();
return;
}
- OpResult result = cast<OpResult>(value);
- printAliasOperand(result.getOwner());
+
+ Operation *owner = value.getDefiningOp();
+ assert(owner && "expected non-block argument value to have a defining op");
+
+ printAliasOperand(owner);
+
+ auto result = cast<mlir::OpResult>(value);
llvm::errs() << "#" << result.getResultNumber();
}
-namespace mlir {
-namespace test {
void printAliasResult(AliasResult result, Value lhs, Value rhs) {
printAliasOperand(lhs);
llvm::errs() << " <-> ";
>From 1b2ed96b6c291c781f7d8953c62e847edbd22f89 Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Wed, 8 Jul 2026 17:15:19 +0800
Subject: [PATCH 3/3] [mlir][vector] Fix crash in narrow-type emulation when
mask has no defining op
---
.../Transforms/VectorEmulateNarrowType.cpp | 5 ++++-
.../emulate-narrow-type-unsupported.mlir | 18 ++++++++++++++++++
2 files changed, 22 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
index 9faaebdcf8f35..9aa28ff899bc4 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
@@ -98,10 +98,13 @@ static FailureOr<Operation *> getCompressedMaskOp(OpBuilder &rewriter,
if (auto extractOp = dyn_cast<vector::ExtractOp>(maskOp)) {
maskOp = extractOp.getSource().getDefiningOp();
extractOps.push_back(extractOp);
+ }else {
+ return failure();
}
}
- 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..17bd24e2c1959 100644
--- a/mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir
+++ b/mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir
@@ -109,3 +109,21 @@ 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
+
+// -----
+
+///----------------------------------------------------------------------------------------
+/// vector.maskedload
+///----------------------------------------------------------------------------------------
+
+func.func @vector_maskedload_blockarg_mask_negative(%arg0: memref<16xi8>, %arg1: vector<8xi1>, %arg2: vector<8xi8>) -> vector<8xi8> {
+ %c0 = arith.constant 0 : index
+ %0 = vector.maskedload %arg0[%c0], %arg1, %arg2 : memref<16xi8>, vector<8xi1>, vector<8xi8> into vector<8xi8>
+ return %0 : vector<8xi8>
+}
+// The mask is a block argument with no defining op, so it cannot be
+// compressed - expect no conversion (and, crucially, no crash).
+// CHECK-LABEL: func @vector_maskedload_blockarg_mask_negative
+// CHECK-SAME: %{{.*}}: memref<16xi8>, %[[MASK:.+]]: vector<8xi1>
+// CHECK: vector.maskedload %{{.*}}, %[[MASK]]
+// CHECK-NOT: i32
More information about the Mlir-commits
mailing list