[Mlir-commits] [mlir] [mlir][dataflow] Fix LivenessAnalysis/RemoveDeadValues handling of loop induction variables (PR #161117)

[lonely eagle]

https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/161117

>From 59621f88123a7d90d1d4a18898ba72ffaf76c5fe Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Mon, 29 Sep 2025 02:42:42 +0000
Subject: [PATCH 1/2] fix liveness-analysis in the loop ivs case.

---
 .../Analysis/DataFlow/LivenessAnalysis.cpp    | 23 +++++++++++++++----
 mlir/test/Transforms/remove-dead-values.mlir  | 13 +++++++++++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
index d705d8d4c7819..f540870113e3f 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/LoopLikeInterface.h>
 #include <mlir/Interfaces/SideEffectInterfaces.h>
 #include <mlir/Support/LLVM.h>
 
@@ -309,15 +310,29 @@ RunLivenessAnalysis::RunLivenessAnalysis(Operation *op) {
              << " has no liveness info (unreachable), mark dead";
       solver.getOrCreateState<Liveness>(result.value());
     }
+    SmallVector<Value> mustLiveValues;
+    if (auto loopOp = dyn_cast<LoopLikeOpInterface>(op)) {
+      std::optional<SmallVector<Value>> ivs = loopOp.getLoopInductionVars();
+      if (ivs.has_value())
+        mustLiveValues.append(*ivs);
+    }
     for (auto &region : op->getRegions()) {
       for (auto &block : region) {
         for (auto blockArg : llvm::enumerate(block.getArguments())) {
           if (getLiveness(blockArg.value()))
             continue;
-          LDBG() << "Block argument: " << blockArg.index() << " of "
-                 << OpWithFlags(op, OpPrintingFlags().skipRegions())
-                 << " has no liveness info, mark dead";
-          solver.getOrCreateState<Liveness>(blockArg.value());
+          if (llvm::find(mustLiveValues, blockArg.value())) {
+            LDBG() << "Block argument: " << blockArg.index() << " of "
+                   << OpWithFlags(op, OpPrintingFlags().skipRegions())
+                   << " is must value, mark live";
+            (void)solver.getOrCreateState<Liveness>(blockArg.value())
+                ->markLive();
+          } else {
+            LDBG() << "Block argument: " << blockArg.index() << " of "
+                   << OpWithFlags(op, OpPrintingFlags().skipRegions())
+                   << " has no liveness info, mark dead";
+            solver.getOrCreateState<Liveness>(blockArg.value());
+          }
         }
       }
     }
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 56449469dc29f..979851bc3162d 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -649,3 +649,16 @@ func.func @callee(%arg0: index, %arg1: index, %arg2: index) -> index {
   %res = call @mutl_parameter(%arg0, %arg1, %arg2) : (index, index, index) -> (index)
   return %res : index
 }
+
+// -----
+
+// This test verifies that the induction variables in loops are not deleted.
+
+// CHECK-LABEL: func @dead_value_loop_ivs
+func.func @dead_value_loop_ivs(%lb: index, %ub: index, %step: index, %b : i1) -> i1 {
+  %loop_ret = scf.for %iv = %lb to %ub step %step iter_args(%iter = %b) -> (i1) {
+    cf.assert %b, "loop not dead"
+    scf.yield %b : i1
+  }
+  return %loop_ret : i1
+}

>From fbed4d1d1bdbbd01347dce9d5da18828630b75f3 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Thu, 2 Oct 2025 02:56:13 +0000
Subject: [PATCH 2/2] add LoopLikeOpInterface visit logic in the visit
 function.

---
 .../Analysis/DataFlow/LivenessAnalysis.cpp    | 35 +++++++++----------
 1 file changed, 17 insertions(+), 18 deletions(-)

diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
index f540870113e3f..3611d07695bf9 100644
--- a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
@@ -196,6 +196,18 @@ void LivenessAnalysis::visitBranchOperand(OpOperand &operand) {
           break;
         }
       }
+
+      // If the parentOp is live and it implements the LoopLiveOpInterface, then
+      // set its IV as live.
+      if (mayLive && isa<LoopLikeOpInterface>(parentOp)) {
+        auto loopOp = cast<LoopLikeOpInterface>(parentOp);
+        std::optional<SmallVector<Value>> ivs = loopOp.getLoopInductionVars();
+        if (ivs.has_value()) {
+          for (auto iv : *ivs) {
+            getLatticeElement(iv)->markLive();
+          }
+        }
+      }
     } else {
       // When the op is a `RegionBranchTerminatorOpInterface`, like an
       // `scf.condition` op or return-like, like an `scf.yield` op, its branch
@@ -310,29 +322,16 @@ RunLivenessAnalysis::RunLivenessAnalysis(Operation *op) {
              << " has no liveness info (unreachable), mark dead";
       solver.getOrCreateState<Liveness>(result.value());
     }
-    SmallVector<Value> mustLiveValues;
-    if (auto loopOp = dyn_cast<LoopLikeOpInterface>(op)) {
-      std::optional<SmallVector<Value>> ivs = loopOp.getLoopInductionVars();
-      if (ivs.has_value())
-        mustLiveValues.append(*ivs);
-    }
+
     for (auto &region : op->getRegions()) {
       for (auto &block : region) {
         for (auto blockArg : llvm::enumerate(block.getArguments())) {
           if (getLiveness(blockArg.value()))
             continue;
-          if (llvm::find(mustLiveValues, blockArg.value())) {
-            LDBG() << "Block argument: " << blockArg.index() << " of "
-                   << OpWithFlags(op, OpPrintingFlags().skipRegions())
-                   << " is must value, mark live";
-            (void)solver.getOrCreateState<Liveness>(blockArg.value())
-                ->markLive();
-          } else {
-            LDBG() << "Block argument: " << blockArg.index() << " of "
-                   << OpWithFlags(op, OpPrintingFlags().skipRegions())
-                   << " has no liveness info, mark dead";
-            solver.getOrCreateState<Liveness>(blockArg.value());
-          }
+          LDBG() << "Block argument: " << blockArg.index() << " of "
+                 << OpWithFlags(op, OpPrintingFlags().skipRegions())
+                 << " has no liveness info, mark dead";
+          solver.getOrCreateState<Liveness>(blockArg.value());
         }
       }
     }