[MLIR][RemoveDeadValues] Fix affine.for induction variable incorrectly removed #172612
Conversation
|
@llvm/pr-subscribers-mlir Author: Francisco Geiman Thiesen (FranciscoThiesen) ChangesSummaryThe fix in PR #161117 for issue #157934 addressed For
Since there are no non-forwarded operands, The FixThis patch fixes the issue by proactively marking non-successor-input block arguments (like IVs) as live during The fix walks all Test Plan
Fixes #172610 Full diff: https://github.com/llvm/llvm-project/pull/172612.diff 2 Files Affected:
diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
index 20be50c8e8a5b..878ae39e2a5fd 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/ControlFlowInterfaces.h>
#include <mlir/Interfaces/SideEffectInterfaces.h>
#include <mlir/Support/LLVM.h>
@@ -327,6 +328,32 @@ RunLivenessAnalysis::RunLivenessAnalysis(Operation *op) {
solver.load<LivenessAnalysis>(symbolTable);
LDBG() << "Initializing and running solver";
(void)solver.initializeAndRun(op);
+
+ // Mark block arguments of RegionBranchOpInterface ops that are NOT successor
+ // inputs as live. These include induction variables (IVs) like those in
+ // affine.for or scf.for. The fix in visitBranchOperand() only handles this
+ // when non-forwarded operands exist, but ops like affine.for with constant
+ // bounds have no non-forwarded operands, so visitBranchOperand() is never
+ // called. We must handle this case here during initialization.
+ op->walk([&](RegionBranchOpInterface regionBranchOp) {
+ SmallVector<RegionSuccessor> successors;
+ regionBranchOp.getSuccessorRegions(RegionBranchPoint::parent(), successors);
+ for (RegionSuccessor &successor : successors) {
+ Region *successorRegion = successor.getSuccessor();
+ if (!successorRegion || successorRegion->empty())
+ continue;
+ Block &entryBlock = successorRegion->front();
+ ValueRange inputs = successor.getSuccessorInputs();
+ for (BlockArgument arg : entryBlock.getArguments()) {
+ if (llvm::find(inputs, arg) == inputs.end()) {
+ // This is an IV or similar non-successor-input argument.
+ // Mark it live unconditionally.
+ LDBG() << "Marking non-successor-input block argument live: " << arg;
+ solver.getOrCreateState<Liveness>(arg)->markLive();
+ }
+ }
+ }
+ });
LDBG() << "RunLivenessAnalysis initialized for op: " << op->getName()
<< " check on unreachable code now:";
// The framework doesn't visit operations in dead blocks, so we need to
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 71306676d48e9..d00831dae9a14 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -688,6 +688,25 @@ func.func @dead_value_loop_ivs_no_result(%lb: index, %ub: index, %step: index, %
// -----
+// This test verifies that the induction variable in affine.for with constant
+// bounds is not deleted. This is a regression test for a bug where affine.for
+// with constant bounds (no operands for lb/ub/step) would have its IV deleted
+// because visitBranchOperand() was never called (no non-forwarded operands).
+
+// CHECK-LABEL: func @affine_for_iv_constant_bounds
+// CHECK: affine.for %{{.*}} = 0 to 1024 iter_args(%{{.*}} = %{{.*}}) -> (i32)
+func.func @affine_for_iv_constant_bounds() -> i32 {
+ %c1_i32 = arith.constant 1 : i32
+ %c0_i32 = arith.constant 0 : i32
+ %0 = affine.for %iv = 0 to 1024 iter_args(%arg = %c0_i32) -> (i32) {
+ %1 = arith.addi %arg, %c1_i32 : i32
+ affine.yield %1 : i32
+ }
+ return %0 : i32
+}
+
+// -----
+
// CHECK-LABEL: func @op_block_have_dead_arg
func.func @op_block_have_dead_arg(%arg0: index, %arg1: index, %arg2: i1) {
scf.execute_region {
|
🐧 Linux x64 Test Results
✅ The build succeeded and all tests passed. |
FranciscoThiesen
force-pushed
the
fix-affine-for-iv-removal
branch
from
6275d43 to
d503589
Compare
|
|
||
| // Mark block arguments of RegionBranchOpInterface ops that are NOT successor | ||
| // inputs as live. These include induction variables (IVs) like those in | ||
| // affine.for or scf.for. The fix in visitBranchOperand() only handles this |
There was a problem hiding this comment.
I wouldn't expect to read "The fix in visitBranchOperand()" in a code comment, this is too contextual to your historical investigation.
Also can you clarify how this new code interacts with the fix that was made to visitBranchOperand()?
Is the handling in visitBranchOperand() still relevant or is this superseding it?
There was a problem hiding this comment.
Thanks @joker-eph for putting forward your point of view; I am considering how to respond to this PR.
I thank this is make sense https://discourse.llvm.org/t/rfc-add-visitbranchregionargument-interface-to-sparsedataflowanalysis/89061.
There was a problem hiding this comment.
I've updated the comment.
In a bit more detail, after the last commit:
- visitBranchOperand() handles region-to-region transitions but is only called when non-forwarded operands exist
- This fix handles parent-to-region transitions unconditionally
- They are complementary: visitBranchOperand() is still relevant for ops that have non-forwarded operands, and it does other things beyond
marking IVs (memory effect checks, result liveness propagation) - This fix ensures IVs are caught regardless of whether visitBranchOperand() runs
There was a problem hiding this comment.
hey are complementary: visitBranchOperand() is still relevant for ops that have non-forwarded operands, and it does other things beyond marking IVs (memory effect checks, result liveness propagation)
Of course I wasn't asking about visitBranchOperand as a whole, but specifically about #161117
There was a problem hiding this comment.
Sorry for that.
For any loop (scf.for, affine.for), the IV is in the entry block and is NOT a successor input in BOTH transition types. So both approaches would identify and mark the same IV.
The key insight: My fix runs unconditionally after the analysis completes. So:
- For scf.for with variable bounds: [mlir][dataflow] Fix LivenessAnalysis/RemoveDeadValues handling of loop induction variables #161117 marks IV live during analysis, then my fix marks it live again after
(redundant but harmless) - For affine.for with constant bounds: visitBranchOperand() never called, so [mlir][dataflow] Fix LivenessAnalysis/RemoveDeadValues handling of loop induction variables #161117 never runs. My fix marks IV live after analysis (essential)
For the specific purpose of marking IVs live, the #161117 fix is now redundant - my fix will mark IVs live for ALL
RegionBranchOpInterface ops regardless of whether visitBranchOperand() runs. The #161117 code could theoretically be removed, but it's harmless to keep (marking something live twice is idempotent) and provides defense-in-depth.
The only theoretical case where #161117 would still matter is if there's a RegionBranchOpInterface op with an implicit block argument that appears ONLY in internal region-to-region transitions (not in the initial parent-to-region entry). I'm not aware of any such op in MLIR today.
There was a problem hiding this comment.
I think it's worth removing the #161117 code here. Let me do this as part of this PR. @joker-eph or @linuxlonelyeagle if you'd prefer me to remove the redundant code in a follow-up PR, let me know,
FranciscoThiesen
force-pushed
the
fix-affine-for-iv-removal
branch
from
d503589 to
b9b6a5b
Compare
…y removed The liveness analysis marks block arguments that are not successor inputs (like induction variables) as live via visitBranchOperand(). However, visitBranchOperand() is only invoked when non-forwarded operands exist. For affine.for with constant bounds: - Lower/upper bounds are encoded in affine maps (not SSA operands) - Step is a constant attribute (not an operand) - Inits are forwarded operands (not non-forwarded) Since there are no non-forwarded operands, visitBranchOperand() is never called, and the IV is never marked live, leading to incorrect removal. This patch fixes the issue by marking non-successor-input block arguments as live during post-processing of the liveness analysis. This handles all RegionBranchOpInterface ops unconditionally. This new approach supersedes the similar logic added in llvm#161117 within visitBranchOperand(), which only handled region-to-region transitions and only ran when visitBranchOperand() was called. That code is now removed as it is redundant. Also merges two separate IR walks into a single traversal for efficiency. Fixes llvm#172610
FranciscoThiesen
force-pushed
the
fix-affine-for-iv-removal
branch
from
b9b6a5b to
f1a4223
Compare
🪟 Windows x64 Test Results
✅ The build succeeded and all tests passed. |
|
Failure seems to be unrelated to my changes... The CI failures on Linux and Windows are both from the same unrelated test (Dialect/XeGPU/propagate-layout-subgroup.mlir). |
It was indeed unrelated, tests are now passing |
ftynse
left a comment
ftynse
left a comment
There was a problem hiding this comment.
There is #169816 that seems to be addressing the same problem. I put a blocker on it for a testing/understanding reason, but generally I prefer that approach. Specifically, here, liveness is set on arguments at the very end of the process and is never propagated. So we will incorrectly mark as dead a result of a side-effect free operation whose operand is such an argument.
Could you folks collaborate with @linuxlonelyeagle to find a good way to test these?
|
Hello @FranciscoThiesen here is a detailed description of the question raised by @ftynse 😉 #172874. |
I agree that #169816 is a better approach. Let me close this down and we can just wait for that. |
5 participants
Summary
The fix in PR #161117 for issue #157934 addressed
scf.forinduction variable deletion by marking IVs as live invisitBranchOperand(). However, this fix doesn't coveraffine.forwith constant bounds becausevisitBranchOperand()is only called for non-forwarded operands.For
affine.forwith constant bounds likeaffine.for %iv = 0 to 1024:Since there are no non-forwarded operands,
visitBranchOperand()is never called, and the IV is never marked live, leading to incorrect removal.The Fix
This patch fixes the issue by proactively marking non-successor-input block arguments (like IVs) as live during
RunLivenessAnalysisinitialization, rather than relying onvisitBranchOperand()being called.The fix walks all
RegionBranchOpInterfaceops and for each region successor from the parent, marks any block arguments that are NOT in the successor inputs as live. This covers IVs in bothscf.forandaffine.for, as well as similar patterns in other region branch ops.Test Plan
affine.forwith constant bounds and unused IVremove-dead-values.mlirtests still passaffine.forloops andaffine.forwithout iter_argsFixes #172610