[MLIR] Extract buffer alias analysis for reuse

[frgossen]

Summary

Extract buffer alias analysis from buffer placement.

[pifon2a]

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

143

nit: tyop!

[rriddle]

mlir/lib/Analysis/BufferAliasAnalysis.cpp

37

Can you please write this as a non-recursive algorithm instead? With unknown bounds, recursive algorithms quickly become unwieldy.

[frgossen]

Address comments

[rriddle]

mlir/lib/Analysis/BufferAliasAnalysis.cpp

24

nit: Use pop_back_val

[frgossen]

Address comment

[frgossen]

Rebase

[jpienaar]

Is this change related to build breakage? (https://buildkite.com/mlir/mlir-core/builds/8813)

[frgossen]

In D89902#2350128, @jpienaar wrote:

Is this change related to build breakage? (https://buildkite.com/mlir/mlir-core/builds/8813)

Yes but I fixed it in https://github.com/llvm/llvm-project/commit/8039b3f966804c6578f8265d5eb7bdde8dec3ddd

mlir/include/mlir/Analysis/BufferAliasAnalysis.h

+//===- BufferAliasAnalysis.h - Buffer alias analysis for MLIR ---*- 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_ANALYSIS_BUFFERALIASANALYSIS_H
+#define MLIR_ANALYSIS_BUFFERALIASANALYSIS_H
+
+#include "mlir/IR/Operation.h"
+#include "llvm/ADT/SmallPtrSet.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 BufferAliasAnalysis {
+public:
+  using ValueSetT = SmallPtrSet<Value, 16>;
+  using ValueMapT = llvm::DenseMap<Value, ValueSetT>;
+
+public:
+  /// Constructs a new alias analysis using the op provided.
+  BufferAliasAnalysis(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;
+};
+
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_BUFFERALIASANALYSIS_H

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

def RegionBranchOpInterface : OpInterface<"RegionBranchOpInterface"> {
 
   let verify = [{
     static_assert(!ConcreteOpType::template hasTrait<OpTrait::ZeroRegion>(),
                   "expected operation to have non-zero regions");
     return success();
   }];
 
   let extraClassDeclaration = [{
+    /// Convenience helper in case none of the operands is known.
+    void getSuccessorRegions(Optional<unsigned> index,
+                             SmallVectorImpl<RegionSuccessor> &regions) {
+       SmallVector<Attribute, 2> nullAttrs(getOperation()->getNumOperands());
+       getSuccessorRegions(index, nullAttrs, regions);
+    }
+
     /// Verify types along control flow edges described by this interface.
     static LogicalResult verifyTypes(Operation *op) {
       return detail::verifyTypesAlongControlFlowEdges(op);
     }
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // ControlFlow Traits
 //===----------------------------------------------------------------------===//
 
 // Op is "return-like".
 def ReturnLike : NativeOpTrait<"ReturnLike">;
 
 #endif // MLIR_INTERFACES_CONTROLFLOWINTERFACES

mlir/include/mlir/Transforms/Bufferize.h

 // those cases, BufferizeConversionPattern and its related classes should be
 // used, which provide access to a BufferizeTypeConverter directly.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TRANSFORMS_BUFFERIZE_H
 #define MLIR_TRANSFORMS_BUFFERIZE_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"
 
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>;

 private:
   /// Maps allocation nodes to their associated blocks.
   AllocEntryList allocs;
 };
 
 /// The base class for all BufferPlacement transformations.
 class BufferPlacementTransformationBase {
 public:
-  using ValueSetT = BufferPlacementAliasAnalysis::ValueSetT;
+  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.
public:
   /// 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;
+  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/Analysis/BufferAliasAnalysis.cpp

+//===- BufferAliasAnalysis.cpp - Buffer alias analysis for MLIR -*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/BufferAliasAnalysis.h"
+
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
+#include "llvm/ADT/SetOperations.h"
+
+using namespace mlir;
+
+/// Constructs a new alias analysis using the op provided.
+BufferAliasAnalysis::BufferAliasAnalysis(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.
+BufferAliasAnalysis::ValueSetT
+BufferAliasAnalysis::resolve(Value rootValue) const {
+  ValueSetT result;
+  SmallVector<Value, 8> queue;
+  queue.push_back(rootValue);
+  while (!queue.empty()) {
+    Value currentValue = queue.pop_back_val();
+    if (result.insert(currentValue).second) {
+      auto it = aliases.find(currentValue);
+      if (it != aliases.end()) {
+        for (Value aliasValue : it->second)
+          queue.push_back(aliasValue);
+      }
+    }
+  }
+  return result;
+}
+
+/// Removes the given values from all alias sets.
+void BufferAliasAnalysis::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 BufferAliasAnalysis::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;
+    regionInterface.getSuccessorRegions(/*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;
+      regionInterface.getSuccessorRegions(region.getRegionNumber(),
+                                          successorRegions);
+      for (RegionSuccessor &successorRegion : successorRegions) {
+        // Iterate over all immediate terminator operations and wire the
+        // successor inputs with the operands of each terminator.
+        for (Block &block : region) {
+          for (Operation &operation : block) {
+            if (operation.hasTrait<OpTrait::ReturnLike>())
+              registerAliases(operation.getOperands(),
+                              successorRegion.getSuccessorInputs());
+          }
+        }
+      }
+    }
+  });
+}

mlir/lib/Analysis/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   AffineAnalysis.cpp
   AffineStructures.cpp
+  BufferAliasAnalysis.cpp
   CallGraph.cpp
   Liveness.cpp
   LoopAnalysis.cpp
   NestedMatcher.cpp
   PresburgerSet.cpp
   SliceAnalysis.cpp
   Utils.cpp
   )
 
 add_mlir_library(MLIRAnalysis
+  BufferAliasAnalysis.cpp
   CallGraph.cpp
   Liveness.cpp
   SliceAnalysis.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Analysis
 
   DEPENDS
add_mlir_library(MLIRLoopAnalysis
   MLIRAffine
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRInferTypeOpInterface
   MLIRPresburger
   MLIRSCF
   )
   
-add_subdirectory(Presburger)
+add_subdirectory(Presburger)
 

mlir/lib/Transforms/BufferDeallocation.cpp

 // Furthermore, buffer placement also adds required new alloc and copy
 // operations to ensure that all buffers are deallocated. The main class is the
 // BufferDeallocationPass class that implements the underlying algorithm. In
 // order to put allocations and deallocations at safe positions, it is
 // significantly important to put them into the correct blocks. However, the
 // liveness analysis does not pay attention to aliases, which can occur due to
 // branches (and their associated block arguments) in general. For this purpose,
 // BufferDeallocation firstly finds all possible aliases for a single value
-// (using the BufferPlacementAliasAnalysis class). Consider the following
+// (using the BufferAliasAnalysis class). Consider the following
 // example:
 //
 // ^bb0(%arg0):
 //   cond_br %cond, ^bb1, ^bb2
 // ^bb1:
 //   br ^exit(%arg0)
 // ^bb2:
 //   %new_value = ...
 //   br ^exit(%new_value)
 // ^exit(%arg1):
 //   return %arg1;
 //
 // We should place the dealloc for %new_value in exit. However, we have to free
 // the buffer in the same block, because it cannot be freed in the post
 // dominator. However, this requires a new copy buffer for %arg1 that will
-// contain the actual contents. Using the class BufferPlacementAliasAnalysis, we
+// contain the actual contents. Using the class BufferAliasAnalysis, we
 // will find out that %new_value has a potential alias %arg1. In order to find
 // the dealloc position we have to find all potential aliases, iterate over
 // their uses and find the common post-dominator block (note that additional
 // copies and buffers remove potential aliases and will influence the placement
 // of the deallocs). In all cases, the computed block can be safely used to free
 // the %new_value buffer (may be exit or bb2) as it will die and we can use
 // liveness information to determine the exact operation after which we have to
 // insert the dealloc. However, the algorithm supports introducing copy buffers
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) {