[llvm] [mlir] Fix affine for fold detached block (PR #209111)
via llvm-commits
llvm-commits at lists.llvm.org
Mon Jul 13 02:02:11 PDT 2026
https://github.com/LouisLu060211 updated https://github.com/llvm/llvm-project/pull/209111
>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/6] [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/6] 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/6] 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;
>From ff058946ea2e5b8f9ef0d84446d7d0b02bd8d3bd Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Mon, 13 Jul 2026 16:49:54 +0800
Subject: [PATCH 4/6] remove
---
mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp | 13 ++++---------
1 file changed, 4 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
index 9a8acbb61cfab..6748e2cf71804 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
@@ -67,9 +67,7 @@ SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
if (!type.hasStaticShape())
return {};
// Make sure the memref contains only a single element.
- std::optional<int64_t> numElements =
- ShapedType::tryGetNumElements(type.getShape());
- if (!numElements || *numElements != 1)
+ if (type.getNumElements() != 1)
return {};
return {MemorySlot{getResult(), type.getElementType()}};
@@ -292,12 +290,9 @@ struct MemRefDestructurableTypeExternalModel
getSubelementIndexMap(Type type) const {
auto memrefType = llvm::cast<MemRefType>(type);
constexpr int64_t maxMemrefSizeForDestructuring = 16;
- if (!memrefType.hasStaticShape())
- return {};
- std::optional<int64_t> numElements =
- ShapedType::tryGetNumElements(memrefType.getShape());
- if (!numElements || *numElements > maxMemrefSizeForDestructuring ||
- *numElements == 0)
+ if (!memrefType.hasStaticShape() ||
+ memrefType.getNumElements() > maxMemrefSizeForDestructuring ||
+ memrefType.getNumElements() == 1)
return {};
DenseMap<Attribute, Type> destructured;
>From 8bd4d3ff49e0120be740eb50bc30b906406b3ad4 Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Mon, 13 Jul 2026 17:00:08 +0800
Subject: [PATCH 5/6] fix
---
a.mlir | 17 -------
.../Analysis/test-alias-analysis-modref.mlir | 14 ------
mlir/test/Dialect/MemRef/mem2reg.mlir | 15 ------
mlir/test/lib/Analysis/TestAliasAnalysis.cpp | 48 +++++--------------
4 files changed, 11 insertions(+), 83 deletions(-)
delete mode 100644 a.mlir
diff --git a/a.mlir b/a.mlir
deleted file mode 100644
index 899e405a1dabb..0000000000000
--- a/a.mlir
+++ /dev/null
@@ -1,17 +0,0 @@
-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/test/Analysis/test-alias-analysis-modref.mlir b/mlir/test/Analysis/test-alias-analysis-modref.mlir
index 9c1fa17c1249a..2a9c612317d53 100644
--- a/mlir/test/Analysis/test-alias-analysis-modref.mlir
+++ b/mlir/test/Analysis/test-alias-analysis-modref.mlir
@@ -112,17 +112,3 @@ 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 9327db24ff39c..8f937c4efe75e 100644
--- a/mlir/test/Dialect/MemRef/mem2reg.mlir
+++ b/mlir/test/Dialect/MemRef/mem2reg.mlir
@@ -309,18 +309,3 @@ 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 8b33a08db4ab1..d59e9854e925a 100644
--- a/mlir/test/lib/Analysis/TestAliasAnalysis.cpp
+++ b/mlir/test/lib/Analysis/TestAliasAnalysis.cpp
@@ -16,63 +16,36 @@
#include "mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Pass/Pass.h"
-#include "mlir/IR/Value.h"
using namespace mlir;
-namespace mlir {
-namespace test {
-
-/// Print an operation that is used as an operand of an alias query.
+/// Print a value that is used as an operand of an alias query.
static void printAliasOperand(Operation *op) {
- 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();
+ llvm::errs() << op->getAttrOfType<StringAttr>("test.ptr").getValue();
}
/// 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;
}
-
- Operation *owner = value.getDefiningOp();
- assert(owner && "expected non-block argument value to have a defining op");
-
- printAliasOperand(owner);
-
- auto result = cast<mlir::OpResult>(value);
+ OpResult result = cast<OpResult>(value);
+ printAliasOperand(result.getOwner());
llvm::errs() << "#" << result.getResultNumber();
}
+namespace mlir {
+namespace test {
void printAliasResult(AliasResult result, Value lhs, Value rhs) {
printAliasOperand(lhs);
llvm::errs() << " <-> ";
@@ -253,3 +226,4 @@ void registerTestAliasAnalysisPass() {
}
} // namespace test
} // namespace mlir
+
>From 40f41ecd6a705fcaa1790d68ab705f250fc5380e Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Mon, 13 Jul 2026 17:01:50 +0800
Subject: [PATCH 6/6] fix issue
---
mlir/test/lib/Analysis/TestAliasAnalysis.cpp | 4 ----
1 file changed, 4 deletions(-)
diff --git a/mlir/test/lib/Analysis/TestAliasAnalysis.cpp b/mlir/test/lib/Analysis/TestAliasAnalysis.cpp
index d59e9854e925a..0125e403272a8 100644
--- a/mlir/test/lib/Analysis/TestAliasAnalysis.cpp
+++ b/mlir/test/lib/Analysis/TestAliasAnalysis.cpp
@@ -23,8 +23,6 @@ using namespace mlir;
static void printAliasOperand(Operation *op) {
llvm::errs() << op->getAttrOfType<StringAttr>("test.ptr").getValue();
}
-
-/// 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();
@@ -35,7 +33,6 @@ static void printAliasOperand(Value value) {
<< ".region" << region->getRegionNumber();
if (parentBlockNumber != 0)
llvm::errs() << ".block" << parentBlockNumber;
-
llvm::errs() << "#" << arg.getArgNumber();
return;
}
@@ -226,4 +223,3 @@ void registerTestAliasAnalysisPass() {
}
} // namespace test
} // namespace mlir
-
More information about the llvm-commits
mailing list