diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp index d705d8d4c7819..b66072fd8c6ad 100644 --- a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp +++ b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp @@ -165,6 +165,28 @@ void LivenessAnalysis::visitBranchOperand(OpOperand &operand) { blocks.push_back(&block); } } + + // In the block of the successor block argument of RegionBranchOpInterface, + // there may be arguments of RegionBranchOpInterface, such as the IV of + // scf.forOp. Explicitly set this argument to live. + auto regionBranchOp = cast(op); + for (size_t i = 0, e = op->getNumRegions(); i < e; ++i) { + SmallVector successors; + regionBranchOp.getSuccessorRegions(op->getRegion(i), successors); + for (RegionSuccessor successor : successors) { + if (successor.isParent()) + continue; + auto arguments = successor.getSuccessor()->getArguments(); + ValueRange regionInputs = successor.getSuccessorInputs(); + for (auto argument : arguments) { + if (llvm::find(regionInputs, argument) == regionInputs.end()) { + (void)getLatticeElement(argument)->markLive(); + LDBG() << "Marking RegionBranchOp's success argument live: " + << argument; + } + } + } + } } else if (isa(op)) { // We cannot track all successor blocks of the branch operation(More // specifically, it's the successor's successor). Additionally, different @@ -309,6 +331,7 @@ RunLivenessAnalysis::RunLivenessAnalysis(Operation *op) { << " has no liveness info (unreachable), mark dead"; solver.getOrCreateState(result.value()); } + for (auto ®ion : op->getRegions()) { for (auto &block : region) { for (auto blockArg : llvm::enumerate(block.getArguments())) { diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir index 56449469dc29f..e7304505c809e 100644 --- a/mlir/test/Transforms/remove-dead-values.mlir +++ b/mlir/test/Transforms/remove-dead-values.mlir @@ -649,3 +649,28 @@ 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, the loop has results. + +// CHECK-LABEL: func @dead_value_loop_ivs +func.func @dead_value_loop_ivs_has_result(%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 +} + +// ----- + +// This test verifies that the induction variables in loops are not deleted, the loop has no results. + +// CHECK-LABEL: func @dead_value_loop_ivs_no_result +func.func @dead_value_loop_ivs_no_result(%lb: index, %ub: index, %step: index, %input: memref, %value: f32, %pos: index) { + scf.for %iv = %lb to %ub step %step { + memref.store %value, %input[%pos] : memref + } + return +}