[Mlir-commits] [mlir] [MLIR][RemoveDeadValues] Fix affine.for induction variable incorrectly removed (PR #172612)

[Francisco Geiman Thiesen]

https://github.com/FranciscoThiesen updated https://github.com/llvm/llvm-project/pull/172612

>From d5035896e13a235b16f1d1970776d2d6e18b6a8b 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] [MLIR][RemoveDeadValues] Fix affine.for induction variable
 incorrectly removed

The fix in PR #161117 for issue #157934 addressed scf.for induction
variable deletion by marking IVs as live in visitBranchOperand().
However, this fix doesn't cover affine.for with constant bounds because
visitBranchOperand() is only called for non-forwarded operands.

For affine.for with constant bounds like `affine.for %iv = 0 to 1024`:
- Lower/upper bounds are encoded in affine maps (no operands)
- Step is a constant attribute (no 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 proactively marking non-successor-input
block arguments (like IVs) as live during RunLivenessAnalysis
initialization, rather than relying on visitBranchOperand() being called.

Fixes #172610
---
 .../Analysis/DataFlow/LivenessAnalysis.cpp    | 27 +++++++++++++++++++
 mlir/test/Transforms/remove-dead-values.mlir  | 19 +++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
index 20be50c8e8a5b..f884c76cf2918 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>
 
@@ -327,6 +328,32 @@ RunLivenessAnalysis::RunLivenessAnalysis(Operation *op) {
   solver.load<LivenessAnalysis>(symbolTable);
   LDBG() << "Initializing and running solver";
   (void)solver.initializeAndRun(op);
+
+  // Mark block arguments of RegionBranchOpInterface ops that are NOT successor
+  // inputs as live. These include induction variables (IVs) like those in
+  // affine.for or scf.for. The fix in visitBranchOperand() only handles this
+  // when non-forwarded operands exist, but ops like affine.for with constant
+  // bounds have no non-forwarded operands, so visitBranchOperand() is never
+  // called. We must handle this case here during initialization.
+  op->walk([&](RegionBranchOpInterface regionBranchOp) {
+    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()) {
+          // This is an IV or similar non-successor-input argument.
+          // Mark it live unconditionally.
+          LDBG() << "Marking non-successor-input block argument live: " << arg;
+          (void)solver.getOrCreateState<Liveness>(arg)->markLive();
+        }
+      }
+    }
+  });
   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
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 {