173 lines · cpp
1//===- OptimizeAllocationLiveness.cpp - impl. optimize allocation liveness pass2//-===//3//4// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.5// See https://llvm.org/LICENSE.txt for license information.6// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception7//8//===----------------------------------------------------------------------===//9//10// This file implements a pass for optimizing allocation liveness.11// The pass moves the deallocation operation after the last user of the12// allocated buffer.13//===----------------------------------------------------------------------===//14 15#include "mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h"16#include "mlir/Dialect/Bufferization/Transforms/Passes.h"17#include "mlir/Dialect/Func/IR/FuncOps.h"18#include "mlir/IR/Operation.h"19#include "mlir/IR/Value.h"20#include "mlir/Interfaces/SideEffectInterfaces.h"21#include "llvm/Support/DebugLog.h"22 23#define DEBUG_TYPE "optimize-allocation-liveness"24 25namespace mlir {26namespace bufferization {27#define GEN_PASS_DEF_OPTIMIZEALLOCATIONLIVENESSPASS28#include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"29} // namespace bufferization30} // namespace mlir31 32using namespace mlir;33 34namespace {35 36//===----------------------------------------------------------------------===//37// Helper functions38//===----------------------------------------------------------------------===//39 40/// Return true if `a` happens before `b`, i.e., `a` or one of its ancestors41/// properly dominates `b` and `b` is not inside `a`.42static bool happensBefore(Operation *a, Operation *b) {43 do {44 if (a->isProperAncestor(b))45 return false;46 if (Operation *bAncestor = a->getBlock()->findAncestorOpInBlock(*b)) {47 return a->isBeforeInBlock(bAncestor);48 }49 } while ((a = a->getParentOp()));50 return false;51}52 53/// This method searches for a user of value that is a dealloc operation.54/// If multiple users with free effect are found, return nullptr.55Operation *findUserWithFreeSideEffect(Value value) {56 Operation *freeOpUser = nullptr;57 for (Operation *user : value.getUsers()) {58 if (MemoryEffectOpInterface memEffectOp =59 dyn_cast<MemoryEffectOpInterface>(user)) {60 SmallVector<MemoryEffects::EffectInstance, 2> effects;61 memEffectOp.getEffects(effects);62 63 for (const auto &effect : effects) {64 if (isa<MemoryEffects::Free>(effect.getEffect())) {65 if (freeOpUser) {66 LDBG() << "Multiple users with free effect found: " << *freeOpUser67 << " and " << *user;68 return nullptr;69 }70 freeOpUser = user;71 }72 }73 }74 }75 return freeOpUser;76}77 78/// Checks if the given op allocates memory.79static bool hasMemoryAllocEffect(MemoryEffectOpInterface memEffectOp) {80 SmallVector<MemoryEffects::EffectInstance, 2> effects;81 memEffectOp.getEffects(effects);82 for (const auto &effect : effects) {83 if (isa<MemoryEffects::Allocate>(effect.getEffect())) {84 return true;85 }86 }87 return false;88}89 90/// Extracts OpResult's with Allocate effects from given op91static SmallVector<OpResult>92collectAllocations(MemoryEffectOpInterface allocOp) {93 SmallVector<MemoryEffects::EffectInstance> effects;94 allocOp.getEffects(effects);95 SmallVector<OpResult> allocResults;96 for (const MemoryEffects::EffectInstance &it : effects)97 if (isa<MemoryEffects::Allocate>(it.getEffect()))98 if (auto val = it.getValue(); val && val.getDefiningOp() == allocOp)99 allocResults.push_back(cast<OpResult>(val));100 return allocResults;101}102 103struct OptimizeAllocationLiveness104 : public bufferization::impl::OptimizeAllocationLivenessPassBase<105 OptimizeAllocationLiveness> {106public:107 OptimizeAllocationLiveness() = default;108 109 void runOnOperation() override {110 func::FuncOp func = getOperation();111 112 if (func.isExternal())113 return;114 115 BufferViewFlowAnalysis analysis = BufferViewFlowAnalysis(func);116 117 func.walk([&](MemoryEffectOpInterface memEffectOp) -> WalkResult {118 if (!hasMemoryAllocEffect(memEffectOp))119 return WalkResult::advance();120 121 auto allocOp = memEffectOp;122 LDBG() << "Checking alloc op: " << allocOp;123 124 SmallVector<OpResult> allocationResults = collectAllocations(allocOp);125 // Multiple allocations from a single op are not considered here yet.126 if (allocationResults.size() != 1)127 return WalkResult::advance();128 129 OpResult allocResult = allocationResults[0];130 LDBG() << "On allocation result: " << allocResult;131 132 auto *deallocOp = findUserWithFreeSideEffect(allocResult);133 if (!deallocOp || (deallocOp->getBlock() != allocOp->getBlock())) {134 // The pass handles allocations that have a single dealloc op in the135 // same block. We also should not hoist the dealloc op out of136 // conditionals.137 return WalkResult::advance();138 }139 140 Operation *lastUser = nullptr;141 const BufferViewFlowAnalysis::ValueSetT &deps =142 analysis.resolve(allocResult);143 for (auto dep : llvm::make_early_inc_range(deps)) {144 for (auto *user : dep.getUsers()) {145 // We are looking for a non dealloc op user.146 // check if user is the dealloc op itself.147 if (user == deallocOp)148 continue;149 150 // find the ancestor of user that is in the same block as the allocOp.151 auto *topUser = allocOp->getBlock()->findAncestorOpInBlock(*user);152 if (!lastUser || happensBefore(lastUser, topUser)) {153 lastUser = topUser;154 }155 }156 }157 if (lastUser == nullptr) {158 return WalkResult::advance();159 }160 LDBG() << "Last user found: " << *lastUser;161 assert(lastUser->getBlock() == allocOp->getBlock());162 assert(lastUser->getBlock() == deallocOp->getBlock());163 // Move the dealloc op after the last user.164 deallocOp->moveAfter(lastUser);165 LDBG() << "Moved dealloc op after: " << *lastUser;166 167 return WalkResult::advance();168 });169 }170};171 172} // end anonymous namespace173