[mlir][linalg][bufferize] Add PostAnalysisStep

[springerm]

Summary

This helper struct allows users of ComprehensiveBufferize to inject "post analysis" steps that are implemented after the analysis but before the bufferization.

Depends On D113457

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h

 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/SetVector.h"
 
 namespace mlir {
 class BlockAndValueMapping;
+class DominanceInfo;
+class FuncOp;
 
 namespace linalg {
 namespace comprehensive_bufferize {
 
 class BufferizationAliasInfo;
 
 /// Specify fine-grain relationship between buffers to enable more analysis.
 enum class BufferRelation {
struct BufferizationState {
   BlockAndValueMapping &tensorToBufferMap;
 };
 
 /// Return the result buffer (memref) for a given OpResult (tensor). Allocate
 /// a new buffer and copy over data from the existing buffer if out-of-place
 /// bufferization is necessary.
 Value getResultBuffer(OpBuilder &b, OpResult result, BufferizationState &state);
 
+/// PostAnalysisSteps can be registered with `BufferizationOptions` and are
+/// executed after the analysis, but before bufferization. They can be used
+/// implement custom dialect-specific optimizations.
+struct PostAnalysisStep {
+  virtual ~PostAnalysisStep() {}
+
+  /// Run the post analysis step. This function may modify the IR, but must keep
+  /// `aliasInfo` consistent. Newly created operations and operations that
+  /// should be re-analyzed must be stored in `newOps`.
+  virtual LogicalResult run(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
+                            DominanceInfo &domInfo,
+                            SmallVector<Operation *> &newOps) = 0;
+};
+
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h.inc"
 
 namespace mlir {
 namespace linalg {

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h

 /// non-null if `op` is a CallOpInterface (resp. GlobalCreator).
 LogicalResult
 bufferizeOp(Operation *op, BufferizationState &state,
             DenseMap<FuncOp, FunctionType> *bufferizedFunctionTypes = nullptr);
 
 /// Register external models implemented for the `BufferizableOpInterface`.
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
+struct BufferizationOptions {
+  BufferizationOptions();
+
+  /// Register a "post analysis" step. Such steps are executed after the
+  /// analysis, but before bufferization.
+  template <typename Step, typename... Args>
+  void addPostAnalysisStep(Args... args) {
+    postAnalysisSteps.emplace_back(
+        std::make_unique<Step>(std::forward<Args>(args)...));
+  }
+
+  std::unique_ptr<AllocationCallbacks> allocationFns;
+  bool allowReturnMemref = false;
+  unsigned analysisFuzzerSeed = 0;
+  bool testAnalysisOnly = false;
+  std::vector<std::unique_ptr<PostAnalysisStep>> postAnalysisSteps;
+};
+
+LogicalResult runComprehensiveBufferize(ModuleOp moduleOp,
+                                        const BufferizationOptions &options);
+
+namespace linalg_ext {
+
+struct InitTensorEliminationStep : public PostAnalysisStep {
   /// Try to eliminate InitTensorOps inside `funcOp`.
   ///
   /// * `rewriteFunc` generates the replacement for the InitTensorOp.
   /// * Only InitTensorOps that are anchored on a matching OpOperand as per
-///   `anchorMatchFunc` are considered. "Anchored" means that there is a path on
-///   the reverse SSA use-def chain, starting from the OpOperand and always
+  ///   `anchorMatchFunc` are considered. "Anchored" means that there is a path
+  ///   on the reverse SSA use-def chain, starting from the OpOperand and always
   ///   following the aliasing  OpOperand, that eventually ends at a single
   ///   InitTensorOp.
   /// * The result of `rewriteFunc` must usually be analyzed for inplacability.
   ///   This analysis can be skipped with `skipAnalysis`.
-LogicalResult initTensorElimination(
+  LogicalResult eliminateInitTensors(
       FuncOp funcOp, BufferizationAliasInfo &aliasInfo, DominanceInfo &domInfo,
       std::function<bool(OpOperand &)> anchorMatchFunc,
       std::function<Value(OpBuilder &, Location, OpOperand &)> rewriteFunc,
-    bool skipAnalysis = false);
+      SmallVector<Operation *> &newOps);
+};
 
 /// Try to eliminate InitTensorOps inside funcOp that are anchored on an
 /// InsertSliceOp, i.e., if it is eventually inserted into another tensor
 /// (and some other conditions are met).
-LogicalResult eliminateInsertSliceAnchoredInitTensorOps(
-    FuncOp funcOp, BufferizationAliasInfo &aliasInfo, DominanceInfo &domInfo);
-
-struct BufferizationOptions {
-  BufferizationOptions();
-
-  std::unique_ptr<AllocationCallbacks> allocationFns;
-  bool allowReturnMemref = false;
-  unsigned analysisFuzzerSeed = 0;
-  bool testAnalysisOnly = false;
+struct InsertSliceAnchoredInitTensorEliminationStep
+    : public InitTensorEliminationStep {
+  LogicalResult run(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
+                    DominanceInfo &domInfo,
+                    SmallVector<Operation *> &newOps) override;
 };
 
-LogicalResult runComprehensiveBufferize(ModuleOp moduleOp,
-                                        const BufferizationOptions &options);
+} // namespace linalg_ext
 
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVE_BUFFERIZE_H

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp

static void layoutPostProcessing(ModuleOp moduleOp) {
   }
 }
 
 /// Try to eliminate InitTensorOps inside funcOp. An InitTensorOp is replaced
 /// with the the result of `rewriteFunc` if it is anchored on a matching
 /// OpOperand. "Anchored" means that there is a path on the reverse SSA use-def
 /// chain, starting from the OpOperand and always following the aliasing
 /// OpOperand, that eventually ends at a single InitTensorOp.
-LogicalResult mlir::linalg::comprehensive_bufferize::initTensorElimination(
-    FuncOp funcOp, BufferizationAliasInfo &aliasInfo, DominanceInfo &domInfo,
+LogicalResult mlir::linalg::comprehensive_bufferize::linalg_ext::
+    InitTensorEliminationStep::eliminateInitTensors(
+        FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
+        DominanceInfo &domInfo,
         std::function<bool(OpOperand &)> anchorMatchFunc,
         std::function<Value(OpBuilder &, Location, OpOperand &)> rewriteFunc,
-    bool skipAnalysis) {
+        SmallVector<Operation *> &newOps) {
   OpBuilder b(funcOp->getContext());
 
   WalkResult status = funcOp->walk([&](Operation *op) {
     for (OpOperand &operand : op->getOpOperands()) {
       // Is this a matching OpOperand?
       if (!anchorMatchFunc(operand))
         continue;
 
for (OpOperand &operand : op->getOpOperands()) {
 
       // Uses of the InitTensorOp are replaced here, but the op is not deleted.
       // InitTensorOps without uses are ignored by the bufferization.
       initTensor.replaceAllUsesWith(replacement);
       aliasInfo.createAliasInfoEntry(replacement);
       aliasInfo.unionAliasSets(initTensor, replacement);
       aliasInfo.unionEquivalenceClasses(initTensor, replacement);
 
-      // Run analysis on the newly created op.
-      if (auto opResult = replacement.dyn_cast<OpResult>()) {
-        if (!skipAnalysis) {
-          SmallVector<Operation *> ops(1, replacement.getDefiningOp());
-          if (failed(inPlaceAnalysis(ops, aliasInfo, domInfo)))
-            return WalkResult::interrupt();
-        }
-      }
+      // Register replacement ops.
+      if (Operation *newOp = replacement.getDefiningOp())
+        newOps.push_back(newOp);
     }
 
     // Advance to the next operation.
     return WalkResult::advance();
   });
 
   return failure(status.wasInterrupted());
 }

 /// * The InsertSliceOp was decided to bufferize inplace.
 /// * On the reverse use-def chain path from the InsertSliceOp to the
 ///   InitTensorOp, all ops were decided to bufferize inplace and the buffer
 ///   relation is "equivalent" (TODO: can be relaxed if needed).
 /// * The reverse use-def chain has exactly one end, which is the InitTensorOp.
 ///
 /// Note that the newly inserted ExtractSliceOp may have to bufferize
 /// out-of-place due to RaW conflicts.
-LogicalResult mlir::linalg::comprehensive_bufferize::
-    eliminateInsertSliceAnchoredInitTensorOps(FuncOp funcOp,
-                                              BufferizationAliasInfo &aliasInfo,
-                                              DominanceInfo &domInfo) {
-  return initTensorElimination(
+LogicalResult mlir::linalg::comprehensive_bufferize::linalg_ext::
+    InsertSliceAnchoredInitTensorEliminationStep::run(
+        FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
+        DominanceInfo &domInfo, SmallVector<Operation *> &newOps) {
+  return eliminateInitTensors(
       funcOp, aliasInfo, domInfo,
       [&](OpOperand &operand) {
         auto insertSliceOp = dyn_cast<InsertSliceOp>(operand.getOwner());
         if (!insertSliceOp)
           return false;
         // Only inplace bufferized InsertSliceOps are eligible.
         if (!aliasInfo.isInPlace(insertSliceOp->getOpResult(0)))
           return false;
         return &operand == &insertSliceOp->getOpOperand(0) /*source*/;
       },
       [](OpBuilder &b, Location loc, OpOperand &operand) {
         auto insertSliceOp = cast<InsertSliceOp>(operand.getOwner());
         auto extractOp = b.create<tensor::ExtractSliceOp>(
             loc, insertSliceOp.dest(), insertSliceOp.getMixedOffsets(),
             insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
         return extractOp.result();
-      });
+      },
+      newOps);
 }
 
 #ifndef NDEBUG
 /// Assert that the current bufferization decisions are consistent.
 static void checkAliasInfoConsistency(FuncOp funcOp,
                                       const DominanceInfo &domInfo,
                                       const BufferizationAliasInfo &aliasInfo) {
   funcOp.walk([&](Operation *op) {
#ifndef NDEBUG
     checkAliasInfoConsistency(funcOp, domInfo, aliasInfo);
 #endif // NDEBUG
 
     // If the analysis fails, just return.
     if (failed(inPlaceAnalysisFuncOpBody(funcOp, aliasInfo, domInfo,
                                          options.analysisFuzzerSeed)))
       return failure();
 
-    // Try to eliminate InitTensorOps to avoid new allocations during the
-    // bufferization phase.
-    if (failed(eliminateInsertSliceAnchoredInitTensorOps(funcOp, aliasInfo,
-                                                         domInfo)))
+    for (const std::unique_ptr<PostAnalysisStep> &step :
+         options.postAnalysisSteps) {
+      SmallVector<Operation *> newOps;
+      if (failed(step->run(funcOp, aliasInfo, domInfo, newOps)))
         return failure();
+      // Analyze ops that were created by the PostAnalysisStep.
+      if (failed(inPlaceAnalysis(newOps, aliasInfo, domInfo)))
+        return failure();
+    }
 
     // Bufferization phase.
     if (!options.testAnalysisOnly) {
       BlockAndValueMapping tensorToBufferMap;
       BufferizationState state(aliasInfo, *options.allocationFns,
                                tensorToBufferMap);
       if (failed(
               bufferizeFuncOpInternals(funcOp, state, bufferizedFunctionTypes)))

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp

options.allocationFns->memCpyFn = [](OpBuilder &b, Location loc, Value from,
                                        Value to) {
     b.create<linalg::CopyOp>(loc, from, to);
   };
 
   options.allowReturnMemref = allowReturnMemref;
   options.analysisFuzzerSeed = analysisFuzzerSeed;
   options.testAnalysisOnly = testAnalysisOnly;
 
+  // Enable InitTensorOp elimination.
+  options.addPostAnalysisStep<
+      linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
+
   ModuleOp moduleOp = getOperation();
   applyEnablingTransformations(moduleOp);
 
   if (failed(runComprehensiveBufferize(moduleOp, options))) {
     signalPassFailure();
     return;
   }