[llvm] [mlir] Fix affine for fold detached block (PR #209111)
via llvm-commits
llvm-commits at lists.llvm.org
Mon Jul 13 01:45:31 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: LouisLu060211
<details>
<summary>Changes</summary>
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
---
Full diff: https://github.com/llvm/llvm-project/pull/209111.diff
6 Files Affected:
- (added) a.mlir (+17)
- (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+3)
- (modified) mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp (+9-4)
- (modified) mlir/test/Analysis/test-alias-analysis-modref.mlir (+14)
- (modified) mlir/test/Dialect/MemRef/mem2reg.mlir (+15)
- (modified) mlir/test/lib/Analysis/TestAliasAnalysis.cpp (+40-10)
``````````diff
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;
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/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/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
+}
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() << " <-> ";
``````````
</details>
https://github.com/llvm/llvm-project/pull/209111
More information about the llvm-commits
mailing list