[VM] Fix ref leak from incorrect MOVE bit on branch block args

compiler/src/iree/compiler/Dialect/VM/Analysis/ValueLiveness.cpp

@@ -397,17 +397,42 @@ bool ValueLiveness::isLastRealValueUse(Value value, Operation *useOp,
       }
     }
     // Check if the value escapes to any successor blocks.
+    // Only check uses in immediate successor blocks. Uses in other blocks
+    // (e.g. predecessors or the definition block) are on prior control-flow
+    // edges and don't affect whether MOVE is safe at this branch point.
+    SmallPtrSet<Block *, 4> successorBlocks;
+    for (unsigned i = 0; i < useOp->getNumSuccessors(); ++i) {
+      Block *succBlock = useOp->getSuccessor(i);
+      successorBlocks.insert(succBlock);
+      // If the value flows THROUGH a successor (both liveIn and liveOut),
+      // it reaches non-immediate successors we can't enumerate here.
+      // Conservatively prevent MOVE. This is precise for current VM pipeline
+      // invariants: MaterializeRefDiscards places exactly one discard per
+      // path at value death points, so liveIn && liveOut implies real
+      // (non-discard) uses downstream.
+      BlockSets &succSets = blockLiveness_[succBlock];
+      if (succSets.liveIn.contains(value) && succSets.liveOut.contains(value)) {
+        return false;
+      }
+    }
     for (auto &use : value.getUses()) {
       Operation *userOp = use.getOwner();
-      if (userOp->getBlock() != useOp->getBlock()) {
-        // Use is in another block. If it's a discard AND the value was
-        // forwarded as a successor operand, skip it (ownership transferred).
-        // Otherwise, it's a real use and the value escapes.
-        bool isDiscardOfForwardedValue =
-            isa<IREE::VM::DiscardRefsOp>(userOp) && valueIsSuccessorOperand;
-        if (!isDiscardOfForwardedValue) {
-          return false;
-        }
+      Block *userBlock = userOp->getBlock();
+      if (userBlock == useOp->getBlock()) {
+        continue;
+      }
+      // Skip uses in non-successor blocks (e.g. predecessor or definition
+      // blocks). These are on prior control-flow edges, not forward ones.
+      if (!successorBlocks.contains(userBlock)) {
+        continue;
+      }
+      // Use is in a successor block. If it's a discard AND the value was
+      // forwarded as a successor operand, skip it (ownership transferred).
+      // Otherwise, it's a real use and the value escapes.
+      bool isDiscardOfForwardedValue =
+          isa<IREE::VM::DiscardRefsOp>(userOp) && valueIsSuccessorOperand;
+      if (!isDiscardOfForwardedValue) {
+        return false;
       }
     }
     // All uses in other blocks were discards of forwarded values. Fall through

compiler/src/iree/compiler/Dialect/VM/Analysis/test/register_allocation_refs.mlir

@@ -923,3 +923,57 @@ vm.module @module_scratch_register_swap {
     vm.return
   }
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Regression Test: Ref Forwarded via Block Arg with Predecessor Use
+//===----------------------------------------------------------------------===//
+
+// When a ref is defined and used (non-discard) in a predecessor block, then
+// forwarded as a block arg via cond_br to a merge block, with a discard on
+// the other branch, the MOVE bit must be set on the block arg operand.
+//
+// Without the fix, isLastRealValueUse checked uses in ALL other blocks
+// (including predecessors) and found a "past" use, returning false and
+// preventing MOVE. This caused a ref leak: retain incremented the ref
+// count, but no discard existed on the taken path to release it.
+
+// CHECK-LABEL: @module_predecessor_use_blockarg_move
+vm.module @module_predecessor_use_blockarg_move {
+
+  vm.import private @produce() -> !vm.buffer
+  vm.import private @use_buffer(%buf : !vm.buffer)
+  vm.import private @check_condition(%buf : !vm.buffer) -> i32
+
+  // CHECK-LABEL: @predecessor_use_blockarg_forward
+  // Ref is used in the predecessor block, then forwarded as a block arg
+  // via cond_br after a branch boundary. Mimics the while-loop ASAN leak.
+  // MOVE must be set on the block arg operand.
+  vm.func @predecessor_use_blockarg_forward() {
+    %ref = vm.call @produce() : () -> !vm.buffer
+    // Use %ref before the branch boundary.
+    // CHECK: vm.call @use_buffer
+    // CHECK-SAME: operand_registers = ["r0"]
+    vm.call @use_buffer(%ref) : (!vm.buffer) -> ()
+    %status = vm.call @check_condition(%ref) : (!vm.buffer) -> i32
+    vm.br ^post_yield(%status : i32)
+
+  ^post_yield(%result: i32):
+    // Forward %ref as block arg on success, discard on error.
+    // %ref is defined in the entry block (predecessor) and used there too.
+    // MOVE must be set on the success branch's block arg operand.
+    // CHECK: vm.cond_br
+    // CHECK-SAME: operand_registers = ["i{{[0-9]+}}", "R0"]
+    vm.cond_br %result, ^error, ^success(%ref : !vm.buffer)
+
+  ^error:
+    vm.discard.refs %ref : !vm.buffer
+    vm.return
+
+  ^success(%buf: !vm.buffer):
+    vm.call @use_buffer(%buf) : (!vm.buffer) -> ()
+    vm.discard.refs %buf : !vm.buffer
+    vm.return
+  }
+}