[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]

Isn't this something where we should instead propagate the liveness of the result backward to the loop carried value instead of marking it unconditionally live?

[joker-eph]

That is: if the loop had a result but unused, we would want to be able to delete it.

In this loop:

  %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
  }

If %loop_ret has no use, then the loop can be turned into:

  scf.for %iv = %lb to %ub step %step {
    cf.assert %b, "loop not dead"
    scf.yield
  }

But maybe RemoveDeadValues can't remove results of an operation like scf.for?

[linuxlonelyeagle]

I've written a test; perhaps it will do.

func.func @dead_value_loop_ivs_has_result(%lb: index, %ub: index, %step: index, %b: i1) -> i1 {
  %loop_ret, %a = scf.for %iv = %lb to %ub step %step iter_args(%iter = %b, %iter1 = %lb) -> (i1, index) {
    cf.assert %b, "loop not dead"
    scf.yield %b, %lb : i1, index
  }
  return %loop_ret : i1
}

[joker-eph]

That's conservative, but it's a bug fix so let's land it.

[github-actions]

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:

git-clang-format --diff origin/main HEAD --extensions cpp -- mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp --diff_from_common_commit

⚠️
The reproduction instructions above might return results for more than one PR
in a stack if you are using a stacked PR workflow. You can limit the results by
changing origin/main to the base branch/commit you want to compare against.
⚠️

View the diff from clang-format here.

diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
index 70e058d2b..5a9c99c0e 100644
--- a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
@@ -247,9 +247,10 @@ void LivenessAnalysis::visitBranchOperand(OpOperand &operand) {
     for (BlockArgument argument : argumentNotOperand) {
       Liveness *argumentLiveness = getLatticeElement(argument);
       LDBG() << "Marking RegionBranchOp's argument live: " << argument;
-      // TODO: this is overly conservative: we should be able to eliminate unused 
-      // values in a RegionBranchOpInterface operation but that may requires removing
-      // operation results which is beyond current capabilities of this pass right now.
+      // TODO: this is overly conservative: we should be able to eliminate
+      // unused values in a RegionBranchOpInterface operation but that may
+      // requires removing operation results which is beyond current
+      // capabilities of this pass right now.
       propagateIfChanged(argumentLiveness, argumentLiveness->markLive());
     }
   }

[linuxlonelyeagle]

That's conservative, but it's a bug fix so let's land it.

Although it's a bug fix, strictly speaking it's my first pull request fixing data flow analysis, so it makes sense.I'm really glad to communicate with you—our conversations always leave me with gains.😘