[Mlir-commits] [mlir] [mlir][CFGToSCF] fix use-after-free segfault in lift-cf-to-scf (PR #227783)
Keyi Zhang
llvmlistbot at llvm.org
Wed Sep 30 09:48:10 PDT 2026
https://github.com/Kuree created https://github.com/llvm/llvm-project/pull/227783
See example input here: https://godbolt.org/z/565nnzrvf
In `transformToReduceLoop`, `checkValue` iterated over the uses of `value` with `make_early_inc_range` and, on the first use outside the loop, appended successor operands to the terminator of every latch predecessor:
```C++
getMutableSuccessorOperands(*iter, iter.getSuccessorIndex())
.append(succOperand);
```
If one of those terminators itself uses `value` (e.g. as its branch condition), the append can reallocate that terminator's operand storage. `make_early_inc_range` only protects the current element from removal; it holds a raw pointer to the next one. After the reallocation, the next pointer is no longer valid, hence the segfault.
The fix splits `checkValue` into three steps: check whether any use escapes the loop, create the block arguments and append the successor operands with no use iterator alive, then replace the escaping uses with `replaceUsesWithIf`.
The added test segfaults on the unfixed code and is reported as a heap-use-after-free by ASan/valgrind.
>From 5562ec90ed78826baec87e8de2f1e63aab45b5d6 Mon Sep 17 00:00:00 2001
From: Keyi Zhang <keyi at efficient.computer>
Date: Wed, 30 Sep 2026 09:12:23 -0700
Subject: [PATCH] [mlir][CFGToSCF] fix use-after-free segfault in
lift-cf-to-scf
---
mlir/lib/Transforms/Utils/CFGToSCF.cpp | 91 ++++++++++---------
.../Conversion/ControlFlowToSCF/test.mlir | 39 ++++++++
2 files changed, 85 insertions(+), 45 deletions(-)
diff --git a/mlir/lib/Transforms/Utils/CFGToSCF.cpp b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
index f947e889b56fa..b3ac2c44e4b46 100644
--- a/mlir/lib/Transforms/Utils/CFGToSCF.cpp
+++ b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
@@ -733,55 +733,56 @@ transformToReduceLoop(Block *loopHeader, Block *exitBlock,
return iter->second;
};
+ // Returns true if `use` is outside the loop body.
+ auto isEscapingUse = [&](OpOperand &use) {
+ // Go through all the parent blocks and find the one part of the region
+ // of the loop. If the block is part of the loop, then the value does
+ // not escape the loop through this use.
+ Block *currBlock = use.getOwner()->getBlock();
+ while (currBlock && currBlock->getParent() != loopHeader->getParent())
+ currBlock = currBlock->getParentOp()->getBlock();
+ return !loopBlocks.contains(currBlock);
+ };
+
auto checkValue = [&](Value value) {
- Value blockArgument;
- for (OpOperand &use : llvm::make_early_inc_range(value.getUses())) {
- // Go through all the parent blocks and find the one part of the region
- // of the loop. If the block is part of the loop, then the value does
- // not escape the loop through this use.
- Block *currBlock = use.getOwner()->getBlock();
- while (currBlock && currBlock->getParent() != loopHeader->getParent())
- currBlock = currBlock->getParentOp()->getBlock();
- if (loopBlocks.contains(currBlock))
- continue;
+ // Block arguments are only created if the value escapes the loop.
+ if (llvm::none_of(value.getUses(), isEscapingUse))
+ return;
+
+ Value blockArgument =
+ exitBlock->addArgument(value.getType(), value.getLoc());
+ loopHeader->addArgument(value.getType(), value.getLoc());
+
+ // `value` might be defined in a block that does not dominate `latch`
+ // but previously dominated an exit block with a use.
+ // In this case, add a block argument to the latch and go through all
+ // predecessors. If the value dominates the predecessor, pass the
+ // value as a successor operand, otherwise pass poison.
+ // The above is unnecessary if the value is a block argument of the
+ // latch or if `value` dominates all predecessors.
+ Value argument = value;
+ if (value.getParentBlock() != latch &&
+ llvm::any_of(latch->getPredecessors(), [&](Block *pred) {
+ return !loopBlockDominates(pred);
+ })) {
+ argument = latch->addArgument(value.getType(), value.getLoc());
+ for (auto iter = latch->pred_begin(); iter != latch->pred_end();
+ ++iter) {
+ Value succOperand = value;
+ if (!loopBlockDominates(*iter))
+ succOperand = getUndefValue(value.getType());
- // Block argument is only created the first time it is required.
- if (!blockArgument) {
- blockArgument =
- exitBlock->addArgument(value.getType(), value.getLoc());
- loopHeader->addArgument(value.getType(), value.getLoc());
-
- // `value` might be defined in a block that does not dominate `latch`
- // but previously dominated an exit block with a use.
- // In this case, add a block argument to the latch and go through all
- // predecessors. If the value dominates the predecessor, pass the
- // value as a successor operand, otherwise pass poison.
- // The above is unnecessary if the value is a block argument of the
- // latch or if `value` dominates all predecessors.
- Value argument = value;
- if (value.getParentBlock() != latch &&
- llvm::any_of(latch->getPredecessors(), [&](Block *pred) {
- return !loopBlockDominates(pred);
- })) {
- argument = latch->addArgument(value.getType(), value.getLoc());
- for (auto iter = latch->pred_begin(); iter != latch->pred_end();
- ++iter) {
- Value succOperand = value;
- if (!loopBlockDominates(*iter))
- succOperand = getUndefValue(value.getType());
-
- getMutableSuccessorOperands(*iter, iter.getSuccessorIndex())
- .append(succOperand);
- }
- }
-
- loopHeaderSuccessorOperands.push_back(argument);
- for (Edge edge : successorEdges(latch))
- edge.getMutableSuccessorOperands().append(argument);
+ getMutableSuccessorOperands(*iter, iter.getSuccessorIndex())
+ .append(succOperand);
}
-
- use.set(blockArgument);
}
+
+ loopHeaderSuccessorOperands.push_back(argument);
+ for (Edge edge : successorEdges(latch))
+ edge.getMutableSuccessorOperands().append(argument);
+
+ // Replace all uses outside the loop body with the exit block argument.
+ value.replaceUsesWithIf(blockArgument, isEscapingUse);
};
if (loopBlock == latch)
diff --git a/mlir/test/Conversion/ControlFlowToSCF/test.mlir b/mlir/test/Conversion/ControlFlowToSCF/test.mlir
index 78ac69aec3fe4..fbdfcca2bbae7 100644
--- a/mlir/test/Conversion/ControlFlowToSCF/test.mlir
+++ b/mlir/test/Conversion/ControlFlowToSCF/test.mlir
@@ -797,3 +797,42 @@ func.func @nested_region_outside_loop_use() {
// CHECK: scf.execute_region
// CHECK-NEXT: "test.foo"(%[[RES]])
+
+// -----
+
+// Test that %v, used by a latch predecessor (^bb3) and outside the inner loop
+// (^bb5), gets properly forwarded out of the loop.
+
+func.func @escaping_value_used_by_latch_predecessor(%cond: i1, %arg: i32) -> i32 {
+ cf.br ^bb0
+^bb0:
+ cf.br ^bb1(%arg, %arg : i32, i32)
+^bb1(%0: i32, %1: i32):
+ cf.cond_br %cond, ^bb2, ^bb6
+^bb2:
+ %v = arith.cmpi eq, %arg, %arg : i32
+ cf.cond_br %cond, ^bb3, ^bb5
+^bb3:
+ cf.cond_br %v, ^bb0, ^bb4(%arg : i32)
+^bb4(%2: i32):
+ cf.br ^bb1(%arg, %arg : i32, i32)
+^bb5:
+ cf.cond_br %v, ^bb0, ^bb6
+^bb6:
+ return %arg : i32
+}
+
+// CHECK-LABEL: func @escaping_value_used_by_latch_predecessor
+// CHECK-SAME: %[[COND:[[:alnum:]]+]]
+// CHECK-SAME: %[[ARG:[[:alnum:]]+]]
+// CHECK: scf.while
+// CHECK: %[[INNER:.*]]:6 = scf.while
+// CHECK: %[[IF:.*]]:7 = scf.if %[[COND]]
+// CHECK: %[[V:.*]] = arith.cmpi eq, %[[ARG]], %[[ARG]]
+// CHECK: scf.if %[[V]]
+// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[V]]
+// CHECK: scf.condition(%{{.*}}) %[[IF]]#0, %[[IF]]#1, %[[IF]]#6, %[[IF]]#2, %[[IF]]#3, %[[IF]]#4
+// CHECK: scf.index_switch
+// CHECK: default
+// CHECK: scf.if %[[INNER]]#2
+// CHECK: return %[[ARG]]
More information about the Mlir-commits
mailing list