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

[lonely eagle]

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

Fix https://github.com/llvm/llvm-project/issues/157934. In liveness analysis, variables that are not analyzed are set as dead variables, but some variables are definitely live.

>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] 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
+}