[mlir] Move buffer alias analysis into BufferUtils.h

[silvas]

Summary

It was in Bufferize.h, but that file is about bufferization conversion
patterns rather than heavyweight analyses needed for optimization
passes. BufferUtils.h is a place where we can put such utilities.

[herhut]

https://reviews.llvm.org/D89902 is already moving it. @frgossen FYI

[silvas]

In D89992#2349485, @herhut wrote:

https://reviews.llvm.org/D89902 is already moving it. @frgossen FYI

Awesome, thanks @frgossen !

mlir/include/mlir/Transforms/BufferUtils.h

+//===- 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<Value, 16>;
+  using ValueMapT = llvm::DenseMap<Value, ValueSetT>;
+
+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<Value> &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<Value, Operation *>;
+
+  /// Represents a list containing all alloc entries.
+  using AllocEntryList = SmallVector<AllocEntry, 8>;
+
+  /// 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 <typename DominatorT>
+  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

mlir/include/mlir/Transforms/Bufferize.h

patterns.insert<
     BufferizeCallOpConverter,
     BufferizeFuncOpConverter,
     BufferizeReturnOpConverter
       <ReturnOpSourceTy, ReturnOpTargetTy, CopyOpTy>
   >(context, converter);
   // 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<Value, 16>;
-  using ValueMapT = llvm::DenseMap<Value, ValueSetT>;
-
-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<Value> &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<Value, Operation *>;
-
-  /// Represents a list containing all alloc entries.
-  using AllocEntryList = SmallVector<AllocEntry, 8>;
-
-  /// 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 <typename DominatorT>
-  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

mlir/lib/Transforms/BufferDeallocation.cpp

 
 #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/Bufferize.h"
+#include "mlir/Transforms/BufferUtils.h"
 #include "mlir/Transforms/Passes.h"
