compiler/src/iree/compiler/Dialect/VM/Analysis/ValueLiveness.cpp - 3p/openxla/iree - Git at Google

 // Copyright 2019 The IREE Authors
 //
 // Licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

 #include "iree/compiler/Dialect/VM/Analysis/ValueLiveness.h"

 #include <algorithm>
 #include <cstring>

 #include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"

 namespace mlir::iree_compiler {

 // static
 LogicalResult ValueLiveness::annotateIR(IREE::VM::FuncOp funcOp) {
   ValueLiveness liveness;
   if (failed(liveness.recalculate(funcOp))) {
     funcOp.emitOpError()
         << "failed to compute value liveness information for function";
     return failure();
   }

   // Build mapping of Operations to values live during them.
   // This is not needed normally so we calculate it only in this debugging case.
   DenseMap<Operation *, llvm::SmallSetVector<Value, 8>> livePerOp;
   for (auto &liveRange : liveness.liveRanges_) {
     auto value = liveRange.getFirst();
     auto &bitVector = liveRange.getSecond();
     for (int opIndex : bitVector.set_bits()) {
       auto *op = liveness.opsInOrder_[opIndex];
       livePerOp[op].insert(value);
     }
   }

   // Block names are their order in the function.
   DenseMap<Block *, int> blockOrdinals;
   for (auto &block : funcOp.getBlocks()) {
     int ordinal = blockOrdinals.size();
     blockOrdinals[std::addressof(block)] = ordinal;
   }

   // Keep asm state to make getting the SSA value names fast.
   OpPrintingFlags printingFlags;
   printingFlags.elideLargeElementsAttrs(1);
   AsmState asmState(funcOp, printingFlags);

   // Gather attributes for each op before we actually add them. We do this so
   // that it's easier to slice out the results for printing without also
   // including the attributes we ourselves are trying to add.
   Builder builder(funcOp.getContext());
   DenseMap<Operation *, NamedAttrList> livenessAttrs(livePerOp.size());
   auto addValuesAttr = [&](Operation &op, StringRef attrName,
                            const llvm::SmallSetVector<Value, 8> &values) {
     SmallVector<StringAttr, 8> valueNames;
     for (auto value : values) {
       std::string str;
       if (auto blockArg = dyn_cast<BlockArgument>(value)) {
         if (blockArg.getOwner()->isEntryBlock()) {
           str = llvm::formatv("%arg{}", blockArg.getArgNumber());
         } else {
           str = llvm::formatv("%bb{}_arg{}", blockOrdinals[blockArg.getOwner()],
                               blockArg.getArgNumber());
         }
       } else {
         llvm::raw_string_ostream os(str);
         value.print(os, asmState);
         str = os.str();
       }

       int equalsIndex = str.find(" =");
       if (equalsIndex != std::string::npos) { // heh
         auto results = str.substr(0, equalsIndex);
         valueNames.push_back(builder.getStringAttr(results));
       } else {
         valueNames.push_back(builder.getStringAttr(str));
       }
     }

     // Sort attributes by name (as SmallSetVector is unordered).
     std::sort(
         valueNames.begin(), valueNames.end(),
         +[](const StringAttr &a, const StringAttr &b) {
           return a.getValue().compare_insensitive(b.getValue()) < 0;
         });
     SmallVector<Attribute, 8> valueNameAttrs;
     for (auto attr : valueNames) {
       valueNameAttrs.push_back(attr);
     }

     livenessAttrs[&op].set(attrName, builder.getArrayAttr(valueNameAttrs));
   };

   for (auto &block : funcOp.getBlocks()) {
     auto &blockLiveness = liveness.blockLiveness_[&block];

     addValuesAttr(block.front(), "block_live_in", blockLiveness.liveIn);
     addValuesAttr(block.front(), "block_live", blockLiveness.live);
     addValuesAttr(block.front(), "block_live_out", blockLiveness.liveOut);
     addValuesAttr(block.front(), "block_defined", blockLiveness.defined);

     // Add per-op live values.
     for (auto &op : block.getOperations()) {
       addValuesAttr(op, "live", livePerOp[&op]);
     }
   }

   // Markup all ops with their attributes.
   for (auto &opAttrs : livenessAttrs) {
     for (auto nameAttr : opAttrs.getSecond().getAttrs()) {
       opAttrs.getFirst()->setAttr(nameAttr.getName(), nameAttr.getValue());
     }
   }

   return success();
 }

 LogicalResult ValueLiveness::recalculate(IREE::VM::FuncOp funcOp) {
   opsInOrder_.clear();
   opOrdering_.clear();
   blockLiveness_.clear();
   liveRanges_.clear();

   calculateOpOrdering(funcOp);
   if (failed(computeLivenessSets(funcOp))) {
     return funcOp.emitError() << "failed to compute liveness sets";
   }
   if (failed(computeLiveIntervals(funcOp))) {
     return funcOp.emitError() << "failed to compute live intervals";
   }

   return success();
 }

 void ValueLiveness::calculateOpOrdering(IREE::VM::FuncOp funcOp) {
   int nextOrdinal = 0;
   for (auto &block : funcOp.getBlocks()) {
     for (auto &op : block.getOperations()) {
       opOrdering_[&op] = nextOrdinal++;
       opsInOrder_.push_back(&op);
     }
   }
 }

 LogicalResult
 ValueLiveness::computeInitialLivenessSets(IREE::VM::FuncOp funcOp) {
   for (auto &block : funcOp.getBlocks()) {
     auto &blockSets = blockLiveness_[&block];

     // Block arguments are defined within the block, technically.
     for (auto blockArg : block.getArguments()) {
       blockSets.defined.insert(blockArg);
     }

     // Add all operands as live uses and all results as definitions.
     for (auto &op : block.getOperations()) {
       for (auto &operand : op.getOpOperands()) {
         blockSets.live.insert(operand.get());
       }
       for (auto result : op.getResults()) {
         blockSets.defined.insert(result);
         for (auto &use : result.getUses()) {
           if (use.getOwner()->getBlock() != &block) {
             // Value escapes this block.
             blockSets.liveOut.insert(result);
           }
         }
       }
     }
   }
   return success();
 }

 LogicalResult ValueLiveness::computeLivenessSets(IREE::VM::FuncOp funcOp) {
   if (failed(computeInitialLivenessSets(funcOp))) {
     return failure();
   }

   llvm::SmallSetVector<Block *, 32> worklist;
   worklist.insert(&funcOp.getBlocks().front());

   // Compute live-in set for each block.
   for (auto &block : funcOp.getBlocks()) {
     auto &blockSets = blockLiveness_[&block];
     blockSets.liveIn = blockSets.live;
     blockSets.liveIn.set_union(blockSets.liveOut);
     blockSets.liveIn.set_subtract(blockSets.defined);

     // If there are live-ins they may need to be propagated to predecessors.
     if (!blockSets.liveIn.empty()) {
       worklist.insert(block.getPredecessors().begin(),
                       block.getPredecessors().end());
     }
   }

   // Propagate liveness until a fixed point is reached.
   while (!worklist.empty()) {
     Block *block = worklist.pop_back_val();
     auto &blockSets = blockLiveness_[block];

     // Compute a new live-out set for the block.
     auto liveOut = blockSets.liveOut;
     for (Block *successor : block->getSuccessors()) {
       liveOut.set_union(blockLiveness_[successor].liveIn);
     }
     blockSets.liveOut = liveOut;

     // Compute the live-in set for the block.
     auto liveIn = liveOut;
     liveIn.set_union(blockSets.live);
     liveIn.set_subtract(blockSets.defined);

     // Propagate the liveness to predecessors if the live-in set changed.
     if (blockSets.liveIn.size() != liveIn.size()) {
       blockSets.liveIn = liveIn;
       worklist.insert(block->getPredecessors().begin(),
                       block->getPredecessors().end());
     }
   }

   return success();
 }

 LogicalResult ValueLiveness::computeLiveIntervals(IREE::VM::FuncOp funcOp) {
   // Adds a live range for |value| from |start| to |end|.
   // Both |start| and |end| must be within the same Block.
   auto addLiveRange = [this](Value value, Operation *start, Operation *end) {
     assert(start->getBlock() == end->getBlock());
     auto &bitVector = liveRanges_[value];
     bitVector.resize(opOrdering_.size());
     bitVector.set(opOrdering_[start], opOrdering_[end] + 1);
   };

   for (auto &block : funcOp.getBlocks()) {
     auto &blockSets = blockLiveness_[&block];

     // Handle values that escape the block.
     for (auto value : blockSets.liveOut) {
       if (blockSets.liveIn.count(value)) {
         // Live in and live out covers the entire block.
         addLiveRange(value, &block.front(), &block.back());
       } else {
         // Live out but not live in implies defined in the block.
         Operation *firstUse =
             value.getDefiningOp() ? value.getDefiningOp() : &block.front();
         addLiveRange(value, firstUse, &block.back());
       }
     }

     // Handle values entering the block and dying within.
     for (auto value : blockSets.liveIn) {
       if (blockSets.liveOut.contains(value)) {
         continue;
       }
       Operation *lastUse = &block.front();
       for (auto &use : value.getUses()) {
         if (use.getOwner()->getBlock() != &block) {
           continue;
         }
         if (lastUse == use.getOwner()) {
           continue;
         }
         if (lastUse->isBeforeInBlock(use.getOwner())) {
           lastUse = use.getOwner();
         }
       }
       addLiveRange(value, &block.front(), lastUse);
     }

     // Handle values defined within the block and not escaping.
     for (auto value : blockSets.defined) {
       if (blockSets.liveOut.contains(value)) {
         continue;
       }
       Operation *firstUse =
           value.getDefiningOp() ? value.getDefiningOp() : &block.front();
       Operation *lastUse = firstUse;
       for (auto &use : value.getUses()) {
         if (use.getOwner()->getBlock() != &block) {
           continue;
         }
         if (lastUse->isBeforeInBlock(use.getOwner())) {
           lastUse = use.getOwner();
         }
       }
       addLiveRange(value, firstUse, lastUse);
     }
   }

   return success();
 }

 ArrayRef<Value> ValueLiveness::getBlockLiveIns(Block *block) {
   auto &blockSets = blockLiveness_[block];
   return blockSets.liveIn.getArrayRef();
 }

 ArrayRef<Value> ValueLiveness::getBlockLiveOuts(Block *block) {
   auto &blockSets = blockLiveness_[block];
   return blockSets.liveOut.getArrayRef();
 }

 bool ValueLiveness::isLastValueUse(Value value, Operation *useOp) {
   auto &blockSets = blockLiveness_[useOp->getBlock()];
   if (blockSets.liveOut.contains(value)) {
     // Value is escapes the block the useOp is in so it is definitely not the
     // last use.
     return false;
   }
   int opOrdinal = opOrdering_[useOp];
   auto &liveRange = liveRanges_[value];
   if (!useOp->hasTrait<OpTrait::IsTerminator>() &&
       liveRange.test(opOrdinal + 1)) {
     // The value is still live within the block after the useOp.
     return false;
   }
   return true;
 }

 bool ValueLiveness::isLastValueUse(Value value, Operation *useOp,
                                    int operandIndex) {
   if (!isLastValueUse(value, useOp)) {
     return false;
   }
   for (auto &operand : llvm::reverse(useOp->getOpOperands())) {
     if (operand.get() == value) {
       // Compare the queried operand index with the last use index.
       return operandIndex >= operand.getOperandNumber();
     }
   }
   assert(false && "value not used by operand");
   return false;
 }

 // Determines if an operand is the last "real" use of a ref value.
 //
 // "Real" uses exclude vm.discard_refs operations, which are cleanup markers
 // rather than actual uses. This distinction is critical for MOVE bit logic:
 // - If the only uses after a branch are discard ops, the branch IS the last
 //   real use and should get MOVE semantics.
 // - Without this, we'd incorrectly retain refs that are about to be discarded.
 //
 // Returns true when all conditions are met:
 // 1. useOp is not a discard op (discards are never "real" uses)
 // 2. value is not in the block's live-out set (doesn't escape)
 // 3. No non-discard uses exist after useOp in the same block
 // 4. operandIndex is >= the last operand using this value (handles multi-use)
 //
 // The multi-use check (condition 4) ensures that for ops like:
 //   vm.call @foo(%ref, %ref)  // same ref used twice
 // Only the LAST operand (index 1) is considered the "last use", so MOVE
 // is only applied once.
 bool ValueLiveness::isLastRealValueUse(Value value, Operation *useOp,
                                        int operandIndex) {
   // Discards are never "real" uses - they just clean up.
   if (isa<IREE::VM::DiscardRefsOp>(useOp)) {
     return false;
   }

   // Check if value escapes block.
   auto &blockSets = blockLiveness_[useOp->getBlock()];
   if (blockSets.liveOut.contains(value)) {
     // Value is in liveOut. For non-terminators, this means the value has uses
     // after this block, so it's not at last use.
     if (!useOp->hasTrait<OpTrait::IsTerminator>()) {
       return false;
     }
     // For branch terminators where the value is forwarded as a successor
     // operand, discards in successors are just cleanup - they don't count as
     // "real" uses since the branch transfers ownership. But for other
     // terminators (like vm.call.yieldable), the value is NOT forwarded -
     // discards in successors ARE real uses of the original value.
     bool valueIsSuccessorOperand = false;
     if (auto branchOp = dyn_cast<BranchOpInterface>(useOp)) {
       for (unsigned i = 0; i < useOp->getNumSuccessors(); ++i) {
         SuccessorOperands succOperands = branchOp.getSuccessorOperands(i);
         for (Value operand : succOperands.getForwardedOperands()) {
           if (operand == value) {
             valueIsSuccessorOperand = true;
             break;
           }
         }
         if (valueIsSuccessorOperand) {
           break;
         }
       }
     }
     // 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();
       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
     // to check if this is the last use within the block.
   }

   // Walk forward to see if any non-discard uses exist after this op.
   for (auto it = ++Block::iterator(useOp); it != useOp->getBlock()->end();
        ++it) {
     for (OpOperand &operand : it->getOpOperands()) {
       if (operand.get() == value && !isa<IREE::VM::DiscardRefsOp>(&*it)) {
         return false; // Real use found after this op
       }
     }
   }

   // Handle same-value multiple operands (only last operand is "last use").
   for (auto &operand : llvm::reverse(useOp->getOpOperands())) {
     if (operand.get() == value) {
       return operandIndex >= operand.getOperandNumber();
     }
   }
   return false;
 }

 } // namespace mlir::iree_compiler
	// Copyright 2019 The IREE Authors
	//
	// Licensed under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

	#include "iree/compiler/Dialect/VM/Analysis/ValueLiveness.h"

	#include <algorithm>
	#include <cstring>

	#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
	#include "llvm/ADT/BitVector.h"
	#include "llvm/ADT/STLExtras.h"
	#include "llvm/ADT/SetVector.h"
	#include "llvm/Support/FormatVariadic.h"
	#include "llvm/Support/raw_ostream.h"
	#include "mlir/IR/AsmState.h"
	#include "mlir/IR/Attributes.h"
	#include "mlir/IR/Builders.h"
	#include "mlir/IR/Value.h"
	#include "mlir/Interfaces/ControlFlowInterfaces.h"

	namespace mlir::iree_compiler {

	// static
	LogicalResult ValueLiveness::annotateIR(IREE::VM::FuncOp funcOp) {
	ValueLiveness liveness;
	if (failed(liveness.recalculate(funcOp))) {
	funcOp.emitOpError()
	<< "failed to compute value liveness information for function";
	return failure();
	}

	// Build mapping of Operations to values live during them.
	// This is not needed normally so we calculate it only in this debugging case.
	DenseMap<Operation *, llvm::SmallSetVector<Value, 8>> livePerOp;
	for (auto &liveRange : liveness.liveRanges_) {
	auto value = liveRange.getFirst();
	auto &bitVector = liveRange.getSecond();
	for (int opIndex : bitVector.set_bits()) {
	auto *op = liveness.opsInOrder_[opIndex];
	livePerOp[op].insert(value);
	}
	}

	// Block names are their order in the function.
	DenseMap<Block *, int> blockOrdinals;
	for (auto &block : funcOp.getBlocks()) {
	int ordinal = blockOrdinals.size();
	blockOrdinals[std::addressof(block)] = ordinal;
	}

	// Keep asm state to make getting the SSA value names fast.
	OpPrintingFlags printingFlags;
	printingFlags.elideLargeElementsAttrs(1);
	AsmState asmState(funcOp, printingFlags);

	// Gather attributes for each op before we actually add them. We do this so
	// that it's easier to slice out the results for printing without also
	// including the attributes we ourselves are trying to add.
	Builder builder(funcOp.getContext());
	DenseMap<Operation *, NamedAttrList> livenessAttrs(livePerOp.size());
	auto addValuesAttr = [&](Operation &op, StringRef attrName,
	const llvm::SmallSetVector<Value, 8> &values) {
	SmallVector<StringAttr, 8> valueNames;
	for (auto value : values) {
	std::string str;
	if (auto blockArg = dyn_cast<BlockArgument>(value)) {
	if (blockArg.getOwner()->isEntryBlock()) {
	str = llvm::formatv("%arg{}", blockArg.getArgNumber());
	} else {
	str = llvm::formatv("%bb{}_arg{}", blockOrdinals[blockArg.getOwner()],
	blockArg.getArgNumber());
	}
	} else {
	llvm::raw_string_ostream os(str);
	value.print(os, asmState);
	str = os.str();
	}

	int equalsIndex = str.find(" =");
	if (equalsIndex != std::string::npos) { // heh
	auto results = str.substr(0, equalsIndex);
	valueNames.push_back(builder.getStringAttr(results));
	} else {
	valueNames.push_back(builder.getStringAttr(str));
	}
	}

	// Sort attributes by name (as SmallSetVector is unordered).
	std::sort(
	valueNames.begin(), valueNames.end(),
	+[](const StringAttr &a, const StringAttr &b) {
	return a.getValue().compare_insensitive(b.getValue()) < 0;
	});
	SmallVector<Attribute, 8> valueNameAttrs;
	for (auto attr : valueNames) {
	valueNameAttrs.push_back(attr);
	}

	livenessAttrs[&op].set(attrName, builder.getArrayAttr(valueNameAttrs));
	};

	for (auto &block : funcOp.getBlocks()) {
	auto &blockLiveness = liveness.blockLiveness_[&block];

	addValuesAttr(block.front(), "block_live_in", blockLiveness.liveIn);
	addValuesAttr(block.front(), "block_live", blockLiveness.live);
	addValuesAttr(block.front(), "block_live_out", blockLiveness.liveOut);
	addValuesAttr(block.front(), "block_defined", blockLiveness.defined);

	// Add per-op live values.
	for (auto &op : block.getOperations()) {
	addValuesAttr(op, "live", livePerOp[&op]);
	}
	}

	// Markup all ops with their attributes.
	for (auto &opAttrs : livenessAttrs) {
	for (auto nameAttr : opAttrs.getSecond().getAttrs()) {
	opAttrs.getFirst()->setAttr(nameAttr.getName(), nameAttr.getValue());
	}
	}

	return success();
	}

	LogicalResult ValueLiveness::recalculate(IREE::VM::FuncOp funcOp) {
	opsInOrder_.clear();
	opOrdering_.clear();
	blockLiveness_.clear();
	liveRanges_.clear();

	calculateOpOrdering(funcOp);
	if (failed(computeLivenessSets(funcOp))) {
	return funcOp.emitError() << "failed to compute liveness sets";
	}
	if (failed(computeLiveIntervals(funcOp))) {
	return funcOp.emitError() << "failed to compute live intervals";
	}

	return success();
	}

	void ValueLiveness::calculateOpOrdering(IREE::VM::FuncOp funcOp) {
	int nextOrdinal = 0;
	for (auto &block : funcOp.getBlocks()) {
	for (auto &op : block.getOperations()) {
	opOrdering_[&op] = nextOrdinal++;
	opsInOrder_.push_back(&op);
	}
	}
	}

	LogicalResult
	ValueLiveness::computeInitialLivenessSets(IREE::VM::FuncOp funcOp) {
	for (auto &block : funcOp.getBlocks()) {
	auto &blockSets = blockLiveness_[&block];

	// Block arguments are defined within the block, technically.
	for (auto blockArg : block.getArguments()) {
	blockSets.defined.insert(blockArg);
	}

	// Add all operands as live uses and all results as definitions.
	for (auto &op : block.getOperations()) {
	for (auto &operand : op.getOpOperands()) {
	blockSets.live.insert(operand.get());
	}
	for (auto result : op.getResults()) {
	blockSets.defined.insert(result);
	for (auto &use : result.getUses()) {
	if (use.getOwner()->getBlock() != &block) {
	// Value escapes this block.
	blockSets.liveOut.insert(result);
	}
	}
	}
	}
	}
	return success();
	}

	LogicalResult ValueLiveness::computeLivenessSets(IREE::VM::FuncOp funcOp) {
	if (failed(computeInitialLivenessSets(funcOp))) {
	return failure();
	}

	llvm::SmallSetVector<Block *, 32> worklist;
	worklist.insert(&funcOp.getBlocks().front());

	// Compute live-in set for each block.
	for (auto &block : funcOp.getBlocks()) {
	auto &blockSets = blockLiveness_[&block];
	blockSets.liveIn = blockSets.live;
	blockSets.liveIn.set_union(blockSets.liveOut);
	blockSets.liveIn.set_subtract(blockSets.defined);

	// If there are live-ins they may need to be propagated to predecessors.
	if (!blockSets.liveIn.empty()) {
	worklist.insert(block.getPredecessors().begin(),
	block.getPredecessors().end());
	}
	}

	// Propagate liveness until a fixed point is reached.
	while (!worklist.empty()) {
	Block *block = worklist.pop_back_val();
	auto &blockSets = blockLiveness_[block];

	// Compute a new live-out set for the block.
	auto liveOut = blockSets.liveOut;
	for (Block *successor : block->getSuccessors()) {
	liveOut.set_union(blockLiveness_[successor].liveIn);
	}
	blockSets.liveOut = liveOut;

	// Compute the live-in set for the block.
	auto liveIn = liveOut;
	liveIn.set_union(blockSets.live);
	liveIn.set_subtract(blockSets.defined);

	// Propagate the liveness to predecessors if the live-in set changed.
	if (blockSets.liveIn.size() != liveIn.size()) {
	blockSets.liveIn = liveIn;
	worklist.insert(block->getPredecessors().begin(),
	block->getPredecessors().end());
	}
	}

	return success();
	}

	LogicalResult ValueLiveness::computeLiveIntervals(IREE::VM::FuncOp funcOp) {
	// Adds a live range for \|value\| from \|start\| to \|end\|.
	// Both \|start\| and \|end\| must be within the same Block.
	auto addLiveRange = [this](Value value, Operation start, Operation end) {
	assert(start->getBlock() == end->getBlock());
	auto &bitVector = liveRanges_[value];
	bitVector.resize(opOrdering_.size());
	bitVector.set(opOrdering_[start], opOrdering_[end] + 1);
	};

	for (auto &block : funcOp.getBlocks()) {
	auto &blockSets = blockLiveness_[&block];

	// Handle values that escape the block.
	for (auto value : blockSets.liveOut) {
	if (blockSets.liveIn.count(value)) {
	// Live in and live out covers the entire block.
	addLiveRange(value, &block.front(), &block.back());
	} else {
	// Live out but not live in implies defined in the block.
	Operation *firstUse =
	value.getDefiningOp() ? value.getDefiningOp() : &block.front();
	addLiveRange(value, firstUse, &block.back());
	}
	}

	// Handle values entering the block and dying within.
	for (auto value : blockSets.liveIn) {
	if (blockSets.liveOut.contains(value)) {
	continue;
	}
	Operation *lastUse = &block.front();
	for (auto &use : value.getUses()) {
	if (use.getOwner()->getBlock() != &block) {
	continue;
	}
	if (lastUse == use.getOwner()) {
	continue;
	}
	if (lastUse->isBeforeInBlock(use.getOwner())) {
	lastUse = use.getOwner();
	}
	}
	addLiveRange(value, &block.front(), lastUse);
	}

	// Handle values defined within the block and not escaping.
	for (auto value : blockSets.defined) {
	if (blockSets.liveOut.contains(value)) {
	continue;
	}
	Operation *firstUse =
	value.getDefiningOp() ? value.getDefiningOp() : &block.front();
	Operation *lastUse = firstUse;
	for (auto &use : value.getUses()) {
	if (use.getOwner()->getBlock() != &block) {
	continue;
	}
	if (lastUse->isBeforeInBlock(use.getOwner())) {
	lastUse = use.getOwner();
	}
	}
	addLiveRange(value, firstUse, lastUse);
	}
	}

	return success();
	}

	ArrayRef<Value> ValueLiveness::getBlockLiveIns(Block *block) {
	auto &blockSets = blockLiveness_[block];
	return blockSets.liveIn.getArrayRef();
	}

	ArrayRef<Value> ValueLiveness::getBlockLiveOuts(Block *block) {
	auto &blockSets = blockLiveness_[block];
	return blockSets.liveOut.getArrayRef();
	}

	bool ValueLiveness::isLastValueUse(Value value, Operation *useOp) {
	auto &blockSets = blockLiveness_[useOp->getBlock()];
	if (blockSets.liveOut.contains(value)) {
	// Value is escapes the block the useOp is in so it is definitely not the
	// last use.
	return false;
	}
	int opOrdinal = opOrdering_[useOp];
	auto &liveRange = liveRanges_[value];
	if (!useOp->hasTrait<OpTrait::IsTerminator>() &&
	liveRange.test(opOrdinal + 1)) {
	// The value is still live within the block after the useOp.
	return false;
	}
	return true;
	}

	bool ValueLiveness::isLastValueUse(Value value, Operation *useOp,
	int operandIndex) {
	if (!isLastValueUse(value, useOp)) {
	return false;
	}
	for (auto &operand : llvm::reverse(useOp->getOpOperands())) {
	if (operand.get() == value) {
	// Compare the queried operand index with the last use index.
	return operandIndex >= operand.getOperandNumber();
	}
	}
	assert(false && "value not used by operand");
	return false;
	}

	// Determines if an operand is the last "real" use of a ref value.
	//
	// "Real" uses exclude vm.discard_refs operations, which are cleanup markers
	// rather than actual uses. This distinction is critical for MOVE bit logic:
	// - If the only uses after a branch are discard ops, the branch IS the last
	// real use and should get MOVE semantics.
	// - Without this, we'd incorrectly retain refs that are about to be discarded.
	//
	// Returns true when all conditions are met:
	// 1. useOp is not a discard op (discards are never "real" uses)
	// 2. value is not in the block's live-out set (doesn't escape)
	// 3. No non-discard uses exist after useOp in the same block
	// 4. operandIndex is >= the last operand using this value (handles multi-use)
	//
	// The multi-use check (condition 4) ensures that for ops like:
	// vm.call @foo(%ref, %ref) // same ref used twice
	// Only the LAST operand (index 1) is considered the "last use", so MOVE
	// is only applied once.
	bool ValueLiveness::isLastRealValueUse(Value value, Operation *useOp,
	int operandIndex) {
	// Discards are never "real" uses - they just clean up.
	if (isa<IREE::VM::DiscardRefsOp>(useOp)) {
	return false;
	}

	// Check if value escapes block.
	auto &blockSets = blockLiveness_[useOp->getBlock()];
	if (blockSets.liveOut.contains(value)) {
	// Value is in liveOut. For non-terminators, this means the value has uses
	// after this block, so it's not at last use.
	if (!useOp->hasTrait<OpTrait::IsTerminator>()) {
	return false;
	}
	// For branch terminators where the value is forwarded as a successor
	// operand, discards in successors are just cleanup - they don't count as
	// "real" uses since the branch transfers ownership. But for other
	// terminators (like vm.call.yieldable), the value is NOT forwarded -
	// discards in successors ARE real uses of the original value.
	bool valueIsSuccessorOperand = false;
	if (auto branchOp = dyn_cast<BranchOpInterface>(useOp)) {
	for (unsigned i = 0; i < useOp->getNumSuccessors(); ++i) {
	SuccessorOperands succOperands = branchOp.getSuccessorOperands(i);
	for (Value operand : succOperands.getForwardedOperands()) {
	if (operand == value) {
	valueIsSuccessorOperand = true;
	break;
	}
	}
	if (valueIsSuccessorOperand) {
	break;
	}
	}
	}
	// 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();
	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
	// to check if this is the last use within the block.
	}

	// Walk forward to see if any non-discard uses exist after this op.
	for (auto it = ++Block::iterator(useOp); it != useOp->getBlock()->end();
	++it) {
	for (OpOperand &operand : it->getOpOperands()) {
	if (operand.get() == value && !isa<IREE::VM::DiscardRefsOp>(&*it)) {
	return false; // Real use found after this op
	}
	}
	}

	// Handle same-value multiple operands (only last operand is "last use").
	for (auto &operand : llvm::reverse(useOp->getOpOperands())) {
	if (operand.get() == value) {
	return operandIndex >= operand.getOperandNumber();
	}
	}
	return false;
	}

	} // namespace mlir::iree_compiler