[mlir][bufferization][NFC] Remove post-analysis step infrastructure

[springerm]

Summary

Now that analysis and bufferization are better separated, post-analysis steps are no longer needed. Users can directly interleave analysis and bufferization as needed.

Depends On D126570

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

public:
   SetVector<Value> findLastPrecedingWrite(Value value) const;
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   virtual bool isInPlace(OpOperand &opOperand) const = 0;
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   virtual bool areEquivalentBufferizedValues(Value v1, Value v2) const = 0;
 
+  /// Return true if `v1` and `v2` may bufferize to aliasing buffers.
+  virtual bool areAliasingBufferizedValues(Value v1, Value v2) const = 0;
+
   /// Return `true` if the given tensor has undefined contents.
   virtual bool hasUndefinedContents(OpOperand *opOperand) const = 0;
 
   /// 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.
   ///
   /// Note: In the absence of an analysis, an implementation may return true for
   /// any given tensor.

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

 
 namespace mlir {
 namespace bufferization {
 
 struct OneShotBufferizationOptions;
 class BufferizationAliasInfo;
 class OneShotAnalysisState;
 
-/// PostAnalysisStepFns can be registered with `BufferizationOptions` and are
-/// executed after the analysis, but before bufferization. They can be used to
-/// implement custom dialect-specific optimizations. They may modify the IR, but
-/// must keep `aliasInfo` consistent. Newly created operations and operations
-/// that should be re-analyzed must be added to `newOps`.
-using PostAnalysisStepFn = std::function<LogicalResult(
-    Operation *, AnalysisState &, BufferizationAliasInfo &,
-    SmallVector<Operation *> &)>;
-
-using PostAnalysisStepList = SmallVector<PostAnalysisStepFn>;
-
 /// Options for analysis-enabled bufferization.
 struct OneShotBufferizationOptions : public BufferizationOptions {
   OneShotBufferizationOptions() = default;
 
-  /// Register a "post analysis" step. Such steps are executed after the
-  /// analysis, but before bufferization.
-  void addPostAnalysisStep(PostAnalysisStepFn fn) {
-    postAnalysisSteps.push_back(fn);
-  }
-
-  /// Registered post analysis steps.
-  PostAnalysisStepList postAnalysisSteps;
-
   /// Specifies whether returning newly allocated memrefs should be allowed.
   /// Otherwise, a pass failure is triggered.
   bool allowReturnAllocs = false;
 
   /// Specifies whether buffer return values that are equivalent to a FuncOp
   /// bbArg should be dropped.
   bool dropEquivalentFuncResults = true;
 };
public:
   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;
 
+  /// 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;
 
   /// 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;
 
   /// 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);
 
+  /// 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;
+
 private:
   /// `aliasInfo` keeps track of aliasing and equivalent values. Only internal
   /// functions and `runOneShotBufferize` may access this object.
   BufferizationAliasInfo aliasInfo;
 
   /// A set of all tensors (and maybe aliasing tensors) that yielded from a
   /// block.
   DenseSet<Value> yieldedTensors;

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

public:
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   bool areEquivalentBufferizedValues(Value v1, Value v2) const override {
     // There is no analysis, so we do not know if the values are equivalent. The
     // conservative answer is "false".
     return false;
   }
 
+  /// Return true if `v1` and `v2` may bufferize to aliasing buffers.
+  bool areAliasingBufferizedValues(Value v1, Value v2) const override {
+    // There is no analysis, so we do not know if the values are equivalent. The
+    // conservative answer is "true".
+    return true;
+  }
+
   /// Return `true` if the given tensor has undefined contents.
   bool hasUndefinedContents(OpOperand *opOperand) const override {
     // There is no analysis, so the conservative answer is "false".
     return false;
   }
 
   /// 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.

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

bool OneShotAnalysisState::isInPlace(OpOperand &opOperand) const {
   return aliasInfo.isInPlace(opOperand);
 }
 
 bool OneShotAnalysisState::areEquivalentBufferizedValues(Value v1,
                                                          Value v2) const {
   return aliasInfo.areEquivalentBufferizedValues(v1, v2);
 }
 
+bool OneShotAnalysisState::areAliasingBufferizedValues(Value v1,
+                                                       Value v2) const {
+  return aliasInfo.areAliasingBufferizedValues(v1, v2);
+}
+
 // 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.
 void OneShotAnalysisState::gatherYieldedTensors(Operation *op) {
   op->walk([&](Operation *returnOp) {
     if (!isRegionReturnLike(returnOp) || !getOptions().isOpAllowed(returnOp))
       return WalkResult::advance();

 bool OneShotAnalysisState::hasUndefinedContents(OpOperand *opOperand) const {
   return undefinedTensorUses.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) {
+    for (OpOperand &use : val.getUses())
+      if (isInPlace(use) && bufferizesToMemoryWrite(use))
+        isWritten = true;
+  });
+  return isWritten;
+}
+
 //===----------------------------------------------------------------------===//
 // Bufferization-specific alias analysis.
 //===----------------------------------------------------------------------===//
 
 /// Return true if opOperand has been decided to bufferize in-place.
 static bool isInplaceMemoryWrite(OpOperand &opOperand,
                                  const BufferizationAliasInfo &aliasInfo,
                                  AnalysisState &state) {
if (failed(checkAliasInfoConsistency(op, domInfo, state, aliasInfo)))
     return failure();
 
   // If the analysis fails, just return.
   if (failed(inPlaceAnalysis(op, aliasInfo, state, domInfo,
                              options.analysisFuzzerSeed)))
     return failure();
   equivalenceAnalysis(op, aliasInfo, state);
 
-  for (const PostAnalysisStepFn &fn : options.postAnalysisSteps) {
-    SmallVector<Operation *> newOps;
-    if (failed(fn(op, state, aliasInfo, newOps)))
-      return failure();
-    // Analyze ops that were created by the PostAnalysisStepFn.
-    if (failed(inPlaceAnalysis(newOps, aliasInfo, state, domInfo)))
-      return failure();
-    equivalenceAnalysis(newOps, aliasInfo, state);
-  }
-
   bool failedAnalysis = false;
   if (!options.allowReturnAllocs) {
     SmallVector<Operation *> newOps;
     failedAnalysis |=
         failed(assertDestinationPassingStyle(op, state, aliasInfo, newOps));
   }
 
   // Gather some extra analysis data.

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

 // implementations for FuncOp, CallOp and ReturnOp.
 //
 // Module Bufferization is run via `runOneShotModuleBufferize(ModuleOp, ...)`.
 // This function analyzes the given module and determines the order of analysis
 // and bufferization: Functions that are called are processed before their
 // respective callers.
 //
 // After analyzing a FuncOp, additional information about its bbArgs is
-// gathered through PostAnalysisStepFns and stored in `FuncAnalysisState`.
+// gathered and stored in `FuncAnalysisState`.
 //
 // * `aliasingFuncOpBBArgsAnalysis` determines the equivalent/aliasing bbArgs
 // for
 //   each tensor return value (if any).
 // * `funcOpBbArgReadWriteAnalysis` determines whether or not a tensor bbArg is
 //   read/written.
 //
 // Only tensors that are equivalent to some FuncOp bbArg may be returned.
static void annotateEquivalentReturnBbArg(OpOperand &returnVal,
   equivBbArgs[returnVal.getOperandNumber()] = bbArg.getArgNumber();
 
   OpBuilder b(op->getContext());
   op->setAttr(kEquivalentArgsAttr, b.getI64ArrayAttr(equivBbArgs));
 }
 
 /// Store function BlockArguments that are equivalent to/aliasing a returned
 /// value in FuncAnalysisState.
-static LogicalResult
-aliasingFuncOpBBArgsAnalysis(Operation *op, AnalysisState &state,
-                             BufferizationAliasInfo &aliasInfo,
-                             SmallVector<Operation *> &newOps) {
+static LogicalResult aliasingFuncOpBBArgsAnalysis(FuncOp funcOp,
+                                                  OneShotAnalysisState &state) {
   FuncAnalysisState &funcState = getFuncAnalysisState(state);
 
   // Support only single return-terminated block in the function.
-  auto funcOp = cast<func::FuncOp>(op);
   func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
   assert(returnOp && "expected func with single return op");
 
   for (OpOperand &returnVal : returnOp->getOpOperands())
     if (returnVal.get().getType().isa<RankedTensorType>())
       for (BlockArgument bbArg : funcOp.getArguments())
         if (bbArg.getType().isa<RankedTensorType>()) {
           int64_t returnIdx = returnVal.getOperandNumber();
           int64_t bbArgIdx = bbArg.getArgNumber();
-          if (aliasInfo.areEquivalentBufferizedValues(returnVal.get(), bbArg)) {
+          if (state.areEquivalentBufferizedValues(returnVal.get(), bbArg)) {
             funcState.equivalentFuncArgs[funcOp][returnIdx] = bbArgIdx;
             if (state.getOptions().testAnalysisOnly)
               annotateEquivalentReturnBbArg(returnVal, bbArg);
           }
-          if (aliasInfo.areAliasingBufferizedValues(returnVal.get(), bbArg)) {
+          if (state.areAliasingBufferizedValues(returnVal.get(), bbArg)) {
             funcState.aliasingFuncArgs[funcOp][returnIdx].push_back(bbArgIdx);
             funcState.aliasingReturnVals[funcOp][bbArgIdx].push_back(returnIdx);
           }
         }
 
   return success();
 }
 
-/// Return true if the buffer of the given tensor value is written to. Must not
-/// be called for values inside not yet analyzed functions. (Post-analysis
-/// steps do not have to be run yet, i.e., "in progress" is also OK.)
-static bool isValueWritten(Value value, const AnalysisState &state,
-                           const BufferizationAliasInfo &aliasInfo) {
-#ifndef NDEBUG
-  assert(value.getType().isa<TensorType>() && "expected TensorType");
-  func::FuncOp funcOp;
-  if (auto bbArg = value.dyn_cast<BlockArgument>()) {
-    Operation *owner = bbArg.getOwner()->getParentOp();
-    funcOp = isa<func::FuncOp>(owner) ? cast<func::FuncOp>(owner)
-                                      : owner->getParentOfType<func::FuncOp>();
-  } else {
-    funcOp = value.getDefiningOp()->getParentOfType<func::FuncOp>();
-  }
-  assert(getFuncOpAnalysisState(state, funcOp) !=
-             FuncOpAnalysisState::NotAnalyzed &&
-         "FuncOp must be fully analyzed or analysis in progress");
-#endif // NDEBUG
-
-  bool isWritten = false;
-  aliasInfo.applyOnAliases(value, [&](Value val) {
-    for (OpOperand &use : val.getUses())
-      if (state.isInPlace(use) && state.bufferizesToMemoryWrite(use))
-        isWritten = true;
-  });
-  return isWritten;
-}
-
 static void annotateFuncArgAccess(func::FuncOp funcOp, BlockArgument bbArg,
                                   bool isRead, bool isWritten) {
   OpBuilder b(funcOp.getContext());
   Attribute accessType;
   if (isRead && isWritten) {
     accessType = b.getStringAttr("read-write");
   } else if (isRead) {
     accessType = b.getStringAttr("read");
   } else if (isWritten) {
     accessType = b.getStringAttr("write");
   } else {
     accessType = b.getStringAttr("none");
   }
   funcOp.setArgAttr(bbArg.getArgNumber(), "bufferization.access", accessType);
 }
 
-/// Determine which FuncOp bbArgs are read and which are written. If this
-/// PostAnalysisStepFn is run on a function with unknown ops, it will
-/// conservatively assume that such ops bufferize to a read + write.
-static LogicalResult
-funcOpBbArgReadWriteAnalysis(Operation *op, AnalysisState &state,
-                             BufferizationAliasInfo &aliasInfo,
-                             SmallVector<Operation *> &newOps) {
+/// Determine which FuncOp bbArgs are read and which are written. When run on a
+/// function with unknown ops, we conservatively assume that such ops bufferize
+/// to a read + write.
+static LogicalResult funcOpBbArgReadWriteAnalysis(FuncOp funcOp,
+                                                  OneShotAnalysisState &state) {
   FuncAnalysisState &funcState = getFuncAnalysisState(state);
-  auto funcOp = cast<func::FuncOp>(op);
 
   // If the function has no body, conservatively assume that all args are
   // read + written.
   if (funcOp.getBody().empty()) {
     for (BlockArgument bbArg : funcOp.getArguments()) {
       funcState.readBbArgs[funcOp].insert(bbArg.getArgNumber());
       funcState.writtenBbArgs[funcOp].insert(bbArg.getArgNumber());
     }
 
     return success();
   }
 
   for (BlockArgument bbArg : funcOp.getArguments()) {
     if (!bbArg.getType().isa<TensorType>())
       continue;
     bool isRead = state.isValueRead(bbArg);
-    bool isWritten = isValueWritten(bbArg, state, aliasInfo);
+    bool isWritten = state.isValueWritten(bbArg);
     if (state.getOptions().testAnalysisOnly)
       annotateFuncArgAccess(funcOp, bbArg, isRead, isWritten);
     if (isRead)
       funcState.readBbArgs[funcOp].insert(bbArg.getArgNumber());
     if (isWritten)
       funcState.writtenBbArgs[funcOp].insert(bbArg.getArgNumber());
   }
 
LogicalResult mlir::bufferization::runOneShotModuleBufferize(
   SmallVector<func::FuncOp> orderedFuncOps;
 
   // A mapping of FuncOps to their callers.
   FuncCallerMap callerMap;
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
-  // Collect bbArg/return value information after the analysis.
-  options.addPostAnalysisStep(aliasingFuncOpBBArgsAnalysis);
-  options.addPostAnalysisStep(funcOpBbArgReadWriteAnalysis);
-
   // Analyze ops.
   for (func::FuncOp funcOp : orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.getBody().empty())
       continue;
 
     // Now analyzing function.
     funcState.startFunctionAnalysis(funcOp);
 
     // Gather equivalence info for CallOps.
     equivalenceAnalysis(funcOp, aliasInfo, funcState);
 
     // Analyze funcOp.
     if (failed(analyzeOp(funcOp, analysisState)))
       return failure();
 
+    // Run some extra function analyses.
+    if (failed(aliasingFuncOpBBArgsAnalysis(funcOp, analysisState)) ||
+        failed(funcOpBbArgReadWriteAnalysis(funcOp, analysisState)))
+      return failure();
+
     // Mark op as fully analyzed.
     funcState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::Analyzed;
 
     // Add annotations to function arguments.
     if (options.testAnalysisOnly)
       annotateOpsWithBufferizationMarkers(funcOp, analysisState);
   }