-#include "llvm/ADT/SetOperations.h"
 
 using namespace mlir;
 
 /// Walks over all immediate return-like terminators in the given region.
 template <typename FuncT>
 static void walkReturnOperations(Region *region, const FuncT &func) {
   for (Block &block : *region)
     for (Operation &operation : block) {
static void getSuccessorRegions(RegionBranchOpInterface regionInterface,
   SmallVector<Attribute, 2> 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<void(Value)> 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<Value> &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<RegionSuccessor, 2> 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 &region : regionInterface.getOperation()->getRegions()) {
-      // Iterate over all successor region entries that are reachable from the
-      // current region.
-      SmallVector<RegionSuccessor, 2> 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(&region, [&](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<MemoryEffectOpInterface>(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<MemoryEffects::EffectInstance, 2> effects;
-    effectInterface.getEffectsOnValue(allocValue, effects);
-    return llvm::any_of(effects, [&](MemoryEffects::EffectInstance &it) {
-      return isa<MemoryEffects::Free>(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<MemoryEffects::EffectInstance, 2> effects;
-    opInterface.getEffects(effects);
-
-    SmallVector<MemoryEffects::EffectInstance, 2> allocateResultEffects;
-    llvm::copy_if(
-        effects, std::back_inserter(allocateResultEffects),
-        [=](MemoryEffects::EffectInstance &it) {
-          Value value = it.getValue();
-          return isa<MemoryEffects::Allocate>(it.getEffect()) && value &&
-                 value.isa<OpResult>() &&
-                 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<LoopLikeOpInterface>(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<RegionBranchOpInterface>(op)))
-    return false;
-
-  // Recurses into a region using the current region interface to find potential
-  // cycles.
-  SmallPtrSet<Region *, 4> visitedRegions;
-  std::function<bool(Region *)> 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<RegionSuccessor, 2> successors;
-    getSuccessorRegions(regionInterface, current->getRegionNumber(),
-                        successors);
-    for (RegionSuccessor &regionEntry : successors)
-      if (recurse(regionEntry.getSuccessor()))
-        return true;
-    return false;
-  };
-
-  // Start with all entry regions and test whether they induce a loop.
-  SmallVector<RegionSuccessor, 2> successorRegions;
-  getSuccessorRegions(regionInterface, /*index=*/llvm::None, successorRegions);
-  for (RegionSuccessor &regionEntry : successorRegions) {
-    if (recurse(regionEntry.getSuccessor()))
-      return true;
-    visitedRegions.clear();
-  }
-
-  return false;
-}
-
 namespace {
 
 //===----------------------------------------------------------------------===//
 // Backedges analysis
 //===----------------------------------------------------------------------===//
 
 /// A straight-forward program analysis which detects loop backedges induced by
 /// explicit control flow.

mlir/lib/Transforms/BufferOptimizations.cpp

 //===- BufferOptimizations.cpp - pre-pass optimizations for bufferization -===//
 //
 // 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 logic for two optimization passes. These passes try to
 // hoist alloc nodes to reduce the number of allocations and copies during
 // buffer deallocation.
 
 #include "PassDetail.h"
 #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;
 
 namespace {
 
 //===----------------------------------------------------------------------===//
 // BufferAllocationHoisting

mlir/lib/Transforms/Utils/BufferUtils.cpp

+//===- 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 <typename FuncT>
+static void walkReturnOperations(Region *region, const FuncT &func) {
+  for (Block &block : *region)
+    for (Operation &operation : block) {
+      // Skip non-return-like terminators.
+      if (operation.hasTrait<OpTrait::ReturnLike>())
+        func(&operation);
+    }
+}
+
+/// Wrapper for the actual `RegionBranchOpInterface.getSuccessorRegions`
+/// function that initializes the required `operandAttributes` array.
+static void getSuccessorRegions(RegionBranchOpInterface regionInterface,
+                                llvm::Optional<unsigned> index,
+                                SmallVectorImpl<RegionSuccessor> &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<Attribute, 2> 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<void(Value)> 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<Value> &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<RegionSuccessor, 2> 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 &region : regionInterface.getOperation()->getRegions()) {
+      // Iterate over all successor region entries that are reachable from the
+      // current region.
+      SmallVector<RegionSuccessor, 2> 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(&region, [&](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<MemoryEffectOpInterface>(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<MemoryEffects::EffectInstance, 2> effects;
+    effectInterface.getEffectsOnValue(allocValue, effects);
+    return llvm::any_of(effects, [&](MemoryEffects::EffectInstance &it) {
+      return isa<MemoryEffects::Free>(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<MemoryEffects::EffectInstance, 2> effects;
+    opInterface.getEffects(effects);
+
+    SmallVector<MemoryEffects::EffectInstance, 2> allocateResultEffects;
+    llvm::copy_if(
+        effects, std::back_inserter(allocateResultEffects),
+        [=](MemoryEffects::EffectInstance &it) {
+          Value value = it.getValue();
+          return isa<MemoryEffects::Allocate>(it.getEffect()) && value &&
+                 value.isa<OpResult>() &&
+                 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<LoopLikeOpInterface>(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<RegionBranchOpInterface>(op)))
+    return false;
+
+  // Recurses into a region using the current region interface to find potential
+  // cycles.
+  SmallPtrSet<Region *, 4> visitedRegions;
+  std::function<bool(Region *)> 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<RegionSuccessor, 2> successors;
+    getSuccessorRegions(regionInterface, current->getRegionNumber(),
+                        successors);
+    for (RegionSuccessor &regionEntry : successors)
+      if (recurse(regionEntry.getSuccessor()))
+        return true;
+    return false;
+  };
+
+  // Start with all entry regions and test whether they induce a loop.
+  SmallVector<RegionSuccessor, 2> successorRegions;
+  getSuccessorRegions(regionInterface, /*index=*/llvm::None, successorRegions);
+  for (RegionSuccessor &regionEntry : successorRegions) {
+    if (recurse(regionEntry.getSuccessor()))
+      return true;
+    visitedRegions.clear();
+  }
+
+  return false;
+}

mlir/lib/Transforms/Utils/CMakeLists.txt

 add_mlir_library(MLIRTransformUtils
+  BufferUtils.cpp
   FoldUtils.cpp
   GreedyPatternRewriteDriver.cpp
   InliningUtils.cpp
   LoopFusionUtils.cpp
   LoopUtils.cpp
   RegionUtils.cpp
   Utils.cpp