[mlir][bufferize][NFC] Merge AnalysisState and BufferizationAliasInfo

[springerm]

Summary

There is no longer a need to keep the two separate. This is in preparation of reusing the same AnalysisState for tensor.empty elimination and One-Shot Bufferize (to address performance bottlenecks).

[scotttodd]

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

207–242

nit: "uses uses" typo

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

180–182

How does this interact with reusing the state? Maybe this comment should be updated in https://reviews.llvm.org/D143404 ?

This is in preparation of reusing the same AnalysisState for tensor.empty elimination and One-Shot Bufferize (to address performance bottlenecks).

(I'm also wondering how much of a sharp edge this is and if there are other ways to make the analysis more robust to mutations / usage patterns)

[springerm]

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

180–182

I added a comment to the class, clarifying this: Modifying the IR invalidates the analysis in general, but there are a few things that are safe. (E.g., adding new ops or changing uses of tensor.empty as long as we re-analyze these ops.)

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

 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include <string>
 
 namespace mlir {
 namespace bufferization {
 
 struct OneShotBufferizationOptions;
-class BufferizationAliasInfo;
 struct BufferizationStatistics;
 class OneShotAnalysisState;
 
 /// Options for analysis-enabled bufferization.
 struct OneShotBufferizationOptions : public BufferizationOptions {
   enum class AnalysisHeuristic { BottomUp, TopDown };
 
   OneShotBufferizationOptions() = default;
 
   /// Specifies whether returning newly allocated memrefs should be allowed.
   /// Otherwise, a pass failure is triggered.
   bool allowReturnAllocs = false;
 
   /// The heuristic controls the order in which ops are traversed during the
   /// analysis.
   AnalysisHeuristic analysisHeuristic = AnalysisHeuristic::BottomUp;
 
   /// Specify the functions that should not be analyzed. copyBeforeWrite will be
   /// set to true when bufferizing them.
   llvm::ArrayRef<std::string> noAnalysisFuncFilter;
 };
 
-/// The BufferizationAliasInfo class maintains a list of buffer aliases and
-/// equivalence classes to support bufferization.
-class BufferizationAliasInfo {
-public:
-  explicit BufferizationAliasInfo(Operation *rootOp);
-
-  // BufferizationAliasInfo should be passed as a reference.
-  BufferizationAliasInfo(const BufferizationAliasInfo &) = delete;
-
-  /// Add a new entry for `v` in the `aliasInfo` and `equivalentInfo`. In the
-  /// beginning the alias and equivalence sets only contain `v` itself.
-  void createAliasInfoEntry(Value v);
-
-  /// Insert an info entry for `newValue` and merge its alias set with that of
-  /// `alias`.
-  void insertNewBufferAlias(Value newValue, Value alias);
-
-  /// Insert an info entry for `newValue` and merge its alias set with that of
-  /// `alias`. Additionally, merge their equivalence classes.
-  void insertNewBufferEquivalence(Value newValue, Value alias);
-
-  /// Set the inPlace bufferization spec to true.
-  /// Merge result's and operand's aliasing sets and iterate to a fixed point.
-  void bufferizeInPlace(OpOperand &operand, AnalysisState &state);
-
-  /// Set the inPlace bufferization spec to false.
-  void bufferizeOutOfPlace(OpOperand &operand);
-
-  /// Return true if `v1` and `v2` may bufferize to aliasing buffers.
-  bool areAliasingBufferizedValues(Value v1, Value v2) const {
-    return aliasInfo.isEquivalent(v1, v2);
-  }
-
-  /// Return true if `v1` and `v2` bufferize to equivalent buffers.
-  bool areEquivalentBufferizedValues(Value v1, Value v2) const {
-    return equivalentInfo.isEquivalent(v1, v2);
-  }
-
-  /// Union the alias sets of `v1` and `v2`.
-  void unionAliasSets(Value v1, Value v2) { aliasInfo.unionSets(v1, v2); }
-
-  /// Union the equivalence classes of `v1` and `v2`.
-  void unionEquivalenceClasses(Value v1, Value v2) {
-    equivalentInfo.unionSets(v1, v2);
-  }
-
-  /// Apply `fun` to all the members of the equivalence class of `v`.
-  void applyOnEquivalenceClass(Value v, function_ref<void(Value)> fun) const;
-
-  /// Apply `fun` to all aliases of `v`.
-  void applyOnAliases(Value v, function_ref<void(Value)> fun) const;
-
-  /// Mark a value as in-place bufferized.
-  void markInPlace(OpOperand &o) { inplaceBufferized.insert(&o); }
-
-  /// Return `true` if a value was marked as in-place bufferized.
-  bool isInPlace(OpOperand &opOperand) const;
-
-  int64_t getStatNumTensorOutOfPlace() const { return statNumTensorOutOfPlace; }
-  int64_t getStatNumTensorInPlace() const { return statNumTensorInPlace; }
-
-private:
-  /// llvm::EquivalenceClasses wants comparable elements. This comparator uses
-  /// uses pointer comparison on the defining op. This is a poor man's
-  /// comparison but it's not like UnionFind needs ordering anyway.
-  struct ValueComparator {
-    bool operator()(const Value &lhs, const Value &rhs) const {
-      return lhs.getImpl() < rhs.getImpl();
-    }
-  };
-
-  using EquivalenceClassRangeType = llvm::iterator_range<
-      llvm::EquivalenceClasses<Value, ValueComparator>::member_iterator>;
-  /// Check that aliasInfo for `v` exists and return a reference to it.
-  EquivalenceClassRangeType getAliases(Value v) const;
-
-  /// Set of all OpResults that were decided to bufferize in-place.
-  llvm::DenseSet<OpOperand *> inplaceBufferized;
-
-  /// Auxiliary structure to store all the values a given value may alias with.
-  /// Alias information is "may be" conservative: In the presence of branches, a
-  /// value may alias with one of multiple other values. The concrete aliasing
-  /// value may not even be known at compile time. All such values are
-  /// considered to be aliases.
-  llvm::EquivalenceClasses<Value, ValueComparator> aliasInfo;
-
-  /// Auxiliary structure to store all the equivalent buffer classes. Equivalent
-  /// buffer information is "must be" conservative: Only if two values are
-  /// guaranteed to be equivalent at runtime, they said to be equivalent. It is
-  /// possible that, in the presence of branches, it cannot be determined
-  /// statically if two values are equivalent. In that case, the values are
-  /// considered to be not equivalent.
-  llvm::EquivalenceClasses<Value, ValueComparator> equivalentInfo;
-
-  // Bufferization statistics.
-  int64_t statNumTensorOutOfPlace = 0;
-  int64_t statNumTensorInPlace = 0;
-};
-
 /// State for analysis-enabled bufferization. This class keeps track of alias
-/// (via BufferizationAliasInfo) to decide if tensor OpOperands should bufferize
-/// in-place.
+/// sets, equivalence sets, in-place OpOperands and other things.
+///
+/// Note: Modifying the IR generally invalidates the result of the analysis.
+/// Adding new operations is safe if they are analyzed subsequently.
 class OneShotAnalysisState : public AnalysisState {
 public:
   OneShotAnalysisState(Operation *op,
                        const OneShotBufferizationOptions &options);
 
   OneShotAnalysisState(const OneShotAnalysisState &) = delete;
 
   ~OneShotAnalysisState() override = default;
 
   static bool classof(const AnalysisState *base) {
     return base->getType() == TypeID::get<OneShotAnalysisState>();
   }
 
   /// Return a reference to the BufferizationOptions.
   const OneShotBufferizationOptions &getOptions() const {
     return static_cast<const OneShotBufferizationOptions &>(
         AnalysisState::getOptions());
   }
 
-  /// Return a reference to the BufferizationAliasInfo.
-  BufferizationAliasInfo &getAliasInfo() { return aliasInfo; }
+  /// Apply `fun` to all the members of the equivalence class of `v`.
+  void applyOnEquivalenceClass(Value v, function_ref<void(Value)> fun) const;
 
-  /// Return `true` if the given OpResult has been decided to bufferize inplace.
-  bool isInPlace(OpOperand &opOperand) const override;
+  /// Apply `fun` to all aliases of `v`.
+  void applyOnAliases(Value v, function_ref<void(Value)> fun) const;
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   bool areEquivalentBufferizedValues(Value v1, Value v2) const override;
 
   /// Return true if `v1` and `v2` may bufferize to aliasing buffers.
   bool areAliasingBufferizedValues(Value v1, Value v2) const override;
 
-  /// Return `true` if the given tensor has undefined contents.
-  bool hasUndefinedContents(OpOperand *opOperand) const override;
+  /// Mark the given OpOperand as in-place and merge the results' and operand's
+  /// aliasing sets.
+  void bufferizeInPlace(OpOperand &operand);
 
-  /// Return true if the given tensor (or an aliasing tensor) is yielded from
-  /// the containing block. Also include all aliasing tensors in the same block.
-  bool isTensorYielded(Value tensor) const override;
+  /// Mark the given OpOperand as out-of-place.
+  void bufferizeOutOfPlace(OpOperand &operand);
+
+  /// Add a new entry for `v` in the `aliasInfo` and `equivalentInfo`. In the
+  /// beginning the alias and equivalence sets only contain `v` itself.
+  void createAliasInfoEntry(Value v);
 
   /// Find all tensor values in the given operation that have undefined contents
   /// and store them in `undefinedTensorUses`.
   void gatherUndefinedTensorUses(Operation *op);
 
   /// Find all tensors that are yielded/returned from a block and store them in
   /// `yieldedTensors`. Also include all aliasing tensors in the same block.
   void gatherYieldedTensors(Operation *op);
 
+  int64_t getStatNumTensorOutOfPlace() const { return statNumTensorOutOfPlace; }
+  int64_t getStatNumTensorInPlace() const { return statNumTensorInPlace; }
+
+  /// Return `true` if the given tensor has undefined contents.
+  bool hasUndefinedContents(OpOperand *opOperand) const override;
+
+  /// Return `true` if the given OpResult has been decided to bufferize inplace.
+  bool isInPlace(OpOperand &opOperand) const override;
+
+  /// Return true if the given tensor (or an aliasing tensor) is yielded from
+  /// the containing block. Also include all aliasing tensors in the same block.
+  bool isTensorYielded(Value tensor) const override;
+
   /// Return true if the buffer of the given tensor value is written to. Must
   /// not be called for values inside not yet analyzed functions.
   bool isValueWritten(Value value) const;
 
   /// Return true if the buffer of the given tensor value is writable.
   bool isWritable(Value value) const;
 
+  /// 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
   /// OneShotAnalysisState to contain user-specified information in the state
   /// object. Clients are expected to derive this class, add the desired fields,
   /// and make the derived class compatible with the MLIR TypeID mechanism.
   ///
   /// ```mlir
   /// class MyExtension final : public OneShotAnalysisState::Extension {
   /// public:
public:
 
   /// Returns the extension of the specified type.
   template <typename Ty>
   const Ty *getExtension() const {
     return const_cast<OneShotAnalysisState *>(this)->getExtension<Ty>();
   }
 
 private:
-  /// `aliasInfo` keeps track of aliasing and equivalent values. Only internal
-  /// functions and `runOneShotBufferize` may access this object.
-  BufferizationAliasInfo aliasInfo;
+  /// llvm::EquivalenceClasses wants comparable elements. This comparator uses
+  /// pointer comparison on the defining op. This is a poor man's comparison
+  /// but it's not like UnionFind needs ordering anyway.
+  struct ValueComparator {
+    bool operator()(const Value &lhs, const Value &rhs) const {
+      return lhs.getImpl() < rhs.getImpl();
+    }
+  };
+
+  using EquivalenceClassRangeType = llvm::iterator_range<
+      llvm::EquivalenceClasses<Value, ValueComparator>::member_iterator>;
+  /// Check that aliasInfo for `v` exists and return a reference to it.
+  EquivalenceClassRangeType getAliases(Value v) const;
+
+  /// Set of all OpResults that were decided to bufferize in-place.
+  llvm::DenseSet<OpOperand *> inplaceBufferized;
+
+  /// Auxiliary structure to store all the values a given value may alias with.
+  /// Alias information is "may be" conservative: In the presence of branches, a
+  /// value may alias with one of multiple other values. The concrete aliasing
+  /// value may not even be known at compile time. All such values are
+  /// considered to be aliases.
+  llvm::EquivalenceClasses<Value, ValueComparator> aliasInfo;
+
+  /// Auxiliary structure to store all the equivalent buffer classes. Equivalent
+  /// buffer information is "must be" conservative: Only if two values are
+  /// guaranteed to be equivalent at runtime, they said to be equivalent. It is
+  /// possible that, in the presence of branches, it cannot be determined
+  /// statically if two values are equivalent. In that case, the values are
+  /// considered to be not equivalent.
+  llvm::EquivalenceClasses<Value, ValueComparator> equivalentInfo;
+
+  // Bufferization statistics.
+  int64_t statNumTensorOutOfPlace = 0;
+  int64_t statNumTensorInPlace = 0;
 
   /// A set of all tensors (and maybe aliasing tensors) that yielded from a
   /// block.
   DenseSet<Value> yieldedTensors;
 
   /// A set of uses of tensors that have undefined contents.
   DenseSet<OpOperand *> undefinedTensorUses;
 
   /// Extensions attached to the state, identified by the TypeID of their type.

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

 using namespace mlir;
 using namespace mlir::bufferization;
 
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific attribute manipulation.
 // These are for testing and debugging only. Bufferization information is stored
-// in BufferizationAliasInfo. When run with `testAnalysisOnly`, the IR is
+// in OneShotBufferizationState. When run with `testAnalysisOnly`, the IR is
 // annotated with the results of the analysis, so that they can be checked in
 // tests.
 //===----------------------------------------------------------------------===//
 
 /// Attribute marker to specify op operands that bufferize in-place.
 constexpr StringLiteral kInPlaceOperandsAttrName = "__inplace_operands_attr__";
 
 /// Mark whether OpOperand will be bufferized inplace.
for (OpOperand &opOperand : op->getOpOperands())
         inPlaceVector[opOperand.getOperandNumber()] = "false";
   }
   inPlaceVector[opOperand.getOperandNumber()] = inPlace ? "true" : "false";
   op->setAttr(kInPlaceOperandsAttrName,
               OpBuilder(op).getStrArrayAttr(inPlaceVector));
 }
 
 //===----------------------------------------------------------------------===//
-// BufferizationAliasInfo
+// OneShotAnalysisState
 //===----------------------------------------------------------------------===//
 
-BufferizationAliasInfo::BufferizationAliasInfo(Operation *rootOp) {
-  rootOp->walk([&](Operation *op) {
+OneShotAnalysisState::OneShotAnalysisState(
+    Operation *op, const OneShotBufferizationOptions &options)
+    : AnalysisState(options, TypeID::get<OneShotAnalysisState>()) {
+  // Set up alias sets.
+  op->walk([&](Operation *op) {
     for (Value v : op->getResults())
       if (v.getType().isa<TensorType>())
         createAliasInfoEntry(v);
     for (Region &r : op->getRegions())
       for (Block &b : r.getBlocks())
         for (auto bbArg : b.getArguments())
           if (bbArg.getType().isa<TensorType>())
             createAliasInfoEntry(bbArg);
   });
-}
-
-/// Add a new entry for `v` in the `aliasInfo` and `equivalentInfo`. In the
-/// beginning the alias and equivalence sets only contain `v` itself.
-void BufferizationAliasInfo::createAliasInfoEntry(Value v) {
-  aliasInfo.insert(v);
-  equivalentInfo.insert(v);
-}
-
-/// Insert an info entry for `newValue` and merge its alias set with that of
-/// `alias`.
-void BufferizationAliasInfo::insertNewBufferAlias(Value newValue, Value alias) {
-  createAliasInfoEntry(newValue);
-  aliasInfo.unionSets(newValue, alias);
-}
-
-/// Insert an info entry for `newValue` and merge its alias set with that of
-/// `alias`. Additionally, merge their equivalence classes.
-void BufferizationAliasInfo::insertNewBufferEquivalence(Value newValue,
-                                                        Value alias) {
-  insertNewBufferAlias(newValue, alias);
-  equivalentInfo.unionSets(newValue, alias);
-}
-
-/// Return `true` if a value was marked as in-place bufferized.
-bool BufferizationAliasInfo::isInPlace(OpOperand &operand) const {
-  return inplaceBufferized.contains(&operand);
-}
-
-/// Set the inPlace bufferization spec to true.
-void BufferizationAliasInfo::bufferizeInPlace(OpOperand &operand,
-                                              AnalysisState &state) {
-  if (inplaceBufferized.contains(&operand))
-    return;
-  markInPlace(operand);
-  for (OpResult result : state.getAliasingOpResults(operand))
-    aliasInfo.unionSets(result, operand.get());
-  ++statNumTensorInPlace;
-}
 
-/// Set the inPlace bufferization spec to false.
-void BufferizationAliasInfo::bufferizeOutOfPlace(OpOperand &operand) {
-  assert(!inplaceBufferized.contains(&operand) &&
-         "OpOperand was already decided to bufferize inplace");
-  ++statNumTensorOutOfPlace;
+  // Mark OpOperands in-place that must bufferize in-place.
+  op->walk([&](BufferizableOpInterface bufferizableOp) {
+    if (!options.isOpAllowed(bufferizableOp))
+      return WalkResult::skip();
+    for (OpOperand &opOperand : bufferizableOp->getOpOperands())
+      if (opOperand.get().getType().isa<TensorType>())
+        if (bufferizableOp.mustBufferizeInPlace(opOperand, *this))
+          bufferizeInPlace(opOperand);
+    return WalkResult::advance();
+  });
 }
 
-/// Apply `fun` to all the members of the equivalence class of `v`.
-void BufferizationAliasInfo::applyOnEquivalenceClass(
+void OneShotAnalysisState::applyOnEquivalenceClass(
     Value v, function_ref<void(Value)> fun) const {
   auto leaderIt = equivalentInfo.findLeader(v);
   for (auto mit = leaderIt, meit = equivalentInfo.member_end(); mit != meit;
        ++mit) {
     fun(*mit);
   }
 }
 
-/// Apply `fun` to all aliases of `v`.
-void BufferizationAliasInfo::applyOnAliases(
-    Value v, function_ref<void(Value)> fun) const {
+void OneShotAnalysisState::applyOnAliases(Value v,
+                                          function_ref<void(Value)> fun) const {
   auto leaderIt = aliasInfo.findLeader(v);
   for (auto mit = leaderIt, meit = aliasInfo.member_end(); mit != meit; ++mit) {
     fun(*mit);
   }
 }
 
-BufferizationAliasInfo::EquivalenceClassRangeType
-BufferizationAliasInfo::getAliases(Value v) const {
-  DenseSet<Value> res;
-  auto it = aliasInfo.findValue(aliasInfo.getLeaderValue(v));
-  for (auto mit = aliasInfo.member_begin(it), meit = aliasInfo.member_end();
-       mit != meit; ++mit) {
-    res.insert(static_cast<Value>(*mit));
-  }
-  return BufferizationAliasInfo::EquivalenceClassRangeType(
-      aliasInfo.member_begin(it), aliasInfo.member_end());
+bool OneShotAnalysisState::areEquivalentBufferizedValues(Value v1,
+                                                         Value v2) const {
+  return equivalentInfo.isEquivalent(v1, v2);
 }
 
-//===----------------------------------------------------------------------===//
-// OneShotAnalysisState
-//===----------------------------------------------------------------------===//
-
-OneShotAnalysisState::OneShotAnalysisState(
-    Operation *op, const OneShotBufferizationOptions &options)
-    : AnalysisState(options, TypeID::get<OneShotAnalysisState>()),
-      aliasInfo(op) {
-  // Set up alias sets for OpResults that must bufferize in-place. This should
-  // be done before making any other bufferization decisions.
-  op->walk([&](BufferizableOpInterface bufferizableOp) {
-    if (!options.isOpAllowed(bufferizableOp))
-      return WalkResult::skip();
-    for (OpOperand &opOperand : bufferizableOp->getOpOperands())
-      if (opOperand.get().getType().isa<TensorType>())
-        if (bufferizableOp.mustBufferizeInPlace(opOperand, *this))
-          aliasInfo.bufferizeInPlace(opOperand, *this);
-    return WalkResult::advance();
-  });
+bool OneShotAnalysisState::areAliasingBufferizedValues(Value v1,
+                                                       Value v2) const {
+  return aliasInfo.isEquivalent(v1, v2);
 }
 
-bool OneShotAnalysisState::isInPlace(OpOperand &opOperand) const {
-  return aliasInfo.isInPlace(opOperand);
+void OneShotAnalysisState::bufferizeInPlace(OpOperand &operand) {
+  if (inplaceBufferized.contains(&operand))
+    return;
+  inplaceBufferized.insert(&operand);
+  for (OpResult result : getAliasingOpResults(operand))
+    aliasInfo.unionSets(result, operand.get());
+  ++statNumTensorInPlace;
 }
 
-bool OneShotAnalysisState::areEquivalentBufferizedValues(Value v1,
-                                                         Value v2) const {
-  return aliasInfo.areEquivalentBufferizedValues(v1, v2);
+void OneShotAnalysisState::bufferizeOutOfPlace(OpOperand &operand) {
+  assert(!inplaceBufferized.contains(&operand) &&
+         "OpOperand was already decided to bufferize inplace");
+  ++statNumTensorOutOfPlace;
 }
 
-bool OneShotAnalysisState::areAliasingBufferizedValues(Value v1,
-                                                       Value v2) const {
-  return aliasInfo.areAliasingBufferizedValues(v1, v2);
+void OneShotAnalysisState::createAliasInfoEntry(Value v) {
+  aliasInfo.insert(v);
+  equivalentInfo.insert(v);
 }
 
 // Gather yielded tensors in `yieldedTensors` by querying all aliases. This is
 // to ensure that such information is available during bufferization time.
-// Alias information can no longer be queried through BufferizationAliasInfo
-// once we have started modifying the IR.
+// Alias information can no longer be queried once we have started modifying
+// the IR.
 void OneShotAnalysisState::gatherYieldedTensors(Operation *op) {
   op->walk([&](Operation *returnOp) {
     if (!isRegionReturnLike(returnOp) || !getOptions().isOpAllowed(returnOp))
       return WalkResult::advance();
 
     for (OpOperand &returnValOperand : returnOp->getOpOperands()) {
       Value returnVal = returnValOperand.get();
       // Skip non-tensor values.
       if (!returnVal.getType().isa<TensorType>())
         continue;
 
       // Add all aliases of the returned value. But only the ones that are in
       // the same block.
-      aliasInfo.applyOnAliases(returnVal, [&](Value v) {
+      applyOnAliases(returnVal, [&](Value v) {
         if (auto bbArg = v.dyn_cast<BlockArgument>()) {
           if (bbArg.getOwner()->getParentOp() == returnOp->getParentOp())
             yieldedTensors.insert(bbArg);
           return;
         }
         Operation *definingOp = v.getDefiningOp();
         if (definingOp->getParentOp() == returnOp->getParentOp())
           yieldedTensors.insert(v);
op->walk([&](Operation *op) {
     return WalkResult::advance();
   });
 }
 
 bool OneShotAnalysisState::hasUndefinedContents(OpOperand *opOperand) const {
   return undefinedTensorUses.contains(opOperand);
 }
 
+bool OneShotAnalysisState::isInPlace(OpOperand &opOperand) const {
+  return inplaceBufferized.contains(&opOperand);
+}
+
 bool OneShotAnalysisState::isTensorYielded(Value tensor) const {
   return yieldedTensors.contains(tensor);
 }
 
 bool OneShotAnalysisState::isValueWritten(Value value) const {
   bool isWritten = false;
-  aliasInfo.applyOnAliases(value, [&](Value val) {
+  applyOnAliases(value, [&](Value val) {
     for (OpOperand &use : val.getUses())
       if (isInPlace(use) && bufferizesToMemoryWrite(use))
         isWritten = true;
   });
   return isWritten;
 }
 
 bool OneShotAnalysisState::isWritable(Value value) const {
   // TODO: Out-of-place bufferized value could be considered writable.
   if (auto bufferizableOp = getOptions().dynCastBufferizableOp(value))
     return bufferizableOp.isWritable(value, *this);
 
   // Query BufferizableOpInterface to see if the BlockArgument is writable.
   if (auto bbArg = value.dyn_cast<BlockArgument>())
     if (auto bufferizableOp =
             getOptions().dynCastBufferizableOp(bbArg.getOwner()->getParentOp()))
       return bufferizableOp.isWritable(bbArg, *this);
 
   // Not a bufferizable op: The conservative answer is "not writable".
   return false;
 }
 
+void OneShotAnalysisState::unionAliasSets(Value v1, Value v2) {
+  aliasInfo.unionSets(v1, v2);
+}
+
+void OneShotAnalysisState::unionEquivalenceClasses(Value v1, Value v2) {
+  equivalentInfo.unionSets(v1, v2);
+}
+
 OneShotAnalysisState::Extension::~Extension() = default;
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific alias analysis.
 //===----------------------------------------------------------------------===//
 
 /// Return true if opOperand has been decided to bufferize in-place.
 static bool isInplaceMemoryWrite(OpOperand &opOperand,
-                                 const BufferizationAliasInfo &aliasInfo,
-                                 const AnalysisState &state) {
+                                 const OneShotAnalysisState &state) {
   // OpOperands that do not bufferize to a memory write do not write in-place.
   if (!state.bufferizesToMemoryWrite(opOperand))
     return false;
   // Check current bufferization decisions.
-  return aliasInfo.isInPlace(opOperand);
+  return state.isInPlace(opOperand);
 }
 
 /// Return true if `a` happens before `b`, i.e., `a` or one of its ancestors
 /// properly dominates `b` and `b` is not inside `a`.
 static bool happensBefore(Operation *a, Operation *b,
                           const DominanceInfo &domInfo) {
   do {
     // TODO: Instead of isProperAncestor + properlyDominates, we should use

 
 /// Given sets of uses and writes, return true if there is a RaW conflict under
 /// the assumption that all given reads/writes alias the same buffer and that
 /// all given writes bufferize inplace.
 ///
 /// A conflict is: According to SSA use-def chains, a read R is supposed to read
 /// the result of a definition W1. But because of bufferization decisions, R
 /// actually reads another definition W2.
-static bool hasReadAfterWriteInterference(
-    const DenseSet<OpOperand *> &usesRead,
-    const DenseSet<OpOperand *> &usesWrite, const DominanceInfo &domInfo,
-    AnalysisState &state, const BufferizationAliasInfo &aliasInfo) {
+static bool
+hasReadAfterWriteInterference(const DenseSet<OpOperand *> &usesRead,
+                              const DenseSet<OpOperand *> &usesWrite,
+                              const DominanceInfo &domInfo,
+                              OneShotAnalysisState &state) {
   const BufferizationOptions &options = state.getOptions();
 
   for (OpOperand *uRead : usesRead) {
     Operation *readingOp = uRead->getOwner();
     LLVM_DEBUG(llvm::dbgs() << "\n- check conflict:\n");
     LLVM_DEBUG(llvm::dbgs() << "  uRead = operand " << uRead->getOperandNumber()
                             << " of " << *readingOp << "\n");
 
for (OpOperand *uRead : usesRead) {
     }
   }
 
   return false;
 }
 
 // Helper function to iterate on aliases of `root` and capture the writes.
 static void getAliasingInplaceWrites(DenseSet<OpOperand *> &res, Value root,
-                                     const BufferizationAliasInfo &aliasInfo,
-                                     const AnalysisState &state) {
-  aliasInfo.applyOnAliases(root, [&](Value alias) {
+                                     const OneShotAnalysisState &state) {
+  state.applyOnAliases(root, [&](Value alias) {
     for (auto &use : alias.getUses())
       // Inplace write to a value that aliases root.
-      if (isInplaceMemoryWrite(use, aliasInfo, state))
+      if (isInplaceMemoryWrite(use, state))
         res.insert(&use);
   });
 }
 
 // Helper function to iterate on aliases of `root` and capture the reads.
 static void getAliasingReads(DenseSet<OpOperand *> &res, Value root,
-                             const BufferizationAliasInfo &aliasInfo,
-                             const AnalysisState &state) {
-  aliasInfo.applyOnAliases(root, [&](Value alias) {
+                             const OneShotAnalysisState &state) {
+  state.applyOnAliases(root, [&](Value alias) {
     for (auto &use : alias.getUses()) {
       // Read of a value that aliases root.
       if (state.bufferizesToMemoryRead(use)) {
         res.insert(&use);
         continue;
       }
 
       // Read of a dependent value in the SSA use-def chain. E.g.:

 /// If `checkConsistencyOnly` is true, this function checks if there is a
 /// read-after-write conflict without bufferizing `operand` inplace. This would
 /// indicate a problem with the current inplace bufferization decisions.
 ///
 /// Note: If `checkConsistencyOnly`, this function may be called with a null
 /// OpResult. In that case, only the consistency of bufferization decisions
 /// involving aliases of the given OpOperand are checked.
 static bool wouldCreateReadAfterWriteInterference(
-    OpOperand &operand, const DominanceInfo &domInfo, AnalysisState &state,
-    const BufferizationAliasInfo &aliasInfo,
-    bool checkConsistencyOnly = false) {
+    OpOperand &operand, const DominanceInfo &domInfo,
+    OneShotAnalysisState &state, bool checkConsistencyOnly = false) {
   // Collect reads and writes of all aliases of OpOperand and OpResult.
   DenseSet<OpOperand *> usesRead, usesWrite;
-  getAliasingReads(usesRead, operand.get(), aliasInfo, state);
-  getAliasingInplaceWrites(usesWrite, operand.get(), aliasInfo, state);
+  getAliasingReads(usesRead, operand.get(), state);
+  getAliasingInplaceWrites(usesWrite, operand.get(), state);
   for (OpResult result : state.getAliasingOpResults(operand)) {
-    getAliasingReads(usesRead, result, aliasInfo, state);
-    getAliasingInplaceWrites(usesWrite, result, aliasInfo, state);
+    getAliasingReads(usesRead, result, state);
+    getAliasingInplaceWrites(usesWrite, result, state);
   }
   if (!checkConsistencyOnly && state.bufferizesToMemoryWrite(operand))
     usesWrite.insert(&operand);
 
-  return hasReadAfterWriteInterference(usesRead, usesWrite, domInfo, state,
-                                       aliasInfo);
+  return hasReadAfterWriteInterference(usesRead, usesWrite, domInfo, state);
 }
 
 /// Annotate IR with details about the detected non-writability conflict.
 static void annotateNonWritableTensor(Value value) {
   static int64_t counter = 0;
   OpBuilder b(value.getContext());
   std::string id = "W_" + std::to_string(counter++);
   if (auto opResult = value.dyn_cast<OpResult>()) {

 
 /// Check the reverse SSA use-def chain (following aliasing OpOperands) for
 /// non-writable tensor values. Stop searching when an out-of-place bufferized
 /// OpOperand was found (or when the OpOperand was not bufferized yet).
 /// `currentOpOperand` is assumed to be in-place, even if that decision was not
 /// materialized in `aliasInfo` yet.
 static bool
 hasPrecedingAliasingNonWritableTensor(Value value, OpOperand *currentOpOperand,
-                                      const BufferizationAliasInfo &aliasInfo,
                                       const OneShotAnalysisState &state) {
   SmallVector<Value> worklist;
   worklist.push_back(value);
   while (!worklist.empty()) {
     Value nextVal = worklist.pop_back_val();
     if (!state.isWritable(nextVal)) {
       if (state.getOptions().printConflicts)
         annotateNonWritableTensor(nextVal);
       return true;
     }
 
     // If `nextVal` is not a BlockArgument: End of use-def chain reached.
     auto opResult = nextVal.dyn_cast<OpResult>();
     if (!opResult)
       continue;
 
     // Follow reverse SSA use-def chain.
     AliasingOpOperandList aliasingOpOperands =
         state.getAliasingOpOperands(opResult);
     for (OpOperand *opOperand : aliasingOpOperands)
-      if (aliasInfo.isInPlace(*opOperand) || currentOpOperand == opOperand)
+      if (state.isInPlace(*opOperand) || currentOpOperand == opOperand)
         worklist.push_back(opOperand->get());
   }
   return false;
 }
 
 /// Return true if bufferizing `operand` inplace would create a write to a
 /// non-writable buffer.
-static bool wouldCreateWriteToNonWritableBuffer(
-    OpOperand &operand, const BufferizationAliasInfo &aliasInfo,
-    OneShotAnalysisState &state, bool checkConsistencyOnly = false) {
+static bool
+wouldCreateWriteToNonWritableBuffer(OpOperand &operand,
+                                    OneShotAnalysisState &state,
+                                    bool checkConsistencyOnly = false) {
   // Collect writes of all aliases of OpOperand and OpResult.
   DenseSet<OpOperand *> usesWrite;
-  getAliasingInplaceWrites(usesWrite, operand.get(), aliasInfo, state);
+  getAliasingInplaceWrites(usesWrite, operand.get(), state);
   for (OpResult result : state.getAliasingOpResults(operand)) {
-    getAliasingInplaceWrites(usesWrite, result, aliasInfo, state);
+    getAliasingInplaceWrites(usesWrite, result, state);
   }
   if (!checkConsistencyOnly && state.bufferizesToMemoryWrite(operand))
     usesWrite.insert(&operand);
 
   // Assuming that `operand` bufferizes in-place: For each write (to each
   // alias), check if there is a non-writable tensor in the reverse SSA use-def
   // chain.
   for (OpOperand *uWrite : usesWrite) {
-    if (hasPrecedingAliasingNonWritableTensor(uWrite->get(), &operand,
-                                              aliasInfo, state)) {
+    if (hasPrecedingAliasingNonWritableTensor(uWrite->get(), &operand, state)) {
       LLVM_DEBUG(llvm::dbgs() << "=> NOT WRITABLE\n");
       return true;
     }
   }
 
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization analyses.
 //===----------------------------------------------------------------------===//
 
 /// Determine if `operand` can be bufferized in-place.
-static LogicalResult bufferizableInPlaceAnalysisImpl(
-    OpOperand &operand, BufferizationAliasInfo &aliasInfo,
-    OneShotAnalysisState &state, const DominanceInfo &domInfo) {
+static LogicalResult
+bufferizableInPlaceAnalysisImpl(OpOperand &operand, OneShotAnalysisState &state,
+                                const DominanceInfo &domInfo) {
   LLVM_DEBUG(
       llvm::dbgs() << "//===-------------------------------------------===//\n"
                    << "Analyzing operand #" << operand.getOperandNumber()
                    << " of " << *operand.getOwner() << "\n");
 
   bool foundInterference =
-      wouldCreateWriteToNonWritableBuffer(operand, aliasInfo, state) ||
-      wouldCreateReadAfterWriteInterference(operand, domInfo, state, aliasInfo);
+      wouldCreateWriteToNonWritableBuffer(operand, state) ||
+      wouldCreateReadAfterWriteInterference(operand, domInfo, state);
 
   if (foundInterference)
-    aliasInfo.bufferizeOutOfPlace(operand);
+    state.bufferizeOutOfPlace(operand);
   else
-    aliasInfo.bufferizeInPlace(operand, state);
+    state.bufferizeInPlace(operand);
 
   LLVM_DEBUG(llvm::dbgs()
              << "//===-------------------------------------------===//\n");
   return success();
 }
 
 /// Analyze the `ops` to determine which OpOperands are inplaceable. Walk ops in
 /// reverse and bufferize ops greedily. This is a good starter heuristic.

 ///
 /// When bufferized inplace, an ExtractSliceOp does not by itself create any
 /// read or write from memory. Instead, it has the effect of merging the alias
 /// sets of the source and the result buffers.
 ///
 /// An analysis is required to ensure inplace bufferization would not result in
 /// RaW dependence violations.
 static LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
-                                     BufferizationAliasInfo &aliasInfo,
                                      OneShotAnalysisState &state,
                                      const DominanceInfo &domInfo,
                                      unsigned analysisFuzzerSeed = 0) {
   if (analysisFuzzerSeed) {
     // This is a fuzzer. For testing purposes only. Randomize the order in which
     // operations are analyzed. The bufferization quality is likely worse, but
     // we want to make sure that no assertions are triggered anywhere.
     std::mt19937 g(analysisFuzzerSeed);
     llvm::shuffle(ops.begin(), ops.end(), g);
   }
 
   // Analyze a single op.
   auto analyzeOp = [&](Operation *op) {
     for (OpOperand &opOperand : op->getOpOperands())
       if (opOperand.get().getType().isa<TensorType>())
-        if (failed(bufferizableInPlaceAnalysisImpl(opOperand, aliasInfo, state,
-                                                   domInfo)))
+        if (failed(bufferizableInPlaceAnalysisImpl(opOperand, state, domInfo)))
           return failure();
     return success();
   };
 
   OneShotBufferizationOptions::AnalysisHeuristic heuristic =
       state.getOptions().analysisHeuristic;
   if (heuristic == OneShotBufferizationOptions::AnalysisHeuristic::BottomUp) {
     // Default: Walk ops in reverse for better interference analysis.

 static bool hasTensorSemantics(Operation *op) {
   bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
   bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
   return hasTensorResult || hasTensorOperand;
 }
 
 /// Analyze all ops that are contained in `op`.
 static LogicalResult inPlaceAnalysis(Operation *op,
-                                     BufferizationAliasInfo &aliasInfo,
                                      OneShotAnalysisState &state,
                                      const DominanceInfo &domInfo,
                                      unsigned analysisFuzzerSeed = 0) {
   // Collect ops so we can build our own reverse traversal.
   SmallVector<Operation *> ops;
   op->walk([&](Operation *op) {
     // No tensors => no buffers.
     if (!hasTensorSemantics(op))
       return;
     ops.push_back(op);
   });
 
-  return inPlaceAnalysis(ops, aliasInfo, state, domInfo, analysisFuzzerSeed);
+  return inPlaceAnalysis(ops, state, domInfo, analysisFuzzerSeed);
 }
 
 /// Analyze equivalence of tied OpResult/OpOperand pairs of the given ops.
 static void equivalenceAnalysis(SmallVector<Operation *> &ops,
-                                BufferizationAliasInfo &aliasInfo,
-                                AnalysisState &state) {
+                                OneShotAnalysisState &state) {
   for (Operation *op : ops)
     if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op))
       for (OpResult opResult : op->getOpResults())
         if (opResult.getType().isa<TensorType>())
           for (OpOperand *opOperand :
                bufferizableOp.getAliasingOpOperands(opResult, state))
             if (state.isInPlace(*opOperand))
               if (bufferizableOp.bufferRelation(opResult, state) ==
                   BufferRelation::Equivalent)
-                aliasInfo.unionEquivalenceClasses(opResult, opOperand->get());
+                state.unionEquivalenceClasses(opResult, opOperand->get());
 }
 
 /// Analyze equivalence of tied OpResult/OpOperand pairs of all ops contained
 /// in `op`.
-static void equivalenceAnalysis(Operation *op,
-                                BufferizationAliasInfo &aliasInfo,
-                                AnalysisState &state) {
+static void equivalenceAnalysis(Operation *op, OneShotAnalysisState &state) {
   // Traverse ops in PostOrder: Nested ops first, then enclosing ops.
   SmallVector<Operation *> ops;
   op->walk<WalkOrder::PostOrder>([&](Operation *op) {
     // No tensors => no buffers.
     if (none_of(op->getResultTypes(), isaTensor))
       return;
     ops.push_back(op);
   });
 
-  equivalenceAnalysis(ops, aliasInfo, state);
+  equivalenceAnalysis(ops, state);
 }
 
 /// Assert that the current bufferization decisions are consistent.
-static LogicalResult
-checkAliasInfoConsistency(Operation *op, const DominanceInfo &domInfo,
-                          AnalysisState &state,
-                          const BufferizationAliasInfo &aliasInfo) {
+static LogicalResult checkAliasInfoConsistency(Operation *op,
+                                               const DominanceInfo &domInfo,
+                                               OneShotAnalysisState &state) {
   const BufferizationOptions &options = state.getOptions();
 
   WalkResult walkResult = op->walk([&](BufferizableOpInterface op) {
     // Skip ops that are not in the filter.
     if (!options.isOpAllowed(op.getOperation()))
       return WalkResult::advance();
 
     // Input IR may not contain any ToMemrefOps. These are not supported because
     // the analysis cannot follow the data flow through memrefs.
     if (isa<ToMemrefOp>(op.getOperation())) {
       op->emitError("to_memref ops not supported during One-Shot Analysis");
       return WalkResult::interrupt();
     }
 
     for (OpOperand &opOperand : op->getOpOperands()) {
       if (opOperand.get().getType().isa<TensorType>()) {
         if (wouldCreateReadAfterWriteInterference(
-                opOperand, domInfo, state, aliasInfo,
+                opOperand, domInfo, state,
                 /*checkConsistencyOnly=*/true)) {
           // This error can happen if certain "mustBufferizeInPlace" interface
           // methods are implemented incorrectly, such that the IR already has
           // a RaW conflict before making any bufferization decisions.
           op->emitError("input IR has RaW conflict");
           return WalkResult::interrupt();
         }
       }
     }
 
     return WalkResult::advance();
   });
 
   return success(!walkResult.wasInterrupted());
 }
 
 /// Annotate the IR with the result of the analysis. For testing/debugging only.
 static void
 annotateOpsWithBufferizationMarkers(Operation *op,
-                                    const BufferizationAliasInfo &aliasInfo,
-                                    const BufferizationOptions &options) {
+                                    const OneShotAnalysisState &state) {
   // Add __inplace_operands_attr__.
   op->walk([&](Operation *op) {
     for (OpOperand &opOperand : op->getOpOperands())
       if (opOperand.get().getType().isa<TensorType>())
-        setInPlaceOpOperand(opOperand, aliasInfo.isInPlace(opOperand));
+        setInPlaceOpOperand(opOperand, state.isInPlace(opOperand));
   });
 }
 
 /// Assert that every allocation can be deallocated in the same block. I.e.,
 /// every value that is returned or yielded from a block is:
 /// * guaranteed to be aliasing a bbArg of that block or a parent block, or
 /// * guaranteed to be aliasing an OpResult of a op in a parent block.
 ///

 /// ```
 ///
 /// In the above example, the second scf.yield op is problematic because the
 /// yielded value %t is defined in the same block as the scf.yield op and
 /// and bufferizes to a new allocation.
 // TODO: Remove buffer deallocation from One-Shot Bufferize and fix the buffer
 // deallocation pass.
 static LogicalResult assertNoAllocsReturned(Operation *op,
-                                            const BufferizationOptions &options,
-                                            BufferizationAliasInfo &aliasInfo) {
+                                            const OneShotAnalysisState &state) {
   LogicalResult status = success();
   DominanceInfo domInfo(op);
   op->walk([&](Operation *returnOp) {
-    if (!isRegionReturnLike(returnOp) || !options.isOpAllowed(returnOp))
+    if (!isRegionReturnLike(returnOp) ||
+        !state.getOptions().isOpAllowed(returnOp))
       return WalkResult::advance();
 
     for (OpOperand &returnValOperand : returnOp->getOpOperands()) {
       Value returnVal = returnValOperand.get();
       // Skip non-tensor values.
       if (!returnVal.getType().isa<TensorType>())
         continue;
 
       bool foundEquivValue = false;
-      aliasInfo.applyOnEquivalenceClass(returnVal, [&](Value equivVal) {
+      state.applyOnEquivalenceClass(returnVal, [&](Value equivVal) {
         if (auto bbArg = equivVal.dyn_cast<BlockArgument>()) {
           Operation *definingOp = bbArg.getOwner()->getParentOp();
           if (definingOp->isProperAncestor(returnOp))
             foundEquivValue = true;
           return;
         }
 
         Operation *definingOp = equivVal.getDefiningOp();
static LogicalResult assertNoAllocsReturned(Operation *op,
 
   return status;
 }
 
 LogicalResult bufferization::analyzeOp(Operation *op,
                                        OneShotAnalysisState &state,
                                        BufferizationStatistics *statistics) {
   DominanceInfo domInfo(op);
-  BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
   const OneShotBufferizationOptions &options = state.getOptions();
 
-  if (failed(checkAliasInfoConsistency(op, domInfo, state, aliasInfo)))
+  if (failed(checkAliasInfoConsistency(op, domInfo, state)))
     return failure();
 
   // If the analysis fails, just return.
-  if (failed(inPlaceAnalysis(op, aliasInfo, state, domInfo,
-                             options.analysisFuzzerSeed)))
+  if (failed(inPlaceAnalysis(op, state, domInfo, options.analysisFuzzerSeed)))
     return failure();
 
   if (statistics) {
-    statistics->numTensorInPlace = aliasInfo.getStatNumTensorInPlace();
-    statistics->numTensorOutOfPlace = aliasInfo.getStatNumTensorOutOfPlace();
+    statistics->numTensorInPlace = state.getStatNumTensorInPlace();
+    statistics->numTensorOutOfPlace = state.getStatNumTensorOutOfPlace();
   }
 
-  equivalenceAnalysis(op, aliasInfo, state);
+  equivalenceAnalysis(op, state);
 
   bool failedAnalysis = false;
   if (!options.allowReturnAllocs)
-    failedAnalysis |= failed(assertNoAllocsReturned(op, options, aliasInfo));
+    failedAnalysis |= failed(assertNoAllocsReturned(op, state));
 
   // Gather some extra analysis data.
   state.gatherYieldedTensors(op);
   state.gatherUndefinedTensorUses(op);
 
   // Analysis verification: After setting up alias/equivalence sets, each op
   // can check for expected invariants/limitations and fail the analysis if
   // necessary.
   op->walk([&](Operation *op) {
     if (BufferizableOpInterface bufferizableOp =
             options.dynCastBufferizableOp(op))
       failedAnalysis |= failed(bufferizableOp.verifyAnalysis(state));
   });
 
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly)
-    annotateOpsWithBufferizationMarkers(op, aliasInfo, options);
+    annotateOpsWithBufferizationMarkers(op, state);
 
   return success(!failedAnalysis);
 }
 
 LogicalResult
 bufferization::runOneShotBufferize(Operation *op,
                                    const OneShotBufferizationOptions &options,
                                    BufferizationStatistics *statistics) {

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

return dyn_cast_or_null<func::FuncOp>(
       SymbolTable::lookupNearestSymbolFrom(callOp, sym));
 }
 
 /// Gather equivalence info of CallOps.
 /// Note: This only adds new equivalence info if the called function was already
 /// analyzed.
 // TODO: This does not handle cyclic function call graphs etc.
 static void equivalenceAnalysis(func::FuncOp funcOp,
-                                BufferizationAliasInfo &aliasInfo,
                                 OneShotAnalysisState &state,
                                 FuncAnalysisState &funcState) {
   funcOp->walk([&](func::CallOp callOp) {
     func::FuncOp calledFunction = getCalledFunction(callOp);
     assert(calledFunction && "could not retrieved called func::FuncOp");
 
     // No equivalence info available for the called function.
     if (!funcState.equivalentFuncArgs.count(calledFunction))
       return WalkResult::skip();
 
     for (auto it : funcState.equivalentFuncArgs[calledFunction]) {
       int64_t returnIdx = it.first;
       int64_t bbargIdx = it.second;
       if (!state.isInPlace(callOp->getOpOperand(bbargIdx)))
         continue;
       Value returnVal = callOp.getResult(returnIdx);
       Value argVal = callOp->getOperand(bbargIdx);
-      aliasInfo.unionEquivalenceClasses(returnVal, argVal);
+      state.unionEquivalenceClasses(returnVal, argVal);
     }
 
     return WalkResult::advance();
   });
 }
 
 /// Store all functions of the `moduleOp` in `orderedFuncOps`, sorted by
 /// callee-caller order (i.e. callees without callers first).

 
 LogicalResult
 mlir::bufferization::analyzeModuleOp(ModuleOp moduleOp,
                                      OneShotAnalysisState &state,
                                      BufferizationStatistics *statistics) {
   assert(state.getOptions().bufferizeFunctionBoundaries &&
          "expected that function boundary bufferization is activated");
   FuncAnalysisState &funcState = getOrCreateFuncAnalysisState(state);
-  BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
 
   // A list of functions in the order in which they are analyzed + bufferized.
   SmallVector<func::FuncOp> orderedFuncOps;
 
   // A mapping of FuncOps to their callers.
   FuncCallerMap callerMap;
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
   // Analyze ops.
   for (func::FuncOp funcOp : orderedFuncOps) {
     if (!state.getOptions().isOpAllowed(funcOp))
       continue;
 
     // Now analyzing function.
     funcState.startFunctionAnalysis(funcOp);
 
     // Gather equivalence info for CallOps.
-    equivalenceAnalysis(funcOp, aliasInfo, state, funcState);
+    equivalenceAnalysis(funcOp, state, funcState);
 
     // Analyze funcOp.
     if (failed(analyzeOp(funcOp, state, statistics)))
       return failure();
 
     // Run some extra function analyses.
     if (failed(aliasingFuncOpBBArgsAnalysis(funcOp, state, funcState)) ||
         failed(funcOpBbArgReadWriteAnalysis(funcOp, state, funcState)))