[MLIR] Removing dead values for branches

mlir/lib/Transforms/RemoveDeadValues.cpp

@@ -165,6 +165,17 @@ static SmallVector<OpOperand *> operandsToOpOperands(OperandRange operands) {
   return opOperands;
 }
 
+// Check if any of the operations implements BranchOpInterface
+template <typename UserRange>
+static bool anyBranchUsers(const UserRange &users) {
+  for (auto user : users) {
+    if (auto subBranchOp = dyn_cast<BranchOpInterface>(user)) {
+      return true;
+    }
+  }
+  return false;
+}
+
 /// Clean a simple op `op`, given the liveness analysis information in `la`.
 /// Here, cleaning means:
 ///   (1) Dropping all its uses, AND
@@ -175,7 +186,8 @@ static SmallVector<OpOperand *> operandsToOpOperands(OperandRange operands) {
 /// symbol op, a symbol-user op, a region branch op, a branch op, a region
 /// branch terminator op, or return-like.
 static void cleanSimpleOp(Operation *op, RunLivenessAnalysis &la) {
-  if (!isMemoryEffectFree(op) || hasLive(op->getResults(), la))
+  if (!isMemoryEffectFree(op) || hasLive(op->getResults(), la) ||
+      anyBranchUsers(op->getUsers()))
     return;
 
   op->dropAllUses();
@@ -563,6 +575,48 @@ static void cleanRegionBranchOp(RegionBranchOpInterface regionBranchOp,
   dropUsesAndEraseResults(regionBranchOp.getOperation(), resultsToKeep.flip());
 }
 
+// 1. Iterate over each successor block of the given BranchOpInterface
+// operation.
+// 2. For each successor block:
+//    a. Retrieve the operands passed to the successor.
+//    b. Use the provided liveness analysis (`RunLivenessAnalysis`) to determine
+//    which
+//       operands are live in the successor block.
+//    c. Mark each operand as live or dead based on the analysis.
+// 3. Remove dead operands from the branch operation and arguments accordingly
+
+static void cleanBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la) {
+  unsigned numSuccessors = branchOp->getNumSuccessors();
+
+  // Do (1)
+  for (unsigned succIdx = 0; succIdx < numSuccessors; ++succIdx) {
+    Block *successorBlock = branchOp->getSuccessor(succIdx);
+
+    // Do (2)
+    SuccessorOperands successorOperands =
+        branchOp.getSuccessorOperands(succIdx);
+    SmallVector<Value> operandValues;
+    for (unsigned operandIdx = 0; operandIdx < successorOperands.size();
+         ++operandIdx) {
+      operandValues.push_back(successorOperands[operandIdx]);
+    }
+
+    BitVector successorLiveOperands = markLives(operandValues, la);
+
+    // Do (3)
+    for (int argIdx = successorLiveOperands.size() - 1; argIdx >= 0; --argIdx) {
+      if (!successorLiveOperands[argIdx]) {
+        if (anyBranchUsers(successorBlock->getArgument(argIdx).getUsers())) {
+          continue;
+        }
+
+        successorOperands.erase(argIdx);
+        successorBlock->eraseArgument(argIdx);
+      }
+    }
+  }
+}
+
 struct RemoveDeadValues : public impl::RemoveDeadValuesBase<RemoveDeadValues> {
   void runOnOperation() override;
 };
@@ -572,26 +626,13 @@ void RemoveDeadValues::runOnOperation() {
   auto &la = getAnalysis<RunLivenessAnalysis>();
   Operation *module = getOperation();
 
-  // The removal of non-live values is performed iff there are no branch ops,
-  // and all symbol user ops present in the IR are call-like.
-  WalkResult acceptableIR = module->walk([&](Operation *op) {
-    if (op == module)
-      return WalkResult::advance();
-    if (isa<BranchOpInterface>(op)) {
-      op->emitError() << "cannot optimize an IR with branch ops\n";
-      return WalkResult::interrupt();
-    }
-    return WalkResult::advance();
-  });
-
-  if (acceptableIR.wasInterrupted())
-    return signalPassFailure();
-
   module->walk([&](Operation *op) {
     if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
       cleanFuncOp(funcOp, module, la);
     } else if (auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op)) {
       cleanRegionBranchOp(regionBranchOp, la);
+    } else if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {
+      cleanBranchOp(branchOp, la);
     } else if (op->hasTrait<::mlir::OpTrait::IsTerminator>()) {
       // Nothing to do here because this is a terminator op and it should be
       // honored with respect to its parent

mlir/test/Transforms/remove-dead-values.mlir

@@ -28,22 +28,59 @@ module @named_module_acceptable {
 
 // -----
 
-// The IR remains untouched because of the presence of a branch op `cf.cond_br`.
+// The IR is optimized regardless of the presence of a branch op `cf.cond_br`.
 //
-func.func @dont_touch_unacceptable_ir_has_cleanable_simple_op_with_branch_op(%arg0: i1) {
+func.func @acceptable_ir_has_cleanable_simple_op_with_branch_op(%arg0: i1) {
   %non_live = arith.constant 0 : i32
-  // expected-error @+1 {{cannot optimize an IR with branch ops}}
+  // CHECK-NOT: non_live
   cf.cond_br %arg0, ^bb1(%non_live : i32), ^bb2(%non_live : i32)
 ^bb1(%non_live_0 : i32):
+  // CHECK-NOT: non_live_0
   cf.br ^bb3
 ^bb2(%non_live_1 : i32):
+  // CHECK-NOT: non_live_1
   cf.br ^bb3
 ^bb3:
   return
 }
 
 // -----
 
+// Arguments of unconditional branch op `cf.br` are properly removed.
+//
+func.func @acceptable_ir_has_cleanable_simple_op_with_unconditional_branch_op(%arg0: i1) {
+  %non_live = arith.constant 0 : i32
+  // CHECK-NOT: non_live
+  cf.br ^bb1(%non_live : i32)
+^bb1(%non_live_1 : i32):
+  // CHECK-NOT: non_live_1
+  cf.br ^bb3(%non_live_1 : i32)
+  // CHECK-NOT: non_live_2
+^bb3(%non_live_2 : i32):
+  return
+}
+
+// -----
+
+// Checking that iter_args are properly handled
+//
+func.func @cleanable_loop_iter_args_value(%arg0: index) -> index {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+  %non_live = arith.constant 0 : index
+  // CHECK-NOT: non_live
+  %result, %result_non_live = scf.for %i = %c0 to %c10 step %c1 iter_args(%live_arg = %arg0, %non_live_arg = %non_live) -> (index, index) {
+    %new_live = arith.addi %live_arg, %i : index
+    // CHECK: scf.for %[[ARG_0:.*]] = %c0 to %c10 step %c1 iter_args(%[[ARG_1:.*]] = %arg0)
+    scf.yield %new_live, %non_live_arg : index, index
+  }
+  // CHECK-NOT: result_non_live
+  return %result : index
+}
+
+// -----
+
 // Note that this cleanup cannot be done by the `canonicalize` pass.
 //
 // CHECK-LABEL: func.func private @clean_func_op_remove_argument_and_return_value() {