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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: lonely eagle (linuxlonelyeagle)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/161117.diff


2 Files Affected:

- (modified) mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp (+19-4) 
- (modified) mlir/test/Transforms/remove-dead-values.mlir (+13) 


``````````diff
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
+}

``````````

</details>


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