diff --git a/mlir/lib/Transforms/Utils/RegionUtils.cpp b/mlir/lib/Transforms/Utils/RegionUtils.cpp --- a/mlir/lib/Transforms/Utils/RegionUtils.cpp +++ b/mlir/lib/Transforms/Utils/RegionUtils.cpp @@ -334,37 +334,67 @@ } } -static void propagateLiveness(Operation *op, LiveMap &liveMap) { - // Recurse on any regions the op has. - for (Region ®ion : op->getRegions()) - propagateLiveness(region, liveMap); +static void propagateLiveness(Region ®ion, LiveMap &liveMap) { + if (region.empty()) + return; - // Process terminator operations. - if (op->hasTrait()) - return propagateTerminatorLiveness(op, liveMap); + SmallVector>>> + worklist; + auto addToWorklist = [&](Region ®) { + for (Block &block : reg.getBlocks()) { + SmallVector> opsInBlock; + for (Operation &op : block.getOperations()) + opsInBlock.push_back({&op, op.getNumRegions() != 0}); + worklist.push_back({&block, opsInBlock}); + } + }; + addToWorklist(region); - // Don't reprocess live operations. - if (liveMap.wasProvenLive(op)) - return; + while (!worklist.empty()) { + size_t lastIdx = worklist.size() - 1; + Block *block = worklist[lastIdx].first; + while (!worklist[lastIdx].second.empty()) { + Operation *op = worklist[lastIdx].second.back().first; + auto &processRegion = worklist[lastIdx].second.back().second; + if (processRegion) { + processRegion = false; + // Process any regions the op has if they have not been processed. + for (Region ® : op->getRegions()) + addToWorklist(reg); + break; + } - // Process the op itself. - if (!wouldOpBeTriviallyDead(op)) - return liveMap.setProvedLive(op); + // Pop the op that is about to be processed. + worklist[lastIdx].second.pop_back(); - // If the op isn't intrinsically alive, check it's results. - for (Value value : op->getResults()) - processValue(value, liveMap); -} + // Process terminator operations. + if (op->hasTrait()) { + propagateTerminatorLiveness(op, liveMap); + continue; + } -static void propagateLiveness(Region ®ion, LiveMap &liveMap) { - if (region.empty()) - return; + // Don't reprocess live operations. + if (liveMap.wasProvenLive(op)) + continue; + + // Process the op itself. + if (!wouldOpBeTriviallyDead(op)) { + liveMap.setProvedLive(op); + continue; + } + + // If the op isn't intrinsically alive, check it's results. + for (Value value : op->getResults()) + processValue(value, liveMap); + } + + // Process the regions from ops in this block (newly added when processing + // these ops) first. + if (!worklist[lastIdx].second.empty()) + continue; - for (Block *block : llvm::post_order(®ion.front())) { - // We process block arguments after the ops in the block, to promote - // faster convergence to a fixed point (we try to visit uses before defs). - for (Operation &op : llvm::reverse(block->getOperations())) - propagateLiveness(&op, liveMap); + // Pop the block that is about to be processed. + worklist.pop_back(); // We currently do not remove entry block arguments, so there is no need to // track their liveness.