[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, e.g., for buffer placement
+/// where safe positions allocs and deallocs must be determined.
+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) { build(op); }
+
+  /// Find all immediate aliases this value could potentially have.
+  ValueMapT::const_iterator find(Value value) const {
+    return aliases.find(value);
+  }
+
+  /// 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.
+  /// 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 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/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;
+
+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;
+}
+
+void BufferAliasAnalysis::remove(const SmallPtrSetImpl<Value> &aliasValues) {
+  for (auto &entry : aliases)
+    llvm::set_subtract(entry.second, aliasValues);
+}
+
+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 return-like 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/BufferPlacement.cpp

 // Furthermore, buffer placement also adds required new alloc and copy
 // operations to ensure that all buffers are deallocated.The main class is the
 // BufferPlacementPass 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, BufferPlacement
 // firstly finds all possible aliases for a single value (using the
-// BufferPlacementAliasAnalysis class). Consider the following example:
+// 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;
 //
 // Using liveness information on its own would cause us to place the allocs and
 // deallocs in the wrong block. This is due to the fact that %new_value will not
 // be liveOut of its block. Instead, we can place the alloc for %new_value
 // in bb0 and its associated dealloc in exit. Alternatively, the alloc can stay
 // (or even has to stay due to additional dependencies) at this location and 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. Finding the alloc position is similar and non-obvious.
 // However, the algorithm supports moving allocs to other places and introducing
 // copy buffers and placing deallocs in safe places to ensure that all buffers
 // will be freed in the end.
 //
 // TODO:
 // The current implementation does not support explicit-control-flow loops and
 // the resulting code will be invalid with respect to program semantics.
 // However, structured control-flow loops are fully supported. Furthermore, it
 // doesn't accept functions which return buffers already.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
+#include "mlir/Analysis/BufferAliasAnalysis.h"
 #include "mlir/Analysis/Liveness.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Pass/Pass.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) {
       // 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);
-}
-
 namespace {
 //===----------------------------------------------------------------------===//
-// BufferPlacementAliasAnalysis
-//===----------------------------------------------------------------------===//
-
-/// 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) { build(op); }
-
-  /// Find all immediate aliases this value could potentially have.
-  ValueMapT::const_iterator find(Value value) const {
-    return aliases.find(value);
-  }
-
-  /// 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 {
-    ValueSetT result;
-    resolveRecursive(value, result);
-    return result;
-  }
-
-  /// Removes the given values from all alias sets.
-  void remove(const SmallPtrSetImpl<Value> &aliasValues) {
-    for (auto &entry : aliases)
-      llvm::set_subtract(entry.second, aliasValues);
-  }
-
-private:
-  /// Recursively determines alias information for the given value. It stores
-  /// all newly found potential aliases in the given result set.
-  void resolveRecursive(Value value, ValueSetT &result) const {
-    if (!result.insert(value).second)
-      return;
-    auto it = aliases.find(value);
-    if (it == aliases.end())
-      return;
-    for (Value alias : it->second)
-      resolveRecursive(alias, result);
-  }
-
-  /// 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 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());
-          });
-        }
-      }
-    });
-  }
-
-  /// Maps values to all immediate aliases this value can have.
-  ValueMapT aliases;
-};
-
-//===----------------------------------------------------------------------===//
 // Backedges
 //===----------------------------------------------------------------------===//
 
 /// A straight-forward program analysis which detects loop backedges induced by
 /// explicit control flow.
 class Backedges {
 public:
   using BlockSetT = SmallPtrSet<Block *, 16>;

 
 //===----------------------------------------------------------------------===//
 // BufferPlacement
 //===----------------------------------------------------------------------===//
 
 // The main buffer placement analysis used to place allocs, copies and deallocs.
 class BufferPlacement {
 public:
-  using ValueSetT = BufferPlacementAliasAnalysis::ValueSetT;
+  using ValueSetT = BufferAliasAnalysis::ValueSetT;
 
   /// An intermediate representation of a single allocation node.
   struct AllocEntry {
     /// A reference to the associated allocation node.
     Value allocValue;
 
     /// The associated placement block in which the allocation should be
     /// performed.
introduceCopiesForRegionSuccessors(
           // Find a predecessor of our argRegion.
           return successorRegion.getSuccessor() == argRegion;
         });
 
     // Check whether the block argument belongs to an entry region of the
     // parent operation. In this case, we have to introduce an additional copy
     // for buffer that is passed to the argument.
     SmallVector<RegionSuccessor, 2> successorRegions;
-    getSuccessorRegions(regionInterface, llvm::None, successorRegions);
+    regionInterface.getSuccessorRegions(llvm::None, successorRegions);
     auto *it =
         llvm::find_if(successorRegions, [&](RegionSuccessor &successorRegion) {
           return successorRegion.getSuccessor() == argRegion;
         });
     if (it == successorRegions.end())
       return;
 
     // Determine the actual operand to introduce a copy for and rewire the
private:
   template <typename TPredicate>
   void introduceCopiesForRegionSuccessors(
       RegionBranchOpInterface regionInterface, MutableArrayRef<Region> regions,
       Value argValue, const TPredicate &regionPredicate) {
     for (Region &region : regions) {
       // Query the regionInterface to get all successor regions of the current
       // one.
       SmallVector<RegionSuccessor, 2> successorRegions;
-      getSuccessorRegions(regionInterface, region.getRegionNumber(),
+      regionInterface.getSuccessorRegions(region.getRegionNumber(),
                                           successorRegions);
       // Try to find a matching region successor.
       RegionSuccessor *regionSuccessor =
           llvm::find_if(successorRegions, regionPredicate);
       if (regionSuccessor == successorRegions.end())
         continue;
       // Get the operand index in the context of the current successor input
       // bindings.
       size_t operandIndex =
for (Value childValue : values)
       }
     return dom;
   }
 
   /// The operation this transformation was constructed from.
   Operation *operation;
 
   /// Alias information that can be updated during the insertion of copies.
-  BufferPlacementAliasAnalysis aliases;
+  BufferAliasAnalysis aliases;
 
   /// Maps allocation nodes to their associated blocks.
   AllocEntryList allocs;
 
   // Stores already copied allocations to avoid additional copies of copies.
   ValueSetT copiedValues;
 
   /// The underlying liveness analysis to compute fine grained information