diff --git a/mlir/include/mlir/Transforms/Bufferize.h b/mlir/include/mlir/Transforms/BufferUtils.h copy from mlir/include/mlir/Transforms/Bufferize.h copy to mlir/include/mlir/Transforms/BufferUtils.h --- a/mlir/include/mlir/Transforms/Bufferize.h +++ b/mlir/include/mlir/Transforms/BufferUtils.h @@ -1,4 +1,4 @@ -//===- Bufferize.h - Bufferization utilities --------------------*- C++ -*-===// +//===- BufferUtils.h - Buffer optimization utilities ------------*- C++ -*-===// // // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // See https://llvm.org/LICENSE.txt for license information. @@ -6,59 +6,25 @@ // //===----------------------------------------------------------------------===// // -// We use the term "bufferize" to mean conversion from tensor types to -// memref types. -// -// Generally speaking, for each op that operates on tensor types, a conversion -// pattern needs to be written. The infrastructure in this file assists in -// defining these conversion patterns in a composable way. -// -// Bufferization conversion patterns should generally use the ordinary -// conversion pattern classes (e.g. OpConversionPattern). A TypeConverter -// (accessible with getTypeConverter()) is available if needed for converting -// types. +// This file provides utilities for passes optimizing code that has already +// been converted to buffers. // //===----------------------------------------------------------------------===// -#ifndef MLIR_TRANSFORMS_BUFFERIZE_H -#define MLIR_TRANSFORMS_BUFFERIZE_H +#ifndef MLIR_TRANSFORMS_BUFFERUTILS_H +#define MLIR_TRANSFORMS_BUFFERUTILS_H #include "mlir/Analysis/BufferAliasAnalysis.h" #include "mlir/Analysis/Liveness.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinOps.h" #include "mlir/IR/Dominance.h" +#include "mlir/IR/Function.h" #include "mlir/IR/Operation.h" #include "mlir/Transforms/DialectConversion.h" namespace mlir { -/// A helper type converter class that automatically populates the relevant -/// materializations and type conversions for bufferization. -class BufferizeTypeConverter : public TypeConverter { -public: - BufferizeTypeConverter(); -}; - -/// Marks ops used by bufferization for type conversion materializations as -/// "legal" in the given ConversionTarget. -/// -/// This function should be called by all bufferization passes using -/// BufferizeTypeConverter so that materializations work proprely. One exception -/// is bufferization passes doing "full" conversions, where it can be desirable -/// for even the materializations to remain illegal so that they are eliminated, -/// such as via the patterns in -/// populateEliminateBufferizeMaterializationsPatterns. -void populateBufferizeMaterializationLegality(ConversionTarget &target); - -/// Populate patterns to eliminate bufferize materializations. -/// -/// In particular, these are the tensor_load/tensor_to_memref ops. -void populateEliminateBufferizeMaterializationsPatterns( - MLIRContext *context, BufferizeTypeConverter &typeConverter, - OwningRewritePatternList &patterns); - /// A simple analysis that detects allocation operations. class BufferPlacementAllocs { public: @@ -68,8 +34,8 @@ /// 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. + /// Get the start operation to place the given alloc value within the + /// specified placement block. static Operation *getStartOperation(Value allocValue, Block *placementBlock, const Liveness &liveness); @@ -156,4 +122,4 @@ } // end namespace mlir -#endif // MLIR_TRANSFORMS_BUFFERIZE_H +#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 @@ -59,101 +59,6 @@ MLIRContext *context, BufferizeTypeConverter &typeConverter, OwningRewritePatternList &patterns); -/// 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 = BufferAliasAnalysis::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. - BufferAliasAnalysis 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,7 +58,7 @@ #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" @@ -75,155 +75,6 @@ } } -/// 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); -} - -//===----------------------------------------------------------------------===// -// 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 { //===----------------------------------------------------------------------===// @@ -443,8 +294,7 @@ // parent operation. In this case, we have to introduce an additional copy // for buffer that is passed to the argument. SmallVector successorRegions; - getSuccessorRegions(regionInterface, /*index=*/llvm::None, - successorRegions); + regionInterface.getSuccessorRegions(/*index=*/llvm::None, successorRegions); auto *it = llvm::find_if(successorRegions, [&](RegionSuccessor &successorRegion) { return successorRegion.getSuccessor() == argRegion; @@ -497,8 +347,8 @@ // Query the regionInterface to get all successor regions of the current // one. SmallVector successorRegions; - getSuccessorRegions(regionInterface, region.getRegionNumber(), - successorRegions); + regionInterface.getSuccessorRegions(region.getRegionNumber(), + successorRegions); // Try to find a matching region successor. RegionSuccessor *regionSuccessor = llvm::find_if(successorRegions, regionPredicate); 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 @@ -15,7 +15,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/BufferUtils.cpp b/mlir/lib/Transforms/BufferUtils.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Transforms/BufferUtils.cpp @@ -0,0 +1,156 @@ +//===- BufferUtils.cpp - buffer transformation utilities ------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// +// +// This file implements utilties for buffer optimization passes. +// +//===----------------------------------------------------------------------===// + +#include "mlir/Transforms/BufferUtils.h" +#include "PassDetail.h" +#include "mlir/Dialect/Linalg/IR/LinalgOps.h" +#include "mlir/Dialect/StandardOps/IR/Ops.h" +#include "mlir/IR/Operation.h" +#include "mlir/Interfaces/ControlFlowInterfaces.h" +#include "mlir/Interfaces/LoopLikeInterface.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/Passes.h" +#include "llvm/ADT/SetOperations.h" + +using namespace mlir; + +//===----------------------------------------------------------------------===// +// 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; + regionInterface.getSuccessorRegions(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; + regionInterface.getSuccessorRegions(/*index=*/llvm::None, successorRegions); + for (RegionSuccessor ®ionEntry : successorRegions) { + if (recurse(regionEntry.getSuccessor())) + return true; + visitedRegions.clear(); + } + + return false; +} diff --git a/mlir/lib/Transforms/CMakeLists.txt b/mlir/lib/Transforms/CMakeLists.txt --- a/mlir/lib/Transforms/CMakeLists.txt +++ b/mlir/lib/Transforms/CMakeLists.txt @@ -4,6 +4,7 @@ BufferDeallocation.cpp BufferOptimizations.cpp BufferResultsToOutParams.cpp + BufferUtils.cpp Bufferize.cpp Canonicalizer.cpp CopyRemoval.cpp