[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 b9b6a5bd15621cf05128252b8bd67e79a554a029 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 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 ensures IVs
are marked live regardless of whether the op has non-forwarded operands.

The fix is complementary to visitBranchOperand(): that function handles
region-to-region transitions when it runs, while this fix handles
parent-to-region transitions unconditionally.

Also merges two separate IR walks into a single traversal for efficiency.

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

diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
index 20be50c8e8a5b..81f27af3c3c45 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,28 +328,63 @@ 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. visitBranchOperand() also handles
+  //    this for region-to-region transitions, but it is only invoked when
+  //    non-forwarded operands exist. For ops like affine.for with constant
+  //    bounds (encoded in affine maps rather than as SSA operands),
+  //    visitBranchOperand() is never called. This ensures IVs are marked live
+  //    regardless of whether the op has non-forwarded operands.
+  //
+  // 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 &region : op->getRegions()) {
-      for (auto &block : region) {
-        for (auto blockArg : llvm::enumerate(block.getArguments())) {
-          if (getLiveness(blockArg.value()))
+    for (Region &region : 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 {