diff --git a/mlir/include/mlir/Transforms/BufferUtils.h b/mlir/include/mlir/Transforms/BufferUtils.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Transforms/BufferUtils.h @@ -0,0 +1,159 @@ +//===- BufferUtils.h - Utilities for dealing with buffers -------*- C++ -*-===// +// +// Part of the LLVM Project, 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 +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_TRANSFORMS_BUFFERUTILS_H +#define MLIR_TRANSFORMS_BUFFERUTILS_H + +#include "mlir/Analysis/Liveness.h" +#include "mlir/Dialect/StandardOps/IR/Ops.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/Dominance.h" +#include "mlir/IR/Function.h" +#include "mlir/IR/Operation.h" +#include "mlir/Transforms/DialectConversion.h" + +namespace mlir { + +/// A straight-forward alias analysis which ensures that all aliases of all +/// values will be determined. This is a requirement for the BufferPlacement +/// class since you need to determine safe positions to place alloc and +/// deallocs. +class BufferPlacementAliasAnalysis { +public: + using ValueSetT = SmallPtrSet; + using ValueMapT = llvm::DenseMap; + +public: + /// Constructs a new alias analysis using the op provided. + BufferPlacementAliasAnalysis(Operation *op); + + /// Find all immediate aliases this value could potentially have. + ValueMapT::const_iterator find(Value value) const { + return aliases.find(value); + } + + /// Returns the begin iterator to iterate over all aliases. + ValueMapT::const_iterator begin() const { return aliases.begin(); } + + /// Returns the end iterator that can be used in combination with find. + ValueMapT::const_iterator end() const { return aliases.end(); } + + /// Find all immediate and indirect aliases this value could potentially + /// have. Note that the resulting set will also contain the value provided as + /// it is an alias of itself. + ValueSetT resolve(Value value) const; + + /// Removes the given values from all alias sets. + void remove(const SmallPtrSetImpl &aliasValues); + +private: + /// This function constructs a mapping from values to its immediate aliases. + void build(Operation *op); + + /// Maps values to all immediate aliases this value can have. + ValueMapT aliases; +}; + +/// A simple analysis that detects allocation operations. +class BufferPlacementAllocs { +public: + /// Represents a tuple of allocValue and deallocOperation. + using AllocEntry = std::tuple; + + /// Represents a list containing all alloc entries. + using AllocEntryList = SmallVector; + + /// Get the start operation to place the given alloc value withing the + // specified placement block. + static Operation *getStartOperation(Value allocValue, Block *placementBlock, + const Liveness &liveness); + + /// Find an associated dealloc operation that is linked to the given + /// allocation node (if any). + static Operation *findDealloc(Value allocValue); + +public: + /// Initializes the internal list by discovering all supported allocation + /// nodes. + BufferPlacementAllocs(Operation *op); + + /// Returns the begin iterator to iterate over all allocations. + AllocEntryList::const_iterator begin() const { return allocs.begin(); } + + /// Returns the end iterator that can be used in combination with begin. + AllocEntryList::const_iterator end() const { return allocs.end(); } + + /// Returns the begin iterator to iterate over all allocations. + AllocEntryList::iterator begin() { return allocs.begin(); } + + /// Returns the end iterator that can be used in combination with begin. + AllocEntryList::iterator end() { return allocs.end(); } + + /// Registers a new allocation entry. + void registerAlloc(const AllocEntry &entry) { allocs.push_back(entry); } + +private: + /// Searches for and registers all supported allocation entries. + void build(Operation *op); + +private: + /// Maps allocation nodes to their associated blocks. + AllocEntryList allocs; +}; + +/// The base class for all BufferPlacement transformations. +class BufferPlacementTransformationBase { +public: + using ValueSetT = BufferPlacementAliasAnalysis::ValueSetT; + + /// Finds a common dominator for the given value while taking the positions + /// of the values in the value set into account. It supports dominator and + /// post-dominator analyses via template arguments. + template + static Block *findCommonDominator(Value value, const ValueSetT &values, + const DominatorT &doms) { + // Start with the current block the value is defined in. + Block *dom = value.getParentBlock(); + // Iterate over all aliases and their uses to find a safe placement block + // according to the given dominator information. + for (Value childValue : values) { + for (Operation *user : childValue.getUsers()) { + // Move upwards in the dominator tree to find an appropriate + // dominator block that takes the current use into account. + dom = doms.findNearestCommonDominator(dom, user->getBlock()); + } + // Take values without any users into account. + dom = doms.findNearestCommonDominator(dom, childValue.getParentBlock()); + } + return dom; + } + + /// Returns true if the given operation represents a loop by testing whether + /// it implements the `LoopLikeOpInterface` or the `RegionBranchOpInterface`. + /// In the case of a `RegionBranchOpInterface`, it checks all region-based + /// control-flow edges for cycles. + static bool isLoop(Operation *op); + + /// Constructs a new operation base using the given root operation. + BufferPlacementTransformationBase(Operation *op); + +protected: + /// Alias information that can be updated during the insertion of copies. + BufferPlacementAliasAnalysis aliases; + + /// Stores all internally managed allocations. + BufferPlacementAllocs allocs; + + /// The underlying liveness analysis to compute fine grained information + /// about alloc and dealloc positions. + Liveness liveness; +}; + +} // end namespace mlir + +#endif // MLIR_TRANSFORMS_BUFFERUTILS_H diff --git a/mlir/include/mlir/Transforms/Bufferize.h b/mlir/include/mlir/Transforms/Bufferize.h --- a/mlir/include/mlir/Transforms/Bufferize.h +++ b/mlir/include/mlir/Transforms/Bufferize.h @@ -216,141 +216,6 @@ // clang-format on } -/// A straight-forward alias analysis which ensures that all aliases of all -/// values will be determined. This is a requirement for the BufferPlacement -/// class since you need to determine safe positions to place alloc and -/// deallocs. -class BufferPlacementAliasAnalysis { -public: - using ValueSetT = SmallPtrSet; - using ValueMapT = llvm::DenseMap; - -public: - /// Constructs a new alias analysis using the op provided. - BufferPlacementAliasAnalysis(Operation *op); - - /// Find all immediate aliases this value could potentially have. - ValueMapT::const_iterator find(Value value) const { - return aliases.find(value); - } - - /// Returns the begin iterator to iterate over all aliases. - ValueMapT::const_iterator begin() const { return aliases.begin(); } - - /// Returns the end iterator that can be used in combination with find. - ValueMapT::const_iterator end() const { return aliases.end(); } - - /// Find all immediate and indirect aliases this value could potentially - /// have. Note that the resulting set will also contain the value provided as - /// it is an alias of itself. - ValueSetT resolve(Value value) const; - - /// Removes the given values from all alias sets. - void remove(const SmallPtrSetImpl &aliasValues); - -private: - /// This function constructs a mapping from values to its immediate aliases. - void build(Operation *op); - - /// Maps values to all immediate aliases this value can have. - ValueMapT aliases; -}; - -/// A simple analysis that detects allocation operations. -class BufferPlacementAllocs { -public: - /// Represents a tuple of allocValue and deallocOperation. - using AllocEntry = std::tuple; - - /// Represents a list containing all alloc entries. - using AllocEntryList = SmallVector; - - /// Get the start operation to place the given alloc value withing the - // specified placement block. - static Operation *getStartOperation(Value allocValue, Block *placementBlock, - const Liveness &liveness); - - /// Find an associated dealloc operation that is linked to the given - /// allocation node (if any). - static Operation *findDealloc(Value allocValue); - -public: - /// Initializes the internal list by discovering all supported allocation - /// nodes. - BufferPlacementAllocs(Operation *op); - - /// Returns the begin iterator to iterate over all allocations. - AllocEntryList::const_iterator begin() const { return allocs.begin(); } - - /// Returns the end iterator that can be used in combination with begin. - AllocEntryList::const_iterator end() const { return allocs.end(); } - - /// Returns the begin iterator to iterate over all allocations. - AllocEntryList::iterator begin() { return allocs.begin(); } - - /// Returns the end iterator that can be used in combination with begin. - AllocEntryList::iterator end() { return allocs.end(); } - - /// Registers a new allocation entry. - void registerAlloc(const AllocEntry &entry) { allocs.push_back(entry); } - -private: - /// Searches for and registers all supported allocation entries. - void build(Operation *op); - -private: - /// Maps allocation nodes to their associated blocks. - AllocEntryList allocs; -}; - -/// The base class for all BufferPlacement transformations. -class BufferPlacementTransformationBase { -public: - using ValueSetT = BufferPlacementAliasAnalysis::ValueSetT; - - /// Finds a common dominator for the given value while taking the positions - /// of the values in the value set into account. It supports dominator and - /// post-dominator analyses via template arguments. - template - static Block *findCommonDominator(Value value, const ValueSetT &values, - const DominatorT &doms) { - // Start with the current block the value is defined in. - Block *dom = value.getParentBlock(); - // Iterate over all aliases and their uses to find a safe placement block - // according to the given dominator information. - for (Value childValue : values) { - for (Operation *user : childValue.getUsers()) { - // Move upwards in the dominator tree to find an appropriate - // dominator block that takes the current use into account. - dom = doms.findNearestCommonDominator(dom, user->getBlock()); - } - // Take values without any users into account. - dom = doms.findNearestCommonDominator(dom, childValue.getParentBlock()); - } - return dom; - } - - /// Returns true if the given operation represents a loop by testing whether - /// it implements the `LoopLikeOpInterface` or the `RegionBranchOpInterface`. - /// In the case of a `RegionBranchOpInterface`, it checks all region-based - /// control-flow edges for cycles. - static bool isLoop(Operation *op); - - /// Constructs a new operation base using the given root operation. - BufferPlacementTransformationBase(Operation *op); - -protected: - /// Alias information that can be updated during the insertion of copies. - BufferPlacementAliasAnalysis aliases; - - /// Stores all internally managed allocations. - BufferPlacementAllocs allocs; - - /// The underlying liveness analysis to compute fine grained information - /// about alloc and dealloc positions. - Liveness liveness; -}; - } // end namespace mlir #endif // MLIR_TRANSFORMS_BUFFERIZE_H diff --git a/mlir/lib/Transforms/BufferDeallocation.cpp b/mlir/lib/Transforms/BufferDeallocation.cpp --- a/mlir/lib/Transforms/BufferDeallocation.cpp +++ b/mlir/lib/Transforms/BufferDeallocation.cpp @@ -58,9 +58,8 @@ #include "mlir/Interfaces/ControlFlowInterfaces.h" #include "mlir/Interfaces/LoopLikeInterface.h" #include "mlir/Pass/Pass.h" -#include "mlir/Transforms/Bufferize.h" +#include "mlir/Transforms/BufferUtils.h" #include "mlir/Transforms/Passes.h" -#include "llvm/ADT/SetOperations.h" using namespace mlir; @@ -91,245 +90,6 @@ regionInterface.getSuccessorRegions(index, operandAttributes, successors); } -//===----------------------------------------------------------------------===// -// BufferPlacementAliasAnalysis -//===----------------------------------------------------------------------===// - -/// Constructs a new alias analysis using the op provided. -BufferPlacementAliasAnalysis::BufferPlacementAliasAnalysis(Operation *op) { - build(op); -} - -/// Find all immediate and indirect aliases this value could potentially -/// have. Note that the resulting set will also contain the value provided as -/// it is an alias of itself. -BufferPlacementAliasAnalysis::ValueSetT -BufferPlacementAliasAnalysis::resolve(Value value) const { - ValueSetT result; - - /// Recursively determines alias information for the given value. It stores - /// all newly found potential aliases in the given result set. - std::function resolveRecursive = [&](Value current) { - if (!result.insert(current).second) - return; - auto it = aliases.find(current); - if (it == aliases.end()) - return; - for (Value alias : it->second) - resolveRecursive(alias); - }; - - resolveRecursive(value); - return result; -} - -/// Removes the given values from all alias sets. -void BufferPlacementAliasAnalysis::remove( - const SmallPtrSetImpl &aliasValues) { - for (auto &entry : aliases) - llvm::set_subtract(entry.second, aliasValues); -} - -/// This function constructs a mapping from values to its immediate aliases. -/// It iterates over all blocks, gets their predecessors, determines the -/// values that will be passed to the corresponding block arguments and -/// inserts them into the underlying map. Furthermore, it wires successor -/// regions and branch-like return operations from nested regions. -void BufferPlacementAliasAnalysis::build(Operation *op) { - // Registers all aliases of the given values. - auto registerAliases = [&](auto values, auto aliases) { - for (auto entry : llvm::zip(values, aliases)) - this->aliases[std::get<0>(entry)].insert(std::get<1>(entry)); - }; - - // Add additional aliases created by view changes to the alias list. - op->walk([&](ViewLikeOpInterface viewInterface) { - aliases[viewInterface.getViewSource()].insert( - viewInterface.getOperation()->getResult(0)); - }); - - // Query all branch interfaces to link block argument aliases. - op->walk([&](BranchOpInterface branchInterface) { - Block *parentBlock = branchInterface.getOperation()->getBlock(); - for (auto it = parentBlock->succ_begin(), e = parentBlock->succ_end(); - it != e; ++it) { - // Query the branch op interface to get the successor operands. - auto successorOperands = - branchInterface.getSuccessorOperands(it.getIndex()); - if (!successorOperands.hasValue()) - continue; - // Build the actual mapping of values to their immediate aliases. - registerAliases(successorOperands.getValue(), (*it)->getArguments()); - } - }); - - // Query the RegionBranchOpInterface to find potential successor regions. - op->walk([&](RegionBranchOpInterface regionInterface) { - // Extract all entry regions and wire all initial entry successor inputs. - SmallVector entrySuccessors; - getSuccessorRegions(regionInterface, /*index=*/llvm::None, entrySuccessors); - for (RegionSuccessor &entrySuccessor : entrySuccessors) { - // Wire the entry region's successor arguments with the initial - // successor inputs. - assert(entrySuccessor.getSuccessor() && - "Invalid entry region without an attached successor region"); - registerAliases(regionInterface.getSuccessorEntryOperands( - entrySuccessor.getSuccessor()->getRegionNumber()), - entrySuccessor.getSuccessorInputs()); - } - - // Wire flow between regions and from region exits. - for (Region ®ion : regionInterface.getOperation()->getRegions()) { - // Iterate over all successor region entries that are reachable from the - // current region. - SmallVector successorRegions; - getSuccessorRegions(regionInterface, region.getRegionNumber(), - successorRegions); - for (RegionSuccessor &successorRegion : successorRegions) { - // Iterate over all immediate terminator operations and wire the - // successor inputs with the operands of each terminator. - walkReturnOperations(®ion, [&](Operation *terminator) { - registerAliases(terminator->getOperands(), - successorRegion.getSuccessorInputs()); - }); - } - } - }); -} - -//===----------------------------------------------------------------------===// -// BufferPlacementAllocs -//===----------------------------------------------------------------------===// - -/// Get the start operation to place the given alloc value withing the -// specified placement block. -Operation *BufferPlacementAllocs::getStartOperation(Value allocValue, - Block *placementBlock, - const Liveness &liveness) { - // We have to ensure that we place the alloc before its first use in this - // block. - const LivenessBlockInfo &livenessInfo = *liveness.getLiveness(placementBlock); - Operation *startOperation = livenessInfo.getStartOperation(allocValue); - // Check whether the start operation lies in the desired placement block. - // If not, we will use the terminator as this is the last operation in - // this block. - if (startOperation->getBlock() != placementBlock) { - Operation *opInPlacementBlock = - placementBlock->findAncestorOpInBlock(*startOperation); - startOperation = opInPlacementBlock ? opInPlacementBlock - : placementBlock->getTerminator(); - } - - return startOperation; -} - -/// Finds associated deallocs that can be linked to our allocation nodes (if -/// any). -Operation *BufferPlacementAllocs::findDealloc(Value allocValue) { - auto userIt = llvm::find_if(allocValue.getUsers(), [&](Operation *user) { - auto effectInterface = dyn_cast(user); - if (!effectInterface) - return false; - // Try to find a free effect that is applied to one of our values - // that will be automatically freed by our pass. - SmallVector effects; - effectInterface.getEffectsOnValue(allocValue, effects); - return llvm::any_of(effects, [&](MemoryEffects::EffectInstance &it) { - return isa(it.getEffect()); - }); - }); - // Assign the associated dealloc operation (if any). - return userIt != allocValue.user_end() ? *userIt : nullptr; -} - -/// Initializes the internal list by discovering all supported allocation -/// nodes. -BufferPlacementAllocs::BufferPlacementAllocs(Operation *op) { build(op); } - -/// Searches for and registers all supported allocation entries. -void BufferPlacementAllocs::build(Operation *op) { - op->walk([&](MemoryEffectOpInterface opInterface) { - // Try to find a single allocation result. - SmallVector effects; - opInterface.getEffects(effects); - - SmallVector allocateResultEffects; - llvm::copy_if( - effects, std::back_inserter(allocateResultEffects), - [=](MemoryEffects::EffectInstance &it) { - Value value = it.getValue(); - return isa(it.getEffect()) && value && - value.isa() && - it.getResource() != - SideEffects::AutomaticAllocationScopeResource::get(); - }); - // If there is one result only, we will be able to move the allocation and - // (possibly existing) deallocation ops. - if (allocateResultEffects.size() != 1) - return; - // Get allocation result. - Value allocValue = allocateResultEffects[0].getValue(); - // Find the associated dealloc value and register the allocation entry. - allocs.push_back(std::make_tuple(allocValue, findDealloc(allocValue))); - }); -} - -//===----------------------------------------------------------------------===// -// BufferPlacementTransformationBase -//===----------------------------------------------------------------------===// - -/// Constructs a new transformation base using the given root operation. -BufferPlacementTransformationBase::BufferPlacementTransformationBase( - Operation *op) - : aliases(op), allocs(op), liveness(op) {} - -/// Returns true if the given operation represents a loop by testing whether it -/// implements the `LoopLikeOpInterface` or the `RegionBranchOpInterface`. In -/// the case of a `RegionBranchOpInterface`, it checks all region-based control- -/// flow edges for cycles. -bool BufferPlacementTransformationBase::isLoop(Operation *op) { - // If the operation implements the `LoopLikeOpInterface` it can be considered - // a loop. - if (isa(op)) - return true; - - // If the operation does not implement the `RegionBranchOpInterface`, it is - // (currently) not possible to detect a loop. - RegionBranchOpInterface regionInterface; - if (!(regionInterface = dyn_cast(op))) - return false; - - // Recurses into a region using the current region interface to find potential - // cycles. - SmallPtrSet visitedRegions; - std::function recurse = [&](Region *current) { - if (!current) - return false; - // If we have found a back edge, the parent operation induces a loop. - if (!visitedRegions.insert(current).second) - return true; - // Recurses into all region successors. - SmallVector successors; - getSuccessorRegions(regionInterface, current->getRegionNumber(), - successors); - for (RegionSuccessor ®ionEntry : successors) - if (recurse(regionEntry.getSuccessor())) - return true; - return false; - }; - - // Start with all entry regions and test whether they induce a loop. - SmallVector successorRegions; - getSuccessorRegions(regionInterface, /*index=*/llvm::None, successorRegions); - for (RegionSuccessor ®ionEntry : successorRegions) { - if (recurse(regionEntry.getSuccessor())) - return true; - visitedRegions.clear(); - } - - return false; -} - namespace { //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Transforms/BufferOptimizations.cpp b/mlir/lib/Transforms/BufferOptimizations.cpp --- a/mlir/lib/Transforms/BufferOptimizations.cpp +++ b/mlir/lib/Transforms/BufferOptimizations.cpp @@ -14,7 +14,7 @@ #include "mlir/IR/Operation.h" #include "mlir/Interfaces/LoopLikeInterface.h" #include "mlir/Pass/Pass.h" -#include "mlir/Transforms/Bufferize.h" +#include "mlir/Transforms/BufferUtils.h" #include "mlir/Transforms/Passes.h" using namespace mlir; diff --git a/mlir/lib/Transforms/Utils/BufferUtils.cpp b/mlir/lib/Transforms/Utils/BufferUtils.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Transforms/Utils/BufferUtils.cpp @@ -0,0 +1,280 @@ +//===- BufferUtils.cpp - Utilities for dealing with buffers ---------------===// +// +// Part of the LLVM Project, 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 "mlir/Transforms/BufferUtils.h" +#include "mlir/Interfaces/ControlFlowInterfaces.h" +#include "mlir/Interfaces/LoopLikeInterface.h" +#include "llvm/ADT/SetOperations.h" + +using namespace mlir; + +/// Walks over all immediate return-like terminators in the given region. +template +static void walkReturnOperations(Region *region, const FuncT &func) { + for (Block &block : *region) + for (Operation &operation : block) { + // Skip non-return-like terminators. + if (operation.hasTrait()) + func(&operation); + } +} + +/// Wrapper for the actual `RegionBranchOpInterface.getSuccessorRegions` +/// function that initializes the required `operandAttributes` array. +static void getSuccessorRegions(RegionBranchOpInterface regionInterface, + llvm::Optional index, + SmallVectorImpl &successors) { + // Create a list of null attributes for each operand to comply with the + // `getSuccessorRegions` interface definition that requires a single + // attribute per operand. + SmallVector operandAttributes( + regionInterface.getOperation()->getNumOperands()); + + // Get all successor regions using the temporarily allocated + // `operandAttributes`. + regionInterface.getSuccessorRegions(index, operandAttributes, successors); +} + +//===----------------------------------------------------------------------===// +// BufferPlacementAliasAnalysis +//===----------------------------------------------------------------------===// + +/// Constructs a new alias analysis using the op provided. +BufferPlacementAliasAnalysis::BufferPlacementAliasAnalysis(Operation *op) { + build(op); +} + +/// Find all immediate and indirect aliases this value could potentially +/// have. Note that the resulting set will also contain the value provided as +/// it is an alias of itself. +BufferPlacementAliasAnalysis::ValueSetT +BufferPlacementAliasAnalysis::resolve(Value value) const { + ValueSetT result; + + /// Recursively determines alias information for the given value. It stores + /// all newly found potential aliases in the given result set. + std::function resolveRecursive = [&](Value current) { + if (!result.insert(current).second) + return; + auto it = aliases.find(current); + if (it == aliases.end()) + return; + for (Value alias : it->second) + resolveRecursive(alias); + }; + + resolveRecursive(value); + return result; +} + +/// Removes the given values from all alias sets. +void BufferPlacementAliasAnalysis::remove( + const SmallPtrSetImpl &aliasValues) { + for (auto &entry : aliases) + llvm::set_subtract(entry.second, aliasValues); +} + +/// This function constructs a mapping from values to its immediate aliases. +/// It iterates over all blocks, gets their predecessors, determines the +/// values that will be passed to the corresponding block arguments and +/// inserts them into the underlying map. Furthermore, it wires successor +/// regions and branch-like return operations from nested regions. +void BufferPlacementAliasAnalysis::build(Operation *op) { + // Registers all aliases of the given values. + auto registerAliases = [&](auto values, auto aliases) { + for (auto entry : llvm::zip(values, aliases)) + this->aliases[std::get<0>(entry)].insert(std::get<1>(entry)); + }; + + // Add additional aliases created by view changes to the alias list. + op->walk([&](ViewLikeOpInterface viewInterface) { + aliases[viewInterface.getViewSource()].insert( + viewInterface.getOperation()->getResult(0)); + }); + + // Query all branch interfaces to link block argument aliases. + op->walk([&](BranchOpInterface branchInterface) { + Block *parentBlock = branchInterface.getOperation()->getBlock(); + for (auto it = parentBlock->succ_begin(), e = parentBlock->succ_end(); + it != e; ++it) { + // Query the branch op interface to get the successor operands. + auto successorOperands = + branchInterface.getSuccessorOperands(it.getIndex()); + if (!successorOperands.hasValue()) + continue; + // Build the actual mapping of values to their immediate aliases. + registerAliases(successorOperands.getValue(), (*it)->getArguments()); + } + }); + + // Query the RegionBranchOpInterface to find potential successor regions. + op->walk([&](RegionBranchOpInterface regionInterface) { + // Extract all entry regions and wire all initial entry successor inputs. + SmallVector entrySuccessors; + getSuccessorRegions(regionInterface, /*index=*/llvm::None, entrySuccessors); + for (RegionSuccessor &entrySuccessor : entrySuccessors) { + // Wire the entry region's successor arguments with the initial + // successor inputs. + assert(entrySuccessor.getSuccessor() && + "Invalid entry region without an attached successor region"); + registerAliases(regionInterface.getSuccessorEntryOperands( + entrySuccessor.getSuccessor()->getRegionNumber()), + entrySuccessor.getSuccessorInputs()); + } + + // Wire flow between regions and from region exits. + for (Region ®ion : regionInterface.getOperation()->getRegions()) { + // Iterate over all successor region entries that are reachable from the + // current region. + SmallVector successorRegions; + getSuccessorRegions(regionInterface, region.getRegionNumber(), + successorRegions); + for (RegionSuccessor &successorRegion : successorRegions) { + // Iterate over all immediate terminator operations and wire the + // successor inputs with the operands of each terminator. + walkReturnOperations(®ion, [&](Operation *terminator) { + registerAliases(terminator->getOperands(), + successorRegion.getSuccessorInputs()); + }); + } + } + }); +} + +//===----------------------------------------------------------------------===// +// BufferPlacementAllocs +//===----------------------------------------------------------------------===// + +/// Get the start operation to place the given alloc value withing the +// specified placement block. +Operation *BufferPlacementAllocs::getStartOperation(Value allocValue, + Block *placementBlock, + const Liveness &liveness) { + // We have to ensure that we place the alloc before its first use in this + // block. + const LivenessBlockInfo &livenessInfo = *liveness.getLiveness(placementBlock); + Operation *startOperation = livenessInfo.getStartOperation(allocValue); + // Check whether the start operation lies in the desired placement block. + // If not, we will use the terminator as this is the last operation in + // this block. + if (startOperation->getBlock() != placementBlock) { + Operation *opInPlacementBlock = + placementBlock->findAncestorOpInBlock(*startOperation); + startOperation = opInPlacementBlock ? opInPlacementBlock + : placementBlock->getTerminator(); + } + + return startOperation; +} + +/// Finds associated deallocs that can be linked to our allocation nodes (if +/// any). +Operation *BufferPlacementAllocs::findDealloc(Value allocValue) { + auto userIt = llvm::find_if(allocValue.getUsers(), [&](Operation *user) { + auto effectInterface = dyn_cast(user); + if (!effectInterface) + return false; + // Try to find a free effect that is applied to one of our values + // that will be automatically freed by our pass. + SmallVector effects; + effectInterface.getEffectsOnValue(allocValue, effects); + return llvm::any_of(effects, [&](MemoryEffects::EffectInstance &it) { + return isa(it.getEffect()); + }); + }); + // Assign the associated dealloc operation (if any). + return userIt != allocValue.user_end() ? *userIt : nullptr; +} + +/// Initializes the internal list by discovering all supported allocation +/// nodes. +BufferPlacementAllocs::BufferPlacementAllocs(Operation *op) { build(op); } + +/// Searches for and registers all supported allocation entries. +void BufferPlacementAllocs::build(Operation *op) { + op->walk([&](MemoryEffectOpInterface opInterface) { + // Try to find a single allocation result. + SmallVector effects; + opInterface.getEffects(effects); + + SmallVector allocateResultEffects; + llvm::copy_if( + effects, std::back_inserter(allocateResultEffects), + [=](MemoryEffects::EffectInstance &it) { + Value value = it.getValue(); + return isa(it.getEffect()) && value && + value.isa() && + it.getResource() != + SideEffects::AutomaticAllocationScopeResource::get(); + }); + // If there is one result only, we will be able to move the allocation and + // (possibly existing) deallocation ops. + if (allocateResultEffects.size() != 1) + return; + // Get allocation result. + Value allocValue = allocateResultEffects[0].getValue(); + // Find the associated dealloc value and register the allocation entry. + allocs.push_back(std::make_tuple(allocValue, findDealloc(allocValue))); + }); +} + +//===----------------------------------------------------------------------===// +// BufferPlacementTransformationBase +//===----------------------------------------------------------------------===// + +/// Constructs a new transformation base using the given root operation. +BufferPlacementTransformationBase::BufferPlacementTransformationBase( + Operation *op) + : aliases(op), allocs(op), liveness(op) {} + +/// Returns true if the given operation represents a loop by testing whether it +/// implements the `LoopLikeOpInterface` or the `RegionBranchOpInterface`. In +/// the case of a `RegionBranchOpInterface`, it checks all region-based control- +/// flow edges for cycles. +bool BufferPlacementTransformationBase::isLoop(Operation *op) { + // If the operation implements the `LoopLikeOpInterface` it can be considered + // a loop. + if (isa(op)) + return true; + + // If the operation does not implement the `RegionBranchOpInterface`, it is + // (currently) not possible to detect a loop. + RegionBranchOpInterface regionInterface; + if (!(regionInterface = dyn_cast(op))) + return false; + + // Recurses into a region using the current region interface to find potential + // cycles. + SmallPtrSet visitedRegions; + std::function recurse = [&](Region *current) { + if (!current) + return false; + // If we have found a back edge, the parent operation induces a loop. + if (!visitedRegions.insert(current).second) + return true; + // Recurses into all region successors. + SmallVector successors; + getSuccessorRegions(regionInterface, current->getRegionNumber(), + successors); + for (RegionSuccessor ®ionEntry : successors) + if (recurse(regionEntry.getSuccessor())) + return true; + return false; + }; + + // Start with all entry regions and test whether they induce a loop. + SmallVector successorRegions; + getSuccessorRegions(regionInterface, /*index=*/llvm::None, successorRegions); + for (RegionSuccessor ®ionEntry : successorRegions) { + if (recurse(regionEntry.getSuccessor())) + return true; + visitedRegions.clear(); + } + + return false; +} diff --git a/mlir/lib/Transforms/Utils/CMakeLists.txt b/mlir/lib/Transforms/Utils/CMakeLists.txt --- a/mlir/lib/Transforms/Utils/CMakeLists.txt +++ b/mlir/lib/Transforms/Utils/CMakeLists.txt @@ -1,4 +1,5 @@ add_mlir_library(MLIRTransformUtils + BufferUtils.cpp FoldUtils.cpp GreedyPatternRewriteDriver.cpp InliningUtils.cpp