[mlir] Split BufferUtils.h out of Bufferize.h

[silvas]

Summary

These utilities are more closely associated with the buffer
optimizations and buffer deallocation than with the dialect conversion
stuff in Bufferize.h. So move them out.

This makes Bufferize.h very easy to understand and completely focused on
dialect conversion.

[dfki-jugr]

mlir/lib/Transforms/BufferUtils.cpp

28

Is it useful to have a copy here and in BufferDeallocation since it introduces code duplication?

47

missing third /
withing -> within

[silvas]

Address comments.

mlir/lib/Transforms/BufferUtils.cpp

28

Actually, the helper is redundant. I replaced it with regionInterface.getSuccessorRegions(index, successors)

mlir/include/mlir/Transforms/BufferUtils.h

-//===- 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.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
-// 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/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:
   /// 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>;
protected:
 
   /// The underlying liveness analysis to compute fine grained information
   /// about alloc and dealloc positions.
   Liveness liveness;
 };
 
 } // end namespace mlir
 
-#endif // MLIR_TRANSFORMS_BUFFERIZE_H
+#endif // MLIR_TRANSFORMS_BUFFERUTILS_H

mlir/include/mlir/Transforms/Bufferize.h

 
 /// 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:
-  /// 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 = 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 <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.
-  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

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) {
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);
 }
 
-//===----------------------------------------------------------------------===//
-// 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

 // passes try to move alloc nodes out of blocks to reduce the number of
 // allocations and copies during buffer deallocation. The third pass tries to
 // convert heap-based allocations to stack-based allocations, if possible.
 
 #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;
 
 /// Returns true if the given operation implements a known high-level region-
 /// based control-flow interface.
 static bool isKnownControlFlowInterface(Operation *op) {
   return isa<LoopLikeOpInterface, RegionBranchOpInterface>(op);

mlir/lib/Transforms/BufferUtils.cpp

+//===- 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;
+
+/// 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);
+}
+
+//===----------------------------------------------------------------------===//
+// 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/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_mlir_library(MLIRTransforms
   BufferDeallocation.cpp
   BufferOptimizations.cpp
   BufferResultsToOutParams.cpp
+  BufferUtils.cpp
   Bufferize.cpp
   Canonicalizer.cpp
   CopyRemoval.cpp
   CSE.cpp
   Inliner.cpp
   LocationSnapshot.cpp
   LoopCoalescing.cpp
   LoopFusion.cpp