919 lines · cpp
1//===- RemoveDeadValues.cpp - Remove Dead Values --------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// The goal of this pass is optimization (reducing runtime) by removing10// unnecessary instructions. Unlike other passes that rely on local information11// gathered from patterns to accomplish optimization, this pass uses a full12// analysis of the IR, specifically, liveness analysis, and is thus more13// powerful.14//15// Currently, this pass performs the following optimizations:16// (A) Removes function arguments that are not live,17// (B) Removes function return values that are not live across all callers of18// the function,19// (C) Removes unneccesary operands, results, region arguments, and region20// terminator operands of region branch ops, and,21// (D) Removes simple and region branch ops that have all non-live results and22// don't affect memory in any way,23//24// iff25//26// the IR doesn't have any non-function symbol ops, non-call symbol user ops and27// branch ops.28//29// Here, a "simple op" refers to an op that isn't a symbol op, symbol-user op,30// region branch op, branch op, region branch terminator op, or return-like.31//32//===----------------------------------------------------------------------===//33 34#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"35#include "mlir/Analysis/DataFlow/LivenessAnalysis.h"36#include "mlir/IR/Builders.h"37#include "mlir/IR/BuiltinAttributes.h"38#include "mlir/IR/Dialect.h"39#include "mlir/IR/Operation.h"40#include "mlir/IR/OperationSupport.h"41#include "mlir/IR/SymbolTable.h"42#include "mlir/IR/Value.h"43#include "mlir/IR/ValueRange.h"44#include "mlir/IR/Visitors.h"45#include "mlir/Interfaces/CallInterfaces.h"46#include "mlir/Interfaces/ControlFlowInterfaces.h"47#include "mlir/Interfaces/FunctionInterfaces.h"48#include "mlir/Interfaces/SideEffectInterfaces.h"49#include "mlir/Pass/Pass.h"50#include "mlir/Support/LLVM.h"51#include "mlir/Transforms/FoldUtils.h"52#include "mlir/Transforms/Passes.h"53#include "llvm/ADT/STLExtras.h"54#include "llvm/Support/Debug.h"55#include "llvm/Support/DebugLog.h"56#include <cassert>57#include <cstddef>58#include <memory>59#include <optional>60#include <vector>61 62#define DEBUG_TYPE "remove-dead-values"63 64namespace mlir {65#define GEN_PASS_DEF_REMOVEDEADVALUES66#include "mlir/Transforms/Passes.h.inc"67} // namespace mlir68 69using namespace mlir;70using namespace mlir::dataflow;71 72//===----------------------------------------------------------------------===//73// RemoveDeadValues Pass74//===----------------------------------------------------------------------===//75 76namespace {77 78// Set of structures below to be filled with operations and arguments to erase.79// This is done to separate analysis and tree modification phases,80// otherwise analysis is operating on half-deleted tree which is incorrect.81 82struct FunctionToCleanUp {83 FunctionOpInterface funcOp;84 BitVector nonLiveArgs;85 BitVector nonLiveRets;86};87 88struct OperationToCleanup {89 Operation *op;90 BitVector nonLive;91 Operation *callee =92 nullptr; // Optional: For CallOpInterface ops, stores the callee function93};94 95struct BlockArgsToCleanup {96 Block *b;97 BitVector nonLiveArgs;98};99 100struct SuccessorOperandsToCleanup {101 BranchOpInterface branch;102 unsigned successorIndex;103 BitVector nonLiveOperands;104};105 106struct RDVFinalCleanupList {107 SmallVector<Operation *> operations;108 SmallVector<Value> values;109 SmallVector<FunctionToCleanUp> functions;110 SmallVector<OperationToCleanup> operands;111 SmallVector<OperationToCleanup> results;112 SmallVector<BlockArgsToCleanup> blocks;113 SmallVector<SuccessorOperandsToCleanup> successorOperands;114};115 116// Some helper functions...117 118/// Return true iff at least one value in `values` is live, given the liveness119/// information in `la`.120static bool hasLive(ValueRange values, const DenseSet<Value> &nonLiveSet,121 RunLivenessAnalysis &la) {122 for (Value value : values) {123 if (nonLiveSet.contains(value)) {124 LDBG() << "Value " << value << " is already marked non-live (dead)";125 continue;126 }127 128 const Liveness *liveness = la.getLiveness(value);129 if (!liveness) {130 LDBG() << "Value " << value131 << " has no liveness info, conservatively considered live";132 return true;133 }134 if (liveness->isLive) {135 LDBG() << "Value " << value << " is live according to liveness analysis";136 return true;137 } else {138 LDBG() << "Value " << value << " is dead according to liveness analysis";139 }140 }141 return false;142}143 144/// Return a BitVector of size `values.size()` where its i-th bit is 1 iff the145/// i-th value in `values` is live, given the liveness information in `la`.146static BitVector markLives(ValueRange values, const DenseSet<Value> &nonLiveSet,147 RunLivenessAnalysis &la) {148 BitVector lives(values.size(), true);149 150 for (auto [index, value] : llvm::enumerate(values)) {151 if (nonLiveSet.contains(value)) {152 lives.reset(index);153 LDBG() << "Value " << value154 << " is already marked non-live (dead) at index " << index;155 continue;156 }157 158 const Liveness *liveness = la.getLiveness(value);159 // It is important to note that when `liveness` is null, we can't tell if160 // `value` is live or not. So, the safe option is to consider it live. Also,161 // the execution of this pass might create new SSA values when erasing some162 // of the results of an op and we know that these new values are live163 // (because they weren't erased) and also their liveness is null because164 // liveness analysis ran before their creation.165 if (!liveness) {166 LDBG() << "Value " << value << " at index " << index167 << " has no liveness info, conservatively considered live";168 continue;169 }170 if (!liveness->isLive) {171 lives.reset(index);172 LDBG() << "Value " << value << " at index " << index173 << " is dead according to liveness analysis";174 } else {175 LDBG() << "Value " << value << " at index " << index176 << " is live according to liveness analysis";177 }178 }179 180 return lives;181}182 183/// Collects values marked as "non-live" in the provided range and inserts them184/// into the nonLiveSet. A value is considered "non-live" if the corresponding185/// index in the `nonLive` bit vector is set.186static void collectNonLiveValues(DenseSet<Value> &nonLiveSet, ValueRange range,187 const BitVector &nonLive) {188 for (auto [index, result] : llvm::enumerate(range)) {189 if (!nonLive[index])190 continue;191 nonLiveSet.insert(result);192 LDBG() << "Marking value " << result << " as non-live (dead) at index "193 << index;194 }195}196 197/// Drop the uses of the i-th result of `op` and then erase it iff toErase[i]198/// is 1.199static void dropUsesAndEraseResults(Operation *op, BitVector toErase) {200 assert(op->getNumResults() == toErase.size() &&201 "expected the number of results in `op` and the size of `toErase` to "202 "be the same");203 204 std::vector<Type> newResultTypes;205 for (OpResult result : op->getResults())206 if (!toErase[result.getResultNumber()])207 newResultTypes.push_back(result.getType());208 OpBuilder builder(op);209 builder.setInsertionPointAfter(op);210 OperationState state(op->getLoc(), op->getName().getStringRef(),211 op->getOperands(), newResultTypes, op->getAttrs());212 for (unsigned i = 0, e = op->getNumRegions(); i < e; ++i)213 state.addRegion();214 Operation *newOp = builder.create(state);215 for (const auto &[index, region] : llvm::enumerate(op->getRegions())) {216 Region &newRegion = newOp->getRegion(index);217 // Move all blocks of `region` into `newRegion`.218 Block *temp = new Block();219 newRegion.push_back(temp);220 while (!region.empty())221 region.front().moveBefore(temp);222 temp->erase();223 }224 225 unsigned indexOfNextNewCallOpResultToReplace = 0;226 for (auto [index, result] : llvm::enumerate(op->getResults())) {227 assert(result && "expected result to be non-null");228 if (toErase[index]) {229 result.dropAllUses();230 } else {231 result.replaceAllUsesWith(232 newOp->getResult(indexOfNextNewCallOpResultToReplace++));233 }234 }235 op->erase();236}237 238/// Convert a list of `Operand`s to a list of `OpOperand`s.239static SmallVector<OpOperand *> operandsToOpOperands(OperandRange operands) {240 OpOperand *values = operands.getBase();241 SmallVector<OpOperand *> opOperands;242 for (unsigned i = 0, e = operands.size(); i < e; i++)243 opOperands.push_back(&values[i]);244 return opOperands;245}246 247/// Process a simple operation `op` using the liveness analysis `la`.248/// If the operation has no memory effects and none of its results are live:249/// 1. Add the operation to a list for future removal, and250/// 2. Mark all its results as non-live values251///252/// The operation `op` is assumed to be simple. A simple operation is one that253/// is NOT:254/// - Function-like255/// - Call-like256/// - A region branch operation257/// - A branch operation258/// - A region branch terminator259/// - Return-like260static void processSimpleOp(Operation *op, RunLivenessAnalysis &la,261 DenseSet<Value> &nonLiveSet,262 RDVFinalCleanupList &cl) {263 if (!isMemoryEffectFree(op) || hasLive(op->getResults(), nonLiveSet, la)) {264 LDBG() << "Simple op is not memory effect free or has live results, "265 "preserving it: "266 << OpWithFlags(op, OpPrintingFlags().skipRegions());267 return;268 }269 270 LDBG()271 << "Simple op has all dead results and is memory effect free, scheduling "272 "for removal: "273 << OpWithFlags(op, OpPrintingFlags().skipRegions());274 cl.operations.push_back(op);275 collectNonLiveValues(nonLiveSet, op->getResults(),276 BitVector(op->getNumResults(), true));277}278 279/// Process a function-like operation `funcOp` using the liveness analysis `la`280/// and the IR in `module`. If it is not public or external:281/// (1) Adding its non-live arguments to a list for future removal.282/// (2) Marking their corresponding operands in its callers for removal.283/// (3) Identifying and enqueueing unnecessary terminator operands284/// (return values that are non-live across all callers) for removal.285/// (4) Enqueueing the non-live arguments and return values for removal.286/// (5) Collecting the uses of these return values in its callers for future287/// removal.288/// (6) Marking all its results as non-live values.289static void processFuncOp(FunctionOpInterface funcOp, Operation *module,290 RunLivenessAnalysis &la, DenseSet<Value> &nonLiveSet,291 RDVFinalCleanupList &cl) {292 LDBG() << "Processing function op: "293 << OpWithFlags(funcOp, OpPrintingFlags().skipRegions());294 if (funcOp.isPublic() || funcOp.isExternal()) {295 LDBG() << "Function is public or external, skipping: "296 << funcOp.getOperation()->getName();297 return;298 }299 300 // Get the list of unnecessary (non-live) arguments in `nonLiveArgs`.301 SmallVector<Value> arguments(funcOp.getArguments());302 BitVector nonLiveArgs = markLives(arguments, nonLiveSet, la);303 nonLiveArgs = nonLiveArgs.flip();304 305 // Do (1).306 for (auto [index, arg] : llvm::enumerate(arguments))307 if (arg && nonLiveArgs[index]) {308 cl.values.push_back(arg);309 nonLiveSet.insert(arg);310 }311 312 // Do (2). (Skip creating generic operand cleanup entries for call ops.313 // Call arguments will be removed in the call-site specific segment-aware314 // cleanup, avoiding generic eraseOperands bitvector mechanics.)315 SymbolTable::UseRange uses = *funcOp.getSymbolUses(module);316 for (SymbolTable::SymbolUse use : uses) {317 Operation *callOp = use.getUser();318 assert(isa<CallOpInterface>(callOp) && "expected a call-like user");319 // Push an empty operand cleanup entry so that call-site specific logic in320 // cleanUpDeadVals runs (it keys off CallOpInterface). The BitVector is321 // intentionally all false to avoid generic erasure.322 // Store the funcOp as the callee to avoid expensive symbol lookup later.323 cl.operands.push_back({callOp, BitVector(callOp->getNumOperands(), false),324 funcOp.getOperation()});325 }326 327 // Do (3).328 // Get the list of unnecessary terminator operands (return values that are329 // non-live across all callers) in `nonLiveRets`. There is a very important330 // subtlety here. Unnecessary terminator operands are NOT the operands of the331 // terminator that are non-live. Instead, these are the return values of the332 // callers such that a given return value is non-live across all callers. Such333 // corresponding operands in the terminator could be live. An example to334 // demonstrate this:335 // func.func private @f(%arg0: memref<i32>) -> (i32, i32) {336 // %c0_i32 = arith.constant 0 : i32337 // %0 = arith.addi %c0_i32, %c0_i32 : i32338 // memref.store %0, %arg0[] : memref<i32>339 // return %c0_i32, %0 : i32, i32340 // }341 // func.func @main(%arg0: i32, %arg1: memref<i32>) -> (i32) {342 // %1:2 = call @f(%arg1) : (memref<i32>) -> i32343 // return %1#0 : i32344 // }345 // Here, we can see that %1#1 is never used. It is non-live. Thus, @f doesn't346 // need to return %0. But, %0 is live. And, still, we want to stop it from347 // being returned, in order to optimize our IR. So, this demonstrates how we348 // can make our optimization strong by even removing a live return value (%0),349 // since it forwards only to non-live value(s) (%1#1).350 size_t numReturns = funcOp.getNumResults();351 BitVector nonLiveRets(numReturns, true);352 for (SymbolTable::SymbolUse use : uses) {353 Operation *callOp = use.getUser();354 assert(isa<CallOpInterface>(callOp) && "expected a call-like user");355 BitVector liveCallRets = markLives(callOp->getResults(), nonLiveSet, la);356 nonLiveRets &= liveCallRets.flip();357 }358 359 // Note that in the absence of control flow ops forcing the control to go from360 // the entry (first) block to the other blocks, the control never reaches any361 // block other than the entry block, because every block has a terminator.362 for (Block &block : funcOp.getBlocks()) {363 Operation *returnOp = block.getTerminator();364 if (returnOp && returnOp->getNumOperands() == numReturns)365 cl.operands.push_back({returnOp, nonLiveRets});366 }367 368 // Do (4).369 cl.functions.push_back({funcOp, nonLiveArgs, nonLiveRets});370 371 // Do (5) and (6).372 if (numReturns == 0)373 return;374 for (SymbolTable::SymbolUse use : uses) {375 Operation *callOp = use.getUser();376 assert(isa<CallOpInterface>(callOp) && "expected a call-like user");377 cl.results.push_back({callOp, nonLiveRets});378 collectNonLiveValues(nonLiveSet, callOp->getResults(), nonLiveRets);379 }380}381 382/// Process a region branch operation `regionBranchOp` using the liveness383/// information in `la`. The processing involves two scenarios:384///385/// Scenario 1: If the operation has no memory effects and none of its results386/// are live:387/// (1') Enqueue all its uses for deletion.388/// (2') Enqueue the branch itself for deletion.389///390/// Scenario 2: Otherwise:391/// (1) Collect its unnecessary operands (operands forwarded to unnecessary392/// results or arguments).393/// (2) Process each of its regions.394/// (3) Collect the uses of its unnecessary results (results forwarded from395/// unnecessary operands396/// or terminator operands).397/// (4) Add these results to the deletion list.398///399/// Processing a region includes:400/// (a) Collecting the uses of its unnecessary arguments (arguments forwarded401/// from unnecessary operands402/// or terminator operands).403/// (b) Collecting these unnecessary arguments.404/// (c) Collecting its unnecessary terminator operands (terminator operands405/// forwarded to unnecessary results406/// or arguments).407///408/// Value Flow Note: In this operation, values flow as follows:409/// - From operands and terminator operands (successor operands)410/// - To arguments and results (successor inputs).411static void processRegionBranchOp(RegionBranchOpInterface regionBranchOp,412 RunLivenessAnalysis &la,413 DenseSet<Value> &nonLiveSet,414 RDVFinalCleanupList &cl) {415 LDBG() << "Processing region branch op: "416 << OpWithFlags(regionBranchOp, OpPrintingFlags().skipRegions());417 // Mark live results of `regionBranchOp` in `liveResults`.418 auto markLiveResults = [&](BitVector &liveResults) {419 liveResults = markLives(regionBranchOp->getResults(), nonLiveSet, la);420 };421 422 // Mark live arguments in the regions of `regionBranchOp` in `liveArgs`.423 auto markLiveArgs = [&](DenseMap<Region *, BitVector> &liveArgs) {424 for (Region ®ion : regionBranchOp->getRegions()) {425 if (region.empty())426 continue;427 SmallVector<Value> arguments(region.front().getArguments());428 BitVector regionLiveArgs = markLives(arguments, nonLiveSet, la);429 liveArgs[®ion] = regionLiveArgs;430 }431 };432 433 // Return the successors of `region` if the latter is not null. Else return434 // the successors of `regionBranchOp`.435 auto getSuccessors = [&](RegionBranchPoint point) {436 SmallVector<RegionSuccessor> successors;437 regionBranchOp.getSuccessorRegions(point, successors);438 return successors;439 };440 441 // Return the operands of `terminator` that are forwarded to `successor` if442 // the former is not null. Else return the operands of `regionBranchOp`443 // forwarded to `successor`.444 auto getForwardedOpOperands = [&](const RegionSuccessor &successor,445 Operation *terminator = nullptr) {446 OperandRange operands =447 terminator ? cast<RegionBranchTerminatorOpInterface>(terminator)448 .getSuccessorOperands(successor)449 : regionBranchOp.getEntrySuccessorOperands(successor);450 SmallVector<OpOperand *> opOperands = operandsToOpOperands(operands);451 return opOperands;452 };453 454 // Mark the non-forwarded operands of `regionBranchOp` in455 // `nonForwardedOperands`.456 auto markNonForwardedOperands = [&](BitVector &nonForwardedOperands) {457 nonForwardedOperands.resize(regionBranchOp->getNumOperands(), true);458 for (const RegionSuccessor &successor :459 getSuccessors(RegionBranchPoint::parent())) {460 for (OpOperand *opOperand : getForwardedOpOperands(successor))461 nonForwardedOperands.reset(opOperand->getOperandNumber());462 }463 };464 465 // Mark the non-forwarded terminator operands of the various regions of466 // `regionBranchOp` in `nonForwardedRets`.467 auto markNonForwardedReturnValues =468 [&](DenseMap<Operation *, BitVector> &nonForwardedRets) {469 for (Region ®ion : regionBranchOp->getRegions()) {470 if (region.empty())471 continue;472 // TODO: this isn't correct in face of multiple terminators.473 Operation *terminator = region.front().getTerminator();474 nonForwardedRets[terminator] =475 BitVector(terminator->getNumOperands(), true);476 for (const RegionSuccessor &successor :477 getSuccessors(RegionBranchPoint(478 cast<RegionBranchTerminatorOpInterface>(terminator)))) {479 for (OpOperand *opOperand :480 getForwardedOpOperands(successor, terminator))481 nonForwardedRets[terminator].reset(opOperand->getOperandNumber());482 }483 }484 };485 486 // Update `valuesToKeep` (which is expected to correspond to operands or487 // terminator operands) based on `resultsToKeep` and `argsToKeep`, given488 // `region`. When `valuesToKeep` correspond to operands, `region` is null.489 // Else, `region` is the parent region of the terminator.490 auto updateOperandsOrTerminatorOperandsToKeep =491 [&](BitVector &valuesToKeep, BitVector &resultsToKeep,492 DenseMap<Region *, BitVector> &argsToKeep, Region *region = nullptr) {493 Operation *terminator =494 region ? region->front().getTerminator() : nullptr;495 RegionBranchPoint point =496 terminator497 ? RegionBranchPoint(498 cast<RegionBranchTerminatorOpInterface>(terminator))499 : RegionBranchPoint::parent();500 501 for (const RegionSuccessor &successor : getSuccessors(point)) {502 Region *successorRegion = successor.getSuccessor();503 for (auto [opOperand, input] :504 llvm::zip(getForwardedOpOperands(successor, terminator),505 successor.getSuccessorInputs())) {506 size_t operandNum = opOperand->getOperandNumber();507 bool updateBasedOn =508 successorRegion509 ? argsToKeep[successorRegion]510 [cast<BlockArgument>(input).getArgNumber()]511 : resultsToKeep[cast<OpResult>(input).getResultNumber()];512 valuesToKeep[operandNum] = valuesToKeep[operandNum] | updateBasedOn;513 }514 }515 };516 517 // Recompute `resultsToKeep` and `argsToKeep` based on `operandsToKeep` and518 // `terminatorOperandsToKeep`. Store true in `resultsOrArgsToKeepChanged` if a519 // value is modified, else, false.520 auto recomputeResultsAndArgsToKeep =521 [&](BitVector &resultsToKeep, DenseMap<Region *, BitVector> &argsToKeep,522 BitVector &operandsToKeep,523 DenseMap<Operation *, BitVector> &terminatorOperandsToKeep,524 bool &resultsOrArgsToKeepChanged) {525 resultsOrArgsToKeepChanged = false;526 527 // Recompute `resultsToKeep` and `argsToKeep` based on `operandsToKeep`.528 for (const RegionSuccessor &successor :529 getSuccessors(RegionBranchPoint::parent())) {530 Region *successorRegion = successor.getSuccessor();531 for (auto [opOperand, input] :532 llvm::zip(getForwardedOpOperands(successor),533 successor.getSuccessorInputs())) {534 bool recomputeBasedOn =535 operandsToKeep[opOperand->getOperandNumber()];536 bool toRecompute =537 successorRegion538 ? argsToKeep[successorRegion]539 [cast<BlockArgument>(input).getArgNumber()]540 : resultsToKeep[cast<OpResult>(input).getResultNumber()];541 if (!toRecompute && recomputeBasedOn)542 resultsOrArgsToKeepChanged = true;543 if (successorRegion) {544 argsToKeep[successorRegion][cast<BlockArgument>(input)545 .getArgNumber()] =546 argsToKeep[successorRegion]547 [cast<BlockArgument>(input).getArgNumber()] |548 recomputeBasedOn;549 } else {550 resultsToKeep[cast<OpResult>(input).getResultNumber()] =551 resultsToKeep[cast<OpResult>(input).getResultNumber()] |552 recomputeBasedOn;553 }554 }555 }556 557 // Recompute `resultsToKeep` and `argsToKeep` based on558 // `terminatorOperandsToKeep`.559 for (Region ®ion : regionBranchOp->getRegions()) {560 if (region.empty())561 continue;562 Operation *terminator = region.front().getTerminator();563 for (const RegionSuccessor &successor :564 getSuccessors(RegionBranchPoint(565 cast<RegionBranchTerminatorOpInterface>(terminator)))) {566 Region *successorRegion = successor.getSuccessor();567 for (auto [opOperand, input] :568 llvm::zip(getForwardedOpOperands(successor, terminator),569 successor.getSuccessorInputs())) {570 bool recomputeBasedOn =571 terminatorOperandsToKeep[region.back().getTerminator()]572 [opOperand->getOperandNumber()];573 bool toRecompute =574 successorRegion575 ? argsToKeep[successorRegion]576 [cast<BlockArgument>(input).getArgNumber()]577 : resultsToKeep[cast<OpResult>(input).getResultNumber()];578 if (!toRecompute && recomputeBasedOn)579 resultsOrArgsToKeepChanged = true;580 if (successorRegion) {581 argsToKeep[successorRegion][cast<BlockArgument>(input)582 .getArgNumber()] =583 argsToKeep[successorRegion]584 [cast<BlockArgument>(input).getArgNumber()] |585 recomputeBasedOn;586 } else {587 resultsToKeep[cast<OpResult>(input).getResultNumber()] =588 resultsToKeep[cast<OpResult>(input).getResultNumber()] |589 recomputeBasedOn;590 }591 }592 }593 }594 };595 596 // Mark the values that we want to keep in `resultsToKeep`, `argsToKeep`,597 // `operandsToKeep`, and `terminatorOperandsToKeep`.598 auto markValuesToKeep =599 [&](BitVector &resultsToKeep, DenseMap<Region *, BitVector> &argsToKeep,600 BitVector &operandsToKeep,601 DenseMap<Operation *, BitVector> &terminatorOperandsToKeep) {602 bool resultsOrArgsToKeepChanged = true;603 // We keep updating and recomputing the values until we reach a point604 // where they stop changing.605 while (resultsOrArgsToKeepChanged) {606 // Update the operands that need to be kept.607 updateOperandsOrTerminatorOperandsToKeep(operandsToKeep,608 resultsToKeep, argsToKeep);609 610 // Update the terminator operands that need to be kept.611 for (Region ®ion : regionBranchOp->getRegions()) {612 if (region.empty())613 continue;614 updateOperandsOrTerminatorOperandsToKeep(615 terminatorOperandsToKeep[region.back().getTerminator()],616 resultsToKeep, argsToKeep, ®ion);617 }618 619 // Recompute the results and arguments that need to be kept.620 recomputeResultsAndArgsToKeep(621 resultsToKeep, argsToKeep, operandsToKeep,622 terminatorOperandsToKeep, resultsOrArgsToKeepChanged);623 }624 };625 626 // Scenario 1. This is the only case where the entire `regionBranchOp`627 // is removed. It will not happen in any other scenario. Note that in this628 // case, a non-forwarded operand of `regionBranchOp` could be live/non-live.629 // It could never be live because of this op but its liveness could have been630 // attributed to something else.631 // Do (1') and (2').632 if (isMemoryEffectFree(regionBranchOp.getOperation()) &&633 !hasLive(regionBranchOp->getResults(), nonLiveSet, la)) {634 cl.operations.push_back(regionBranchOp.getOperation());635 return;636 }637 638 // Scenario 2.639 // At this point, we know that every non-forwarded operand of `regionBranchOp`640 // is live.641 642 // Stores the results of `regionBranchOp` that we want to keep.643 BitVector resultsToKeep;644 // Stores the mapping from regions of `regionBranchOp` to their arguments that645 // we want to keep.646 DenseMap<Region *, BitVector> argsToKeep;647 // Stores the operands of `regionBranchOp` that we want to keep.648 BitVector operandsToKeep;649 // Stores the mapping from region terminators in `regionBranchOp` to their650 // operands that we want to keep.651 DenseMap<Operation *, BitVector> terminatorOperandsToKeep;652 653 // Initializing the above variables...654 655 // The live results of `regionBranchOp` definitely need to be kept.656 markLiveResults(resultsToKeep);657 // Similarly, the live arguments of the regions in `regionBranchOp` definitely658 // need to be kept.659 markLiveArgs(argsToKeep);660 // The non-forwarded operands of `regionBranchOp` definitely need to be kept.661 // A live forwarded operand can be removed but no non-forwarded operand can be662 // removed since it "controls" the flow of data in this control flow op.663 markNonForwardedOperands(operandsToKeep);664 // Similarly, the non-forwarded terminator operands of the regions in665 // `regionBranchOp` definitely need to be kept.666 markNonForwardedReturnValues(terminatorOperandsToKeep);667 668 // Mark the values (results, arguments, operands, and terminator operands)669 // that we want to keep.670 markValuesToKeep(resultsToKeep, argsToKeep, operandsToKeep,671 terminatorOperandsToKeep);672 673 // Do (1).674 cl.operands.push_back({regionBranchOp, operandsToKeep.flip()});675 676 // Do (2.a) and (2.b).677 for (Region ®ion : regionBranchOp->getRegions()) {678 if (region.empty())679 continue;680 BitVector argsToRemove = argsToKeep[®ion].flip();681 cl.blocks.push_back({®ion.front(), argsToRemove});682 collectNonLiveValues(nonLiveSet, region.front().getArguments(),683 argsToRemove);684 }685 686 // Do (2.c).687 for (Region ®ion : regionBranchOp->getRegions()) {688 if (region.empty())689 continue;690 Operation *terminator = region.front().getTerminator();691 cl.operands.push_back(692 {terminator, terminatorOperandsToKeep[terminator].flip()});693 }694 695 // Do (3) and (4).696 BitVector resultsToRemove = resultsToKeep.flip();697 collectNonLiveValues(nonLiveSet, regionBranchOp.getOperation()->getResults(),698 resultsToRemove);699 cl.results.push_back({regionBranchOp.getOperation(), resultsToRemove});700}701 702/// Steps to process a `BranchOpInterface` operation:703/// Iterate through each successor block of `branchOp`.704/// (1) For each successor block, gather all operands from all successors.705/// (2) Fetch their associated liveness analysis data and collect for future706/// removal.707/// (3) Identify and collect the dead operands from the successor block708/// as well as their corresponding arguments.709 710static void processBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la,711 DenseSet<Value> &nonLiveSet,712 RDVFinalCleanupList &cl) {713 LDBG() << "Processing branch op: " << *branchOp;714 unsigned numSuccessors = branchOp->getNumSuccessors();715 716 for (unsigned succIdx = 0; succIdx < numSuccessors; ++succIdx) {717 Block *successorBlock = branchOp->getSuccessor(succIdx);718 719 // Do (1)720 SuccessorOperands successorOperands =721 branchOp.getSuccessorOperands(succIdx);722 SmallVector<Value> operandValues;723 for (unsigned operandIdx = 0; operandIdx < successorOperands.size();724 ++operandIdx) {725 operandValues.push_back(successorOperands[operandIdx]);726 }727 728 // Do (2)729 BitVector successorNonLive =730 markLives(operandValues, nonLiveSet, la).flip();731 collectNonLiveValues(nonLiveSet, successorBlock->getArguments(),732 successorNonLive);733 734 // Do (3)735 cl.blocks.push_back({successorBlock, successorNonLive});736 cl.successorOperands.push_back({branchOp, succIdx, successorNonLive});737 }738}739 740/// Removes dead values collected in RDVFinalCleanupList.741/// To be run once when all dead values have been collected.742static void cleanUpDeadVals(RDVFinalCleanupList &list) {743 LDBG() << "Starting cleanup of dead values...";744 745 // 1. Blocks, We must remove the block arguments and successor operands before746 // deleting the operation, as they may reside in the region operation.747 LDBG() << "Cleaning up " << list.blocks.size() << " block argument lists";748 for (auto &b : list.blocks) {749 // blocks that are accessed via multiple codepaths processed once750 if (b.b->getNumArguments() != b.nonLiveArgs.size())751 continue;752 LDBG() << "Erasing " << b.nonLiveArgs.count()753 << " non-live arguments from block: " << b.b;754 // it iterates backwards because erase invalidates all successor indexes755 for (int i = b.nonLiveArgs.size() - 1; i >= 0; --i) {756 if (!b.nonLiveArgs[i])757 continue;758 LDBG() << " Erasing block argument " << i << ": " << b.b->getArgument(i);759 b.b->getArgument(i).dropAllUses();760 b.b->eraseArgument(i);761 }762 }763 764 // 2. Successor Operands765 LDBG() << "Cleaning up " << list.successorOperands.size()766 << " successor operand lists";767 for (auto &op : list.successorOperands) {768 SuccessorOperands successorOperands =769 op.branch.getSuccessorOperands(op.successorIndex);770 // blocks that are accessed via multiple codepaths processed once771 if (successorOperands.size() != op.nonLiveOperands.size())772 continue;773 LDBG() << "Erasing " << op.nonLiveOperands.count()774 << " non-live successor operands from successor "775 << op.successorIndex << " of branch: "776 << OpWithFlags(op.branch, OpPrintingFlags().skipRegions());777 // it iterates backwards because erase invalidates all successor indexes778 for (int i = successorOperands.size() - 1; i >= 0; --i) {779 if (!op.nonLiveOperands[i])780 continue;781 LDBG() << " Erasing successor operand " << i << ": "782 << successorOperands[i];783 successorOperands.erase(i);784 }785 }786 787 // 3. Operations788 LDBG() << "Cleaning up " << list.operations.size() << " operations";789 for (auto &op : list.operations) {790 LDBG() << "Erasing operation: "791 << OpWithFlags(op, OpPrintingFlags().skipRegions());792 op->dropAllUses();793 op->erase();794 }795 796 // 4. Values797 LDBG() << "Cleaning up " << list.values.size() << " values";798 for (auto &v : list.values) {799 LDBG() << "Dropping all uses of value: " << v;800 v.dropAllUses();801 }802 803 // 5. Functions804 LDBG() << "Cleaning up " << list.functions.size() << " functions";805 // Record which function arguments were erased so we can shrink call-site806 // argument segments for CallOpInterface operations (e.g. ops using807 // AttrSizedOperandSegments) in the next phase.808 DenseMap<Operation *, BitVector> erasedFuncArgs;809 for (auto &f : list.functions) {810 LDBG() << "Cleaning up function: " << f.funcOp.getOperation()->getName();811 LDBG() << " Erasing " << f.nonLiveArgs.count() << " non-live arguments";812 LDBG() << " Erasing " << f.nonLiveRets.count()813 << " non-live return values";814 // Some functions may not allow erasing arguments or results. These calls815 // return failure in such cases without modifying the function, so it's okay816 // to proceed.817 if (succeeded(f.funcOp.eraseArguments(f.nonLiveArgs))) {818 // Record only if we actually erased something.819 if (f.nonLiveArgs.any())820 erasedFuncArgs.try_emplace(f.funcOp.getOperation(), f.nonLiveArgs);821 }822 (void)f.funcOp.eraseResults(f.nonLiveRets);823 }824 825 // 6. Operands826 LDBG() << "Cleaning up " << list.operands.size() << " operand lists";827 for (OperationToCleanup &o : list.operands) {828 // Handle call-specific cleanup only when we have a cached callee reference.829 // This avoids expensive symbol lookup and is defensive against future830 // changes.831 bool handledAsCall = false;832 if (o.callee && isa<CallOpInterface>(o.op)) {833 auto call = cast<CallOpInterface>(o.op);834 auto it = erasedFuncArgs.find(o.callee);835 if (it != erasedFuncArgs.end()) {836 const BitVector &deadArgIdxs = it->second;837 MutableOperandRange args = call.getArgOperandsMutable();838 // First, erase the call arguments corresponding to erased callee839 // args. We iterate backwards to preserve indices.840 for (unsigned argIdx : llvm::reverse(deadArgIdxs.set_bits()))841 args.erase(argIdx);842 // If this operand cleanup entry also has a generic nonLive bitvector,843 // clear bits for call arguments we already erased above to avoid844 // double-erasing (which could impact other segments of ops with845 // AttrSizedOperandSegments).846 if (o.nonLive.any()) {847 // Map the argument logical index to the operand number(s) recorded.848 int operandOffset = call.getArgOperands().getBeginOperandIndex();849 for (int argIdx : deadArgIdxs.set_bits()) {850 int operandNumber = operandOffset + argIdx;851 if (operandNumber < static_cast<int>(o.nonLive.size()))852 o.nonLive.reset(operandNumber);853 }854 }855 handledAsCall = true;856 }857 }858 // Perform generic operand erasure for:859 // - Non-call operations860 // - Call operations without cached callee (where handledAsCall is false)861 // But skip call operations that were already handled via segment-aware path862 if (!handledAsCall && o.nonLive.any()) {863 o.op->eraseOperands(o.nonLive);864 }865 }866 867 // 7. Results868 LDBG() << "Cleaning up " << list.results.size() << " result lists";869 for (auto &r : list.results) {870 LDBG() << "Erasing " << r.nonLive.count()871 << " non-live results from operation: "872 << OpWithFlags(r.op, OpPrintingFlags().skipRegions());873 dropUsesAndEraseResults(r.op, r.nonLive);874 }875 LDBG() << "Finished cleanup of dead values";876}877 878struct RemoveDeadValues : public impl::RemoveDeadValuesBase<RemoveDeadValues> {879 void runOnOperation() override;880};881} // namespace882 883void RemoveDeadValues::runOnOperation() {884 auto &la = getAnalysis<RunLivenessAnalysis>();885 Operation *module = getOperation();886 887 // Tracks values eligible for erasure - complements liveness analysis to888 // identify "droppable" values.889 DenseSet<Value> deadVals;890 891 // Maintains a list of Ops, values, branches, etc., slated for cleanup at the892 // end of this pass.893 RDVFinalCleanupList finalCleanupList;894 895 module->walk([&](Operation *op) {896 if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {897 processFuncOp(funcOp, module, la, deadVals, finalCleanupList);898 } else if (auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op)) {899 processRegionBranchOp(regionBranchOp, la, deadVals, finalCleanupList);900 } else if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {901 processBranchOp(branchOp, la, deadVals, finalCleanupList);902 } else if (op->hasTrait<::mlir::OpTrait::IsTerminator>()) {903 // Nothing to do here because this is a terminator op and it should be904 // honored with respect to its parent905 } else if (isa<CallOpInterface>(op)) {906 // Nothing to do because this op is associated with a function op and gets907 // cleaned when the latter is cleaned.908 } else {909 processSimpleOp(op, la, deadVals, finalCleanupList);910 }911 });912 913 cleanUpDeadVals(finalCleanupList);914}915 916std::unique_ptr<Pass> mlir::createRemoveDeadValuesPass() {917 return std::make_unique<RemoveDeadValues>();918}919