[llvm] [mlir] Fix affine for fold detached block (PR #209111)
via llvm-commits
llvm-commits at lists.llvm.org
Mon Jul 13 01:44:42 PDT 2026
https://github.com/LouisLu060211 created https://github.com/llvm/llvm-project/pull/209111
The `affine.for` folder crashes when it runs inside dialect conversion. The test pass `func.func` to `llvm.func`, which detaches the old entry block while values inside the loop still point to the old block arguments. The old block has no parent region at this point, which is a dialect conversion problem. The loop yields one of these old block arguments, and the assertion in the folder calls `isDefinedOutsideOfLoop` on it. The crash is fixed by adding a check before the assertion. So the folder bails out when the value has no parent region. The loop cannot be folded in this state anyway, so the conversion driver leaves the op alone and the pass finishes without crashing.
Assisted by: GPT 5.5
>From 4087c0ff57a953a3984e9f6d5f17684b82a944c6 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 1434c4b3e9a81abbe1c5ec3a4f35fe817e9e01c6 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 7e5e80a6cae5d1546f817f6468759045879d0e66 Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Mon, 13 Jul 2026 16:38:13 +0800
Subject: [PATCH 3/3] Fix AffineForOP creash
---
a.mlir | 17 +++++++++++++++++
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 3 +++
2 files changed, 20 insertions(+)
create mode 100644 a.mlir
diff --git a/a.mlir b/a.mlir
new file mode 100644
index 0000000000000..899e405a1dabb
--- /dev/null
+++ b/a.mlir
@@ -0,0 +1,17 @@
+module {
+ func.func private @fold_vector_transfer(%arg0: index, %arg1: memref<?xf32>, %arg2: index, %arg3: index) {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ scf.for %arg4 = %c0 to %c10 step %c1 {
+ %0 = memref.load %arg1[%arg4] : memref<?xf32>
+ gpu.barrier
+ }
+ scf.for %arg4 = %c0 to %arg3 step %c1 {
+ %0 = memref.load %arg1[%arg4] : memref<?xf32>
+ gpu.barrier
+ }
+ return
+ }
+}
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index f095500495f18..544715e6a169d 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2596,6 +2596,9 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
return {};
if (iterArgIt == iterArgs.end()) {
// `val` is defined outside of the loop.
+ if(!val.getParentRegion()){
+ return{};
+ }
assert(forOp.isDefinedOutsideOfLoop(val) &&
"must be defined outside of the loop");
hasValDefinedOutsideLoop = true;
More information about the llvm-commits
mailing list