[mlir][dataflow] Fix LivenessAnalysis/RemoveDeadValues handling of loop induction variables

[linuxlonelyeagle]

Fix #157934. In liveness analysis, variables that are not analyzed are set as dead variables, but some variables are definitely live.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: lonely eagle (linuxlonelyeagle)

Changes

Fix #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 --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
+}

[joker-eph]

That seems very ad-hoc to me to look for this specific interface, seems like a workaround for something more important going on: likely that all entry-block arguments should be live, unless RegionBranchOpInterface could say they aren't.

[linuxlonelyeagle]

Maybe I should write the visit RegionBranchOpInterface code for the liveness analysis.

[joker-eph]

We should first be conservative and fix the issues. Then adding support for an interface is useful to optimize more.

[linuxlonelyeagle]

I've implemented it in the visit function, though I'm not entirely sure if it's quite right.

[joker-eph]

This is still using the interface: what I meant by conservative is that the baseline implementation should do nothing in the absence of an interface.
Can we fix the issue that way first?

[linuxlonelyeagle]

I think what you said make sense, and I fixed the bug using the method you mentioned.

[joker-eph]

I think you haven't understood what I'm saying: I'd like to see a fix that does not use any interface and treat the program in the most conservative way.

[linuxlonelyeagle]

I understand your point, but to be honest: it would be rather difficult without an interface. Do you have any suggestions?

[linuxlonelyeagle]

Using interfaces allows for clear differentiation between arguments, but if nothing is specified, these arguments become indistinguishable.This is a major reason why it troubles me.

[joker-eph]

Let's step back and start from basics: if you don't know anything about an operation (there is no interface), you still need to treat it conservatively and not modifying it right?

Why are we modifying an op here without knowing that it's OK to do?