[mlir][bufferization] Cache enclosing repetitive region

[maerhart]

Summary

The getEnclosingRepetitiveRegion functions walk the ancestor regions everytime which can be expensive especially when there are multiple regions inbetween. This commit adds a cache to the bufferization analysis to remember the result of the walk.

[springerm]

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

575–577

This cache should be cleared when OneShotAnalysisState::resetCache() is called. We need a virtual resetCache function in this class. (And override + call base impl. in OneShotAnalysisState.)

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

57–61

An enclosing repetitive region can be nullptr. In that case, the stored region in the cache will be nullptr and this check will return "false". I would change this to if (enclosingRepetitiveRegionCache.contains(op)).

[maerhart]

Cache nullptr results as well and consider them in the lookups. I used find_as to avoid looking up the value twice. Also, add the resetCache virtual method.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

 //===- BufferizableOpInterface.h - Bufferizable Ops -------------*- 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_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_
 #define MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_
 
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/LLVM.h"
+#include "llvm/ADT/DenseMapInfoVariant.h"
 #include "llvm/ADT/SetVector.h"
 #include <optional>
 
 #include "mlir/Dialect/Bufferization/IR/BufferizationEnums.h.inc"
 
 namespace mlir {
 class OpBuilder;
 namespace func {
public:
   AnalysisState(const AnalysisState &) = delete;
 
   virtual ~AnalysisState() = default;
 
   static bool classof(const AnalysisState *base) { return true; }
 
   TypeID getType() const { return type; }
 
+  /// Return the closest enclosing repetitive region around the given op.
+  Region *getEnclosingRepetitiveRegion(Operation *op,
+                                       const BufferizationOptions &options);
+
+  /// Return the closest enclosing repetitive region around the place where the
+  /// given value is defined.
+  Region *getEnclosingRepetitiveRegion(Value value,
+                                       const BufferizationOptions &options);
+
+  /// Return the closest enclosing repetitive region around the given block.
+  Region *getEnclosingRepetitiveRegion(Block *block,
+                                       const BufferizationOptions &options);
+
+  virtual void resetCache();
+
 protected:
   AnalysisState(const BufferizationOptions &options, TypeID type);
 
 private:
   /// A reference to current bufferization options.
   const BufferizationOptions &options;
 
   /// The type of analysis.
   TypeID type;
+
+  /// Cache containing closest ancestor repetitive Region.
+  DenseMap<std::variant<Operation *, Block *, Region *, Value>, Region *>
+      enclosingRepetitiveRegionCache;
 };
 
 /// Create an AllocTensorOp for the given shaped value (memref or tensor).
 /// If `copy` is set, the shaped value is copied. Otherwise, a tensor with
 /// undefined contents is allocated.
 FailureOr<Value>
 allocateTensorForShapedValue(OpBuilder &b, Location loc, Value shapedValue,
                              bool escape, const BufferizationOptions &options,

 BaseMemRefType
 getMemRefTypeWithStaticIdentityLayout(TensorType tensorType,
                                       Attribute memorySpace = nullptr);
 
 /// Return the owner of the given value. In case of a BlockArgument that is the
 /// owner of the block. In case of an OpResult that is the defining op.
 Operation *getOwnerOfValue(Value value);
 
-/// Return the closest enclosing repetitive region around the given op.
-Region *getEnclosingRepetitiveRegion(Operation *op,
-                                     const BufferizationOptions &options);
-
-/// Return the closest enclosing repetitive region around the place where the
-/// given value is defined.
-Region *getEnclosingRepetitiveRegion(Value value,
-                                     const BufferizationOptions &options);
-
-/// Return the closest enclosing repetitive region around the given block.
-Region *getEnclosingRepetitiveRegion(Block *block,
-                                     const BufferizationOptions &options);
-
 /// Assuming that the given region is repetitive, find the next enclosing
 /// repetitive region.
 Region *getNextEnclosingRepetitiveRegion(Region *region,
                                          const BufferizationOptions &options);
 
 namespace detail {
 /// This is the default implementation of
 /// BufferizableOpInterface::getAliasingOpOperands. Should not be called from

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h

public:
   /// Return true if the buffer of the given tensor value is writable.
   bool isWritable(Value value) const;
 
   /// Find the definitions of the given tensor value or retrieve them from the
   /// cache.
   const SetVector<Value> &findDefinitionsCached(Value value);
 
   /// Reset cached data structures.
-  void resetCache();
+  void resetCache() override;
 
   /// Union the alias sets of `v1` and `v2`.
   void unionAliasSets(Value v1, Value v2);
 
   /// Union the equivalence classes of `v1` and `v2`.
   void unionEquivalenceClasses(Value v1, Value v2);
 
   /// Base class for OneShotAnalysisState extensions that allow

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

static bool isRepetitiveRegion(Region *region,
                                const BufferizationOptions &options) {
   Operation *op = region->getParentOp();
   if (auto bufferizableOp = options.dynCastBufferizableOp(op))
     if (bufferizableOp.isRepetitiveRegion(region->getRegionNumber()))
       return true;
   return false;
 }
 
-Region *bufferization::getEnclosingRepetitiveRegion(
+Region *AnalysisState::getEnclosingRepetitiveRegion(
     Operation *op, const BufferizationOptions &options) {
   if (!op->getBlock())
     return nullptr;
-  return getEnclosingRepetitiveRegion(op->getBlock(), options);
+  if (auto iter = enclosingRepetitiveRegionCache.find_as(op);
+      iter != enclosingRepetitiveRegionCache.end())
+    return iter->second;
+  return enclosingRepetitiveRegionCache[op] =
+             getEnclosingRepetitiveRegion(op->getBlock(), options);
 }
 
-Region *bufferization::getEnclosingRepetitiveRegion(
+Region *AnalysisState::getEnclosingRepetitiveRegion(
     Value value, const BufferizationOptions &options) {
+  if (auto iter = enclosingRepetitiveRegionCache.find_as(value);
+      iter != enclosingRepetitiveRegionCache.end())
+    return iter->second;
+
   Region *region = value.getParentRegion();
+  // Collect all visited regions since we only know the repetitive region we
+  // want to map it to later on
+  SmallVector<Region *> visitedRegions;
   while (region) {
+    visitedRegions.push_back(region);
     if (isRepetitiveRegion(region, options))
-      return region;
+      break;
     region = region->getParentRegion();
   }
-  return nullptr;
+  enclosingRepetitiveRegionCache[value] = region;
+  for (Region *r : visitedRegions)
+    enclosingRepetitiveRegionCache[r] = region;
+  return region;
 }
 
-Region *bufferization::getEnclosingRepetitiveRegion(
+Region *AnalysisState::getEnclosingRepetitiveRegion(
     Block *block, const BufferizationOptions &options) {
+  if (auto iter = enclosingRepetitiveRegionCache.find_as(block);
+      iter != enclosingRepetitiveRegionCache.end())
+    return iter->second;
+
   Region *region = block->getParent();
   Operation *op = nullptr;
+  // Collect all visited regions since we only know the repetitive region we
+  // want to map it to later on
+  SmallVector<Region *> visitedRegions;
   do {
     op = region->getParentOp();
     if (isRepetitiveRegion(region, options))
-      return region;
+      break;
   } while ((region = op->getParentRegion()));
-  return nullptr;
+
+  enclosingRepetitiveRegionCache[block] = region;
+  for (Region *r : visitedRegions)
+    enclosingRepetitiveRegionCache[r] = region;
+  return region;
 }
 
+void AnalysisState::resetCache() { enclosingRepetitiveRegionCache.clear(); }
+
 Region *bufferization::getNextEnclosingRepetitiveRegion(
     Region *region, const BufferizationOptions &options) {
   assert(isRepetitiveRegion(region, options) && "expected repetitive region");
   while ((region = region->getParentRegion())) {
     if (isRepetitiveRegion(region, options))
       break;
   }
   return region;

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

 /// 1. If the closest enclosing repetitive region of DEF is a proper ancestor of
 ///    a repetitive region that enclosing both READ and WRITE, we cannot rule
 ///    out RaW conflict due to the ordering of ops.
 /// 2. Otherwise: There are no loops that interfere with our analysis; for
 ///    analysis purposes, we can assume that there are no loops/repetitive
 ///    regions. I.e., we can rule out a RaW conflict if READ happensBefore WRITE
 ///    or WRITE happensBefore DEF. (Checked in `hasReadAfterWriteInterference`.)
 ///
-bool canUseOpDominance(OpOperand *uRead, OpOperand *uWrite,
+static bool canUseOpDominance(OpOperand *uRead, OpOperand *uWrite,
                               const SetVector<Value> &definitions,
-                       const AnalysisState &state) {
+                              AnalysisState &state) {
   const BufferizationOptions &options = state.getOptions();
   for (Value def : definitions) {
-    Region *rRead = getEnclosingRepetitiveRegion(uRead->getOwner(), options);
-    Region *rDef = getEnclosingRepetitiveRegion(def, options);
+    Region *rRead =
+        state.getEnclosingRepetitiveRegion(uRead->getOwner(), options);
+    Region *rDef = state.getEnclosingRepetitiveRegion(def, options);
 
     // READ and DEF are in the same repetitive region. `happensBefore` can be
     // used to rule out RaW conflicts due to op ordering.
     if (rRead == rDef)
       continue;
 
     // Find the enclosing repetitive region of READ that is closest to DEF but
     // not the repetitive region of DEF itself.

 // Find the values that define the contents of the given value.
 const llvm::SetVector<Value> &
 OneShotAnalysisState::findDefinitionsCached(Value value) {
   if (!cachedDefinitions.count(value))
     cachedDefinitions[value] = findDefinitions(value);
   return cachedDefinitions[value];
 }
 
-void OneShotAnalysisState::resetCache() { cachedDefinitions.clear(); }
+void OneShotAnalysisState::resetCache() {
+  AnalysisState::resetCache();
+  cachedDefinitions.clear();
+}
 
 /// Determine if `operand` can be bufferized in-place.
 static LogicalResult
 bufferizableInPlaceAnalysisImpl(OpOperand &operand, OneShotAnalysisState &state,
                                 const DominanceInfo &domInfo) {
   LLVM_DEBUG(
       llvm::dbgs() << "//===-------------------------------------------===//\n"
                    << "Analyzing operand #" << operand.getOperandNumber()