[Mlir-commits] [mlir] [MLIR][RemoveDeadValues] Fix affine.for induction variable incorrectly removed (PR #172612)
Francisco Geiman Thiesen
llvmlistbot at llvm.org
Wed Dec 17 21:02:45 PST 2025
https://github.com/FranciscoThiesen updated https://github.com/llvm/llvm-project/pull/172612
>From f1a42238c4cfcd2f9202acc63d47440f9c4589f9 Mon Sep 17 00:00:00 2001
From: Francisco Geiman Thiesen <franciscogthiesen at gmail.com>
Date: Wed, 17 Dec 2025 00:59:36 -0800
Subject: [PATCH 1/2] [MLIR][RemoveDeadValues] Fix affine.for induction
variable incorrectly removed
The liveness analysis marks block arguments that are not successor inputs
(like induction variables) as live via visitBranchOperand(). However,
visitBranchOperand() is only invoked when non-forwarded operands exist.
For affine.for with constant bounds:
- Lower/upper bounds are encoded in affine maps (not SSA operands)
- Step is a constant attribute (not an operand)
- Inits are forwarded operands (not non-forwarded)
Since there are no non-forwarded operands, visitBranchOperand() is never
called, and the IV is never marked live, leading to incorrect removal.
This patch fixes the issue by marking non-successor-input block arguments
as live during post-processing of the liveness analysis. This handles all
RegionBranchOpInterface ops unconditionally.
This new approach supersedes the similar logic added in #161117 within
visitBranchOperand(), which only handled region-to-region transitions and
only ran when visitBranchOperand() was called. That code is now removed
as it is redundant.
Also merges two separate IR walks into a single traversal for efficiency.
Fixes #172610
---
.../Analysis/DataFlow/LivenessAnalysis.cpp | 78 +++++++++++--------
mlir/test/Transforms/remove-dead-values.mlir | 19 +++++
2 files changed, 65 insertions(+), 32 deletions(-)
diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
index 20be50c8e8a5b..159199d1e987a 100644
--- a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
@@ -17,6 +17,7 @@
#include <mlir/IR/Operation.h>
#include <mlir/IR/Value.h>
#include <mlir/Interfaces/CallInterfaces.h>
+#include <mlir/Interfaces/ControlFlowInterfaces.h>
#include <mlir/Interfaces/SideEffectInterfaces.h>
#include <mlir/Support/LLVM.h>
@@ -166,25 +167,6 @@ void LivenessAnalysis::visitBranchOperand(OpOperand &operand) {
blocks.push_back(&block);
}
}
-
- // In the block of the successor block argument of RegionBranchOpInterface,
- // there may be arguments of RegionBranchOpInterface, such as the IV of
- // scf.forOp. Explicitly set this argument to live.
- for (Region ®ion : op->getRegions()) {
- SmallVector<RegionSuccessor> successors;
- regionBranchOp.getSuccessorRegions(region, successors);
- for (RegionSuccessor successor : successors) {
- if (successor.isParent())
- continue;
- auto arguments = successor.getSuccessor()->getArguments();
- ValueRange regionInputs = successor.getSuccessorInputs();
- for (auto argument : arguments) {
- if (llvm::find(regionInputs, argument) == regionInputs.end()) {
- argumentNotOperand.push_back(argument);
- }
- }
- }
- }
} else if (isa<BranchOpInterface>(op)) {
// We cannot track all successor blocks of the branch operation(More
// specifically, it's the successor's successor). Additionally, different
@@ -327,28 +309,60 @@ RunLivenessAnalysis::RunLivenessAnalysis(Operation *op) {
solver.load<LivenessAnalysis>(symbolTable);
LDBG() << "Initializing and running solver";
(void)solver.initializeAndRun(op);
+
LDBG() << "RunLivenessAnalysis initialized for op: " << op->getName()
- << " check on unreachable code now:";
- // The framework doesn't visit operations in dead blocks, so we need to
- // explicitly mark them as dead.
+ << ", post-processing values:";
+ // Post-process the analysis results in a single walk:
+ //
+ // 1. Mark non-successor-input block arguments as live. These are arguments
+ // like induction variables (IVs) that are implicitly defined by the op
+ // rather than forwarded from operands. For example, affine.for encodes
+ // its bounds in affine maps rather than as SSA operands, so the IV is
+ // not a successor input and must be explicitly marked live here.
+ //
+ // 2. Create dead states for values in unreachable blocks. The dataflow
+ // framework doesn't visit operations in dead blocks, so we must
+ // explicitly handle them here.
op->walk([&](Operation *op) {
- for (auto result : llvm::enumerate(op->getResults())) {
- if (getLiveness(result.value()))
+ // (1) Handle RegionBranchOpInterface: mark non-successor-input args live.
+ if (auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op)) {
+ SmallVector<RegionSuccessor> successors;
+ regionBranchOp.getSuccessorRegions(RegionBranchPoint::parent(),
+ successors);
+ for (RegionSuccessor &successor : successors) {
+ Region *successorRegion = successor.getSuccessor();
+ if (!successorRegion || successorRegion->empty())
+ continue;
+ Block &entryBlock = successorRegion->front();
+ ValueRange inputs = successor.getSuccessorInputs();
+ for (BlockArgument arg : entryBlock.getArguments()) {
+ if (llvm::find(inputs, arg) == inputs.end()) {
+ LDBG() << "Marking non-successor-input block argument live: "
+ << arg;
+ (void)solver.getOrCreateState<Liveness>(arg)->markLive();
+ }
+ }
+ }
+ }
+
+ // (2) Handle unreachable code: create dead states for unvisited values.
+ for (OpResult result : op->getResults()) {
+ if (getLiveness(result))
continue;
- LDBG() << "Result: " << result.index() << " of "
+ LDBG() << "Result " << result.getResultNumber() << " of "
<< OpWithFlags(op, OpPrintingFlags().skipRegions())
<< " has no liveness info (unreachable), mark dead";
- solver.getOrCreateState<Liveness>(result.value());
+ solver.getOrCreateState<Liveness>(result);
}
- for (auto ®ion : op->getRegions()) {
- for (auto &block : region) {
- for (auto blockArg : llvm::enumerate(block.getArguments())) {
- if (getLiveness(blockArg.value()))
+ for (Region ®ion : op->getRegions()) {
+ for (Block &block : region) {
+ for (BlockArgument blockArg : block.getArguments()) {
+ if (getLiveness(blockArg))
continue;
- LDBG() << "Block argument: " << blockArg.index() << " of "
+ LDBG() << "Block argument " << blockArg.getArgNumber() << " of "
<< OpWithFlags(op, OpPrintingFlags().skipRegions())
<< " has no liveness info, mark dead";
- solver.getOrCreateState<Liveness>(blockArg.value());
+ solver.getOrCreateState<Liveness>(blockArg);
}
}
}
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 71306676d48e9..d00831dae9a14 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -688,6 +688,25 @@ func.func @dead_value_loop_ivs_no_result(%lb: index, %ub: index, %step: index, %
// -----
+// This test verifies that the induction variable in affine.for with constant
+// bounds is not deleted. This is a regression test for a bug where affine.for
+// with constant bounds (no operands for lb/ub/step) would have its IV deleted
+// because visitBranchOperand() was never called (no non-forwarded operands).
+
+// CHECK-LABEL: func @affine_for_iv_constant_bounds
+// CHECK: affine.for %{{.*}} = 0 to 1024 iter_args(%{{.*}} = %{{.*}}) -> (i32)
+func.func @affine_for_iv_constant_bounds() -> i32 {
+ %c1_i32 = arith.constant 1 : i32
+ %c0_i32 = arith.constant 0 : i32
+ %0 = affine.for %iv = 0 to 1024 iter_args(%arg = %c0_i32) -> (i32) {
+ %1 = arith.addi %arg, %c1_i32 : i32
+ affine.yield %1 : i32
+ }
+ return %0 : i32
+}
+
+// -----
+
// CHECK-LABEL: func @op_block_have_dead_arg
func.func @op_block_have_dead_arg(%arg0: index, %arg1: index, %arg2: i1) {
scf.execute_region {
>From d7c936ce42fd51a6fd26155a67ab6e9c3fccb0b1 Mon Sep 17 00:00:00 2001
From: Francisco Geiman Thiesen <franciscogthiesen at gmail.com>
Date: Wed, 17 Dec 2025 21:02:36 -0800
Subject: [PATCH 2/2] Retrigger CI
More information about the Mlir-commits
mailing list