[mlir][bufferize] Add notifications for AnalysisState extensions

[springerm]

Summary

Notify extensions when an OpOperand was decided to bufferize in-place or out-of-place.

This is in preparation of adding value privatizations in a subsequent change.

Depends On D135051

[springerm]

rebase

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

public:
   /// 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);
+  /// Mark the given OpOperand as in-place.
+  void bufferizeInPlace(OpOperand &operand);
 
-  /// Set the inPlace bufferization spec to false.
+  /// Mark the given OpOperand as out-of-place.
   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.
public:
   }
 
   /// 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;
 
 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 {
public:
   OneShotAnalysisState(const OneShotAnalysisState &) = delete;
 
   ~OneShotAnalysisState() override = default;
 
   static bool classof(const AnalysisState *base) {
     return base->getType() == TypeID::get<OneShotAnalysisState>();
   }
 
+  /// Mark the given OpOperand as in-place. Also merge result's and operand's
+  /// aliasing sets.
+  void bufferizeInPlace(OpOperand &operand);
+
+  /// Mark the given OpOperand as out-of-place.
+  void bufferizeOutOfPlace(OpOperand &operand);
+
   /// Return a reference to the BufferizationAliasInfo.
   BufferizationAliasInfo &getAliasInfo() { return aliasInfo; }
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   bool isInPlace(OpOperand &opOperand) const override;
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   bool areEquivalentBufferizedValues(Value v1, Value v2) const override;
public:
 
   protected:
     /// Constructs an extension of the given TransformState object.
     Extension(OneShotAnalysisState &state) : state(state) {}
 
     /// Provides read-only access to the parent OneShotAnalysisState object.
     const OneShotAnalysisState &getAnalysisState() const { return state; }
 
+    /// Notify this extension of an in-place bufferization decision.
+    virtual void notifyBufferizeInPlace(OpOperand &operand) {}
+
+    /// Notify this extension of an out-of-place bufferization decision.
+    virtual void notifyBufferizeOutOfPlace(OpOperand &operand) {}
+
   private:
     /// Back-reference to the state that is being extended.
     OneShotAnalysisState &state;
   };
 
   /// Adds a new Extension of the type specified as template parameter,
   /// constructing it with the arguments provided. The extension is owned by the
   /// OneShotAnalysisState. It is expected that the state does not already have

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

void BufferizationAliasInfo::insertNewBufferEquivalence(Value newValue,
   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) {
-  markInPlace(operand);
-  for (OpResult result : state.getAliasingOpResult(operand))
-    aliasInfo.unionSets(result, operand.get());
+void BufferizationAliasInfo::bufferizeInPlace(OpOperand &operand) {
+  inplaceBufferized.insert(&operand);
 }
 
-/// Set the inPlace bufferization spec to false.
 void BufferizationAliasInfo::bufferizeOutOfPlace(OpOperand &operand) {
   assert(!inplaceBufferized.contains(&operand) &&
          "OpOperand was already decided to bufferize inplace");
 }
 
 /// Apply `fun` to all the members of the equivalence class of `v`.
 void BufferizationAliasInfo::applyOnEquivalenceClass(
     Value v, function_ref<void(Value)> fun) const {
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)) {
           for (OpResult opResult :
                bufferizableOp.getAliasingOpResult(opOperand, *this))
             aliasInfo.unionAliasSets(opOperand.get(), opResult);
-          aliasInfo.markInPlace(opOperand);
+          aliasInfo.bufferizeInPlace(opOperand);
         }
     }
     return WalkResult::advance();
   });
 }
 
+void OneShotAnalysisState::bufferizeInPlace(OpOperand &operand) {
+  aliasInfo.bufferizeInPlace(operand);
+  for (OpResult result : getAliasingOpResult(operand))
+    aliasInfo.unionAliasSets(result, operand.get());
+  for (auto &it : extensions)
+    it.getSecond()->notifyBufferizeInPlace(operand);
+}
+
+void OneShotAnalysisState::bufferizeOutOfPlace(OpOperand &operand) {
+  aliasInfo.bufferizeOutOfPlace(operand);
+  for (auto &it : extensions)
+    it.getSecond()->notifyBufferizeOutOfPlace(operand);
+}
+
 bool OneShotAnalysisState::isInPlace(OpOperand &opOperand) const {
   return aliasInfo.isInPlace(opOperand);
 }
 
 bool OneShotAnalysisState::areEquivalentBufferizedValues(Value v1,
                                                          Value v2) const {
   return aliasInfo.areEquivalentBufferizedValues(v1, v2);
 }

 static LogicalResult bufferizableInPlaceAnalysisImpl(
     OpOperand &operand, BufferizationAliasInfo &aliasInfo,
     OneShotAnalysisState &state, const DominanceInfo &domInfo) {
   bool foundInterference =
       wouldCreateWriteToNonWritableBuffer(operand, aliasInfo, state) ||
       wouldCreateReadAfterWriteInterference(operand, domInfo, state, aliasInfo);
 
   if (foundInterference)
-    aliasInfo.bufferizeOutOfPlace(operand);
+    state.bufferizeOutOfPlace(operand);
   else
-    aliasInfo.bufferizeInPlace(operand, state);
+    state.bufferizeInPlace(operand);
 
   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.
 ///
 /// Even if an op does not read or write, it may still create an alias when