[llvm-bugs] [Bug 172610] [MLIR][RemoveDeadValues] affine.for induction variable incorrectly removed (related to #157934)

[LLVM Bugs via llvm-bugs]

Issue	172610
Summary	[MLIR][RemoveDeadValues] affine.for induction variable incorrectly removed (related to #157934)
Labels	mlir
Assignees
Reporter	FranciscoThiesen

    ## Description

The fix in PR #161117 for issue #157934 addressed `scf.for` induction variable deletion, but the same bug still occurs with `affine.for`. Running `remove-dead-values` on IR containing an `affine.for` with an unused induction variable causes the pass to incorrectly delete the IV, resulting in a verification failure.

This is a separate issue from #157934 because:
- The root cause is different (no non-forwarded operands trigger `visitBranchOperand`)
- The fix location needs to be different

## Minimal Reproduction

```mlir
module {
  func.func @test_affine_for_iv_removal() -> i32 {
    %c1_i32 = arith.constant 1 : i32
    %c0_i32 = arith.constant 0 : i32
    
    // IV %arg1 is unused, iter_arg %arg2 is used
    %0 = affine.for %arg1 = 0 to 1024 iter_args(%arg2 = %c0_i32) -> (i32) {
      %1 = arith.addi %arg2, %c1_i32 : i32
      affine.yield %1 : i32
    }
    
 return %0 : i32
  }
}
```

**Command:**
```bash
mlir-opt --remove-dead-values test.mlir
```

**Expected:** Pass completes successfully, IR is unchanged (IV must remain even if unused)

**Actual Error:**
```
test.mlir:7:10: error: 'affine.for' op different number of inits and region iter_args: 1 != 0
    %0 = affine.for %arg1 = 0 to 1024 iter_args(%arg2 = %c0_i32) -> (i32) {
         ^
test.mlir:7:10: note: see current operation: 
%2 = "affine.for"(%1) <{lowerBoundMap = affine_map<() -> (0)>, operandSegmentSizes = array<i32: 0, 0, 1>, step = 1 : index, upperBoundMap = affine_map<() -> (1024)>}> ({
^bb0(%arg0: i32):
  %3 = "arith.addi"(%arg0, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
  "affine.yield"(%3) : (i32) -> ()
}) : (i32) -> i32
```

## Root Cause Analysis

The fix in PR #161117 populates `argumentNotOperand` inside `visitBranchOperand()`, which marks non-successor-input block arguments (like IVs) as live. However, `visitBranchOperand()` is only called for non-forwarded operands.

For `affine.for` with constant bounds like `affine.for %iv = 0 to 1024`:
- Lower bound: constant (no operand)
- Upper bound: constant (no operand)  
- Step: constant (no operand)
- Inits: forwarded operands (not non-forwarded)

Since there are no non-forwarded operands, `visitBranchOperand()` is never called, and the IV is never marked live.

In contrast, `scf.for` always has lower/upper/step as operands, so `visitBranchOperand()` gets called and the fix works.

## Proposed Fix

The fix should proactively mark `RegionBranchOpInterface` block arguments that are not successor inputs as live during analysis initialization, rather than relying on `visitBranchOperand()` being called.

Something like:

```cpp
// In RunLivenessAnalysis constructor or similar initialization
op->walk([&](RegionBranchOpInterface regionBranchOp) {
  for (Region &region : regionBranchOp->getRegions()) {
 if (region.empty()) continue;
    Block &entryBlock = region.front();
 
    SmallVector<RegionSuccessor> successors;
 regionBranchOp.getSuccessorRegions(RegionBranchPoint::parent(), successors);
    for (RegionSuccessor &successor : successors) {
      if (successor.getSuccessor() != &region) continue;
      ValueRange inputs = successor.getSuccessorInputs();
      for (BlockArgument arg : entryBlock.getArguments()) {
        if (llvm::find(inputs, arg) == inputs.end()) {
          // This is an IV - mark it live unconditionally
 solver.getOrCreateState<Liveness>(arg)->markLive();
        }
 }
    }
  }
});
```

## Environment

- LLVM main branch (post PR #161117 merge on Oct 21, 2025)
- Also affects downstream projects using MLIR

## Related

- Issue #157934 - Original issue for `scf.for` (closed as fixed)
- PR #161117 - Fix for `scf.for` that doesn't cover `affine.for`

_______________________________________________
llvm-bugs mailing list
llvm-bugs@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-bugs