[MLIR] Fix Liveness analysis handling of unreachable code

[joker-eph]

This patch is forcing all values to be initialized by the LivenessAnalysis, even in dead blocks. The dataflow framework will skip visiting values when its already knows that a block is dynamically unreachable, so this requires specific handling.
Downstream code could consider that the absence of liveness is the same a "dead".
However as the code is mutated, new value can be introduced, and a transformation like "RemoveDeadValue" must conservatively consider that the absence of liveness information meant that we weren't sure if a value was dead (it could be a newly introduced value.

Fixes #153906

[llvmbot]

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

Changes

The code was "conservatively" considering that the absence of liveness information meant that we weren't sure if a value was dead. However the dataflow framework will skip visiting these values when its already knows that a block is dynamically unreachable.
This leds to a crash in the provided test case since the producer operation in a reachable block is actually deleted (its liveness is correctly set to "dead") but the user of the operation was "conservatively" preserved and left with nullptr operands.

Fixes #153906

---

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

2 Files Affected:

• (modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+48-16)
• (modified) mlir/test/Transforms/remove-dead-values.mlir (+21)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp index 4ccb83f3ad298..6bbe9b13dd8c3 100644 --- a/mlir/lib/Transforms/RemoveDeadValues.cpp +++ b/mlir/lib/Transforms/RemoveDeadValues.cpp @@ -50,6 +50,7 @@ #include "mlir/Support/LLVM.h" #include "mlir/Transforms/FoldUtils.h" #include "mlir/Transforms/Passes.h" +#include "mlir/Transforms/RegionUtils.h" #include "llvm/ADT/STLExtras.h" #include "llvm/Support/Debug.h" #include "llvm/Support/DebugLog.h" @@ -124,17 +125,13 @@ static bool hasLive(ValueRange values, const DenseSet<Value> &nonLiveSet, } const Liveness *liveness = la.getLiveness(value); - if (!liveness) { - LDBG() << "Value " << value - << " has no liveness info, conservatively considered live"; - return true; - } - if (liveness->isLive) { - LDBG() << "Value " << value << " is live according to liveness analysis"; - return true; - } else { - LDBG() << "Value " << value << " is dead according to liveness analysis"; + if (!liveness || !liveness->isLive) { + LDBG() << "Value is dead (" << (liveness ? "" : "no ") << "liveness info)" + << value; + continue; } + LDBG() << "Value is live: " << value; + return true; } return false; } @@ -163,6 +160,7 @@ static BitVector markLives(ValueRange values, const DenseSet<Value> &nonLiveSet, if (!liveness) { LDBG() << "Value " << value << " at index " << index << " has no liveness info, conservatively considered live"; + lives.reset(index); continue; } if (!liveness->isLive) { @@ -258,18 +256,17 @@ static SmallVector<OpOperand *> operandsToOpOperands(OperandRange operands) { static void processSimpleOp(Operation *op, RunLivenessAnalysis &la, DenseSet<Value> &nonLiveSet, RDVFinalCleanupList &cl) { - LDBG() << "Processing simple op: " << *op; if (!isMemoryEffectFree(op) || hasLive(op->getResults(), nonLiveSet, la)) { - LDBG() - << "Simple op is not memory effect free or has live results, skipping: " - << *op; + LDBG() << "Simple op is not memory effect free or has live results, " + "preserving it: " + << OpWithFlags(op, OpPrintingFlags().skipRegions()); return; } LDBG() << "Simple op has all dead results and is memory effect free, scheduling " "for removal: " - << *op; + << OpWithFlags(op, OpPrintingFlags().skipRegions()); cl.operations.push_back(op); collectNonLiveValues(nonLiveSet, op->getResults(), BitVector(op->getNumResults(), true)); @@ -728,19 +725,31 @@ static void processBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la, /// Removes dead values collected in RDVFinalCleanupList. /// To be run once when all dead values have been collected. static void cleanUpDeadVals(RDVFinalCleanupList &list) { + LDBG() << "Starting cleanup of dead values..."; + // 1. Operations + LDBG() << "Cleaning up " << list.operations.size() << " operations"; for (auto &op : list.operations) { + LDBG() << "Erasing operation: " + << OpWithFlags(op, OpPrintingFlags().skipRegions()); op->dropAllUses(); op->erase(); } // 2. Values + LDBG() << "Cleaning up " << list.values.size() << " values"; for (auto &v : list.values) { + LDBG() << "Dropping all uses of value: " << v; v.dropAllUses(); } // 3. Functions + LDBG() << "Cleaning up " << list.functions.size() << " functions"; for (auto &f : list.functions) { + LDBG() << "Cleaning up function: " << f.funcOp.getOperation()->getName(); + LDBG() << " Erasing " << f.nonLiveArgs.count() << " non-live arguments"; + LDBG() << " Erasing " << f.nonLiveRets.count() + << " non-live return values"; // Some functions may not allow erasing arguments or results. These calls // return failure in such cases without modifying the function, so it's okay // to proceed. @@ -749,44 +758,67 @@ static void cleanUpDeadVals(RDVFinalCleanupList &list) { } // 4. Operands + LDBG() << "Cleaning up " << list.operands.size() << " operand lists"; for (OperationToCleanup &o : list.operands) { - if (o.op->getNumOperands() > 0) + if (o.op->getNumOperands() > 0) { + LDBG() << "Erasing " << o.nonLive.count() + << " non-live operands from operation: " + << OpWithFlags(o.op, OpPrintingFlags().skipRegions()); o.op->eraseOperands(o.nonLive); + } } // 5. Results + LDBG() << "Cleaning up " << list.results.size() << " result lists"; for (auto &r : list.results) { + LDBG() << "Erasing " << r.nonLive.count() + << " non-live results from operation: " + << OpWithFlags(r.op, OpPrintingFlags().skipRegions()); dropUsesAndEraseResults(r.op, r.nonLive); } // 6. Blocks + LDBG() << "Cleaning up " << list.blocks.size() << " block argument lists"; for (auto &b : list.blocks) { // blocks that are accessed via multiple codepaths processed once if (b.b->getNumArguments() != b.nonLiveArgs.size()) continue; + LDBG() << "Erasing " << b.nonLiveArgs.count() + << " non-live arguments from block: " << b.b; // it iterates backwards because erase invalidates all successor indexes for (int i = b.nonLiveArgs.size() - 1; i >= 0; --i) { if (!b.nonLiveArgs[i]) continue; + LDBG() << " Erasing block argument " << i << ": " << b.b->getArgument(i); b.b->getArgument(i).dropAllUses(); b.b->eraseArgument(i); } } // 7. Successor Operands + LDBG() << "Cleaning up " << list.successorOperands.size() + << " successor operand lists"; for (auto &op : list.successorOperands) { SuccessorOperands successorOperands = op.branch.getSuccessorOperands(op.successorIndex); // blocks that are accessed via multiple codepaths processed once if (successorOperands.size() != op.nonLiveOperands.size()) continue; + LDBG() << "Erasing " << op.nonLiveOperands.count() + << " non-live successor operands from successor " + << op.successorIndex << " of branch: " + << OpWithFlags(op.branch, OpPrintingFlags().skipRegions()); // it iterates backwards because erase invalidates all successor indexes for (int i = successorOperands.size() - 1; i >= 0; --i) { if (!op.nonLiveOperands[i]) continue; + LDBG() << " Erasing successor operand " << i << ": " + << successorOperands[i]; successorOperands.erase(i); } } + + LDBG() << "Finished cleanup of dead values"; } struct RemoveDeadValues : public impl::RemoveDeadValuesBase<RemoveDeadValues> { diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir index 9ded6a30d9c95..0f8d757086e87 100644 --- a/mlir/test/Transforms/remove-dead-values.mlir +++ b/mlir/test/Transforms/remove-dead-values.mlir @@ -571,3 +571,24 @@ module @return_void_with_unused_argument { } } +// ----- + +// CHECK-LABEL: module @dynamically_unreachable +module @dynamically_unreachable { + func.func @dynamically_unreachable() { + // This value is used by an operation in a dynamically unreachable block. + %zero = arith.constant 0 : i64 + + // Dataflow analysis knows from the constant condition that + // ^bb1 is unreachable + %false = arith.constant false + cf.cond_br %false, ^bb1, ^bb4 + ^bb1: + // This unreachable operation should be removed. + // CHECK-NOT: arith.cmpi + %3 = arith.cmpi eq, %zero, %zero : i64 + cf.br ^bb1 + ^bb4: + return + } +} 

[joker-eph]

There is likely a problem with:

 // the execution of this pass might create new SSA values when erasing some // of the results of an op and we know that these new values are live // (because they weren't erased) and also their liveness is null because // liveness analysis ran before their creation. 

But we don't have any test coverage for this!! :(

[Copilot]

Pull Request Overview

This PR fixes the handling of unreachable code in MLIR's RemoveDeadValues pass by ensuring that liveness analysis properly initializes all values, even those in dynamically unreachable blocks. The dataflow framework normally skips visiting values in blocks it knows are unreachable, but the RemoveDeadValues transformation needs explicit liveness information to make conservative decisions about value removal.

• Forces liveness analysis to initialize all values in dead blocks
• Adds debugging output to the cleanup process
• Includes comprehensive tests for unreachable code scenarios

Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp	Adds post-analysis walk to mark values in dead blocks as having explicit liveness state
mlir/lib/Transforms/RemoveDeadValues.cpp	Enhances debug logging throughout the cleanup process
mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp	Adds debug logging to track dataflow analysis execution
mlir/test/Transforms/remove-dead-values.mlir	Adds test case for operations in dynamically unreachable blocks
mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir	Adds test verifying liveness analysis of dead block operations
mlir/test/lib/Analysis/DataFlow/TestLivenessAnalysis.cpp	Removes extraneous whitespace

[joker-eph]

There is likely a problem with:

 // the execution of this pass might create new SSA values when erasing some // of the results of an op and we know that these new values are live // (because they weren't erased) and also their liveness is null because // liveness analysis ran before their creation. 

But we don't have any test coverage for this!! :(

This is fixed in the new approach in this PR. Still sad there is no test for this.

[jpienaar]

Sounds reasonable to flag unreachable as dead explicitly. As mentioned this is conceptually the current state, except downstream consumers can't rely on that for unset.