[mlir][linalg][bufferize] Move FuncOp boundary bufferization to ModuleBufferization

[springerm]

Summary

Depends On D114669

[nicolasvasilache]

Please mark the commit as NFC -

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

 // of ops that have well-defined inplace semantics.
 // This pass caters to high-performance codegen where buffer reuse is deemed
 // critical: the pass should fail if the bufferized form of the function needs
 // to return any buffer.
 // Generic control-flow and branching are unsupported.
 // Composability with extensible set of ops is not a first-class concern.
 //
 // Bufferization occurs by:
-//  a. performing an inPlace analysis `inPlaceAnalysisFuncOpBody`
-//     which marks each operation within the function with the
-//     `kInPlaceResultsAttrName` attribute.
-//  b. traversing each operation in the function and rewriting it in
+//  a. performing an inPlace analysis `inPlaceAnalysis` which marks each
+//     operation within the op with the `kInPlaceResultsAttrName` attribute.
+//  b. traversing each operation in the op and rewriting it in
 //     buffer form and keeping a BlockAndValueMapping mapping of the
 //     rewrites. New allocations are introduced during this step.
 //     TODO: Allocation + depending op hoisting to outermost enclosing
 //     sequential scope.
 //  c. at the end of this bufferization, 3 cases may occur:
 //     i. inplaceable function arguments may be reused in place after the
 //        function itself has been bufferized. This is encoded by IR resembling:
 //
#endif // NDEBUG
   if (!hasWrite)
     return false;
 
   LDBG("->the corresponding buffer is not writeable\n");
   return true;
 }
 
 //===----------------------------------------------------------------------===//
-// Bufferization as simple BlockAndValueMapping rewrites.
-//===----------------------------------------------------------------------===//
-
-/// FuncOp always creates TensorToMemRef ops.
-static LogicalResult bufferizeFuncOp(FuncOp funcOp, BufferizationState &state) {
-  // Take a guard before anything else.
-  OpBuilder b(funcOp->getContext());
-  b.setInsertionPointToStart(&funcOp.body().front());
-
-  // Create BufferCastOps for function args.
-  for (auto bbArg : funcOp.getArguments()) {
-    auto tensorType = bbArg.getType().dyn_cast<TensorType>();
-    if (!tensorType)
-      continue;
-    auto rankedTensorType = tensorType.dyn_cast<RankedTensorType>();
-    // Cast the tensor to the most dynamic buffer possible. Further
-    // canonicalizations will clean up.
-    Type memRefType = rankedTensorType
-                          ? getDynamicMemRefType(rankedTensorType)
-                          : getContiguousOrUnrankedMemRefType(tensorType);
-    Value bufferCast =
-        b.create<bufferization::ToMemrefOp>(funcOp.getLoc(), memRefType, bbArg);
-    state.aliasInfo.insertNewBufferEquivalence(bufferCast, bbArg);
-    state.mapBuffer(bbArg, bufferCast);
-  }
-
-  // Bufferize function body.
-  return bufferize(&funcOp.body(), state);
-}
-
-//===----------------------------------------------------------------------===//
 // Bufferization analyses.
 //===----------------------------------------------------------------------===//
 
 /// Determine if `operand` can be bufferized in-place with `result`.
 static LogicalResult
 bufferizableInPlaceAnalysisImpl(OpOperand &operand, OpResult result,
                                 BufferizationAliasInfo &aliasInfo,
                                 const DominanceInfo &domInfo) {
for (OpOperand &opOperand : op->getOpOperands())
           if (OpResult opResult = bufferizableOp.getAliasingOpResult(opOperand))
             if (failed(bufferizableInPlaceAnalysisImpl(opOperand, opResult,
                                                        aliasInfo, domInfo)))
               return failure();
 
   return success();
 }
 
-/// Analyze the `funcOp` body to determine which OpResults are inplaceable.
-static LogicalResult
-inPlaceAnalysisFuncOpBody(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
-                          const DominanceInfo &domInfo,
-                          unsigned analysisFuzzerSeed = 0) {
-  LLVM_DEBUG(llvm::dbgs() << "\n\n");
-  LDBG("Begin InPlaceAnalysisFuncOpInternals:\n" << funcOp << '\n');
-  assert(funcOp && funcOp->getNumRegions() > 0 && !funcOp.body().empty() &&
-         "expected a funcOp definition with a body");
-
-  // Collect ops so we can build our own reverse traversal.
-  SmallVector<Operation *> ops;
-  funcOp.walk([&](Operation *op) {
-    // No tensors => no buffers.
-    if (none_of(op->getOperandTypes(), isaTensor) &&
-        none_of(op->getResultTypes(), isaTensor))
-      return;
-    ops.push_back(op);
-  });
-
-  // Set the function arguments marked with inplaceable to be known as
-  // bufferizing to a writeable memory.
-  for (BlockArgument bbArg : funcOp.getArguments()) {
-    BoolAttr inplaceAttr = funcOp.getArgAttrOfType<BoolAttr>(
-        bbArg.getArgNumber(), BufferizableOpInterface::kInplaceableAttrName);
-    if (inplaceAttr && inplaceAttr.getValue())
-      aliasInfo.setBufferizesToWritableMemory(bbArg);
-  }
-
-  LogicalResult res =
-      inPlaceAnalysis(ops, aliasInfo, domInfo, analysisFuzzerSeed);
-  LDBG("End InPlaceAnalysisFuncOpInternals:\n" << funcOp << '\n');
-
-  return res;
-}
-
 /// Assert that the current bufferization decisions are consistent.
 static LogicalResult
 checkAliasInfoConsistency(FuncOp funcOp, const DominanceInfo &domInfo,
                           const BufferizationAliasInfo &aliasInfo) {
   Operation *inconsistentOp = nullptr;
   WalkResult walkResult = funcOp.walk([&](Operation *op) {
     if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
       for (OpOperand &opOperand : op->getOpOperands())
LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
   BufferizationAliasInfo &aliasInfo = state.aliasInfo;
 
   if (funcOp.body().empty())
     return success();
 
   if (failed(checkAliasInfoConsistency(funcOp, domInfo, aliasInfo)))
     return failure();
 
+  // Collect ops so we can build our own reverse traversal.
+  SmallVector<Operation *> ops;
+  funcOp.walk([&](Operation *op) {
+    // No tensors => no buffers.
+    if (none_of(op->getOperandTypes(), isaTensor) &&
+        none_of(op->getResultTypes(), isaTensor))
+      return;
+    ops.push_back(op);
+  });
+
   // If the analysis fails, just return.
-  if (failed(inPlaceAnalysisFuncOpBody(funcOp, aliasInfo, domInfo,
-                                       options.analysisFuzzerSeed)))
+  if (failed(
+          inPlaceAnalysis(ops, aliasInfo, domInfo, options.analysisFuzzerSeed)))
     return failure();
 
   for (const std::unique_ptr<PostAnalysisStep> &step :
        options.postAnalysisSteps) {
     SmallVector<Operation *> newOps;
     if (failed(step->run(funcOp, state, newOps)))
       return failure();
     // Analyze ops that were created by the PostAnalysisStep.
     if (failed(inPlaceAnalysis(newOps, aliasInfo, domInfo)))
       return failure();
   }
 
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly) {
     annotateOpsWithBufferizationMarkers(funcOp, aliasInfo);
     return success();
   }
 
   // Bufferize all ops in funcOp.
-  if (failed(bufferizeFuncOp(funcOp, state)))
+  OpBuilder b(funcOp.getContext());
+  auto bufferizableOp =
+      dyn_cast<BufferizableOpInterface>(funcOp.getOperation());
+  assert(bufferizableOp && "must use ModuleBufferization");
+  if (failed(bufferizableOp.bufferize(b, state)))
     return failure();
 
   // Erase all obsolete ops.
   state.eraseObsoleteOps();
 
   return success();
 }

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

for (OpOperand &operand : returnOp->getOpOperands()) {
       operand.set(returnTensor);
       state.aliasInfo.insertNewBufferEquivalence(returnTensor, v);
       state.mapBuffer(returnTensor, v);
     }
     return success();
   }
 };
 
+struct FuncOpInterface
+    : public BufferizableOpInterface::ExternalModel<FuncOpInterface, FuncOp> {
+  LogicalResult bufferize(Operation *op, OpBuilder &b,
+                          BufferizationState &state) const {
+    auto funcOp = cast<FuncOp>(op);
+
+    // Take a guard before anything else.
+    OpBuilder::InsertionGuard g(b);
+    b.setInsertionPointToStart(&funcOp.body().front());
+
+    // Create BufferCastOps for function args.
+    for (auto bbArg : funcOp.getArguments()) {
+      auto tensorType = bbArg.getType().dyn_cast<TensorType>();
+      if (!tensorType)
+        continue;
+      auto rankedTensorType = tensorType.dyn_cast<RankedTensorType>();
+      // Cast the tensor to the most dynamic buffer possible. Further
+      // canonicalizations will clean up.
+      Type memRefType = rankedTensorType
+                            ? getDynamicMemRefType(rankedTensorType)
+                            : getContiguousOrUnrankedMemRefType(tensorType);
+      Value bufferCast = b.create<bufferization::ToMemrefOp>(funcOp.getLoc(),
+                                                             memRefType, bbArg);
+      state.aliasInfo.insertNewBufferEquivalence(bufferCast, bbArg);
+      state.mapBuffer(bbArg, bufferCast);
+    }
+
+    // Bufferize function body.
+    return comprehensive_bufferize::bufferize(&funcOp.body(), state);
+  }
+
+  bool isAllocationHoistingBarrier(Operation *op) const { return true; }
+};
+
 } // namespace std_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 void mlir::linalg::comprehensive_bufferize::std_ext::
     registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
   registry.addOpInterface<CallOp, std_ext::CallOpInterface>();
   registry.addOpInterface<ReturnOp, std_ext::ReturnOpInterface>();
-  registry.addOpInterface<FuncOp, AllocationHoistingBarrierOnly<FuncOp>>();
+  registry.addOpInterface<FuncOp, std_ext::FuncOpInterface>();
 }
 
 LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
     ModuleOp moduleOp, const BufferizationOptions &options) {
   SmallVector<FuncOp> orderedFuncOps;
   DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
for (FuncOp funcOp : orderedFuncOps) {
     // If the function is called, we can allocate on the caller side which lets
     // us force inplace arguments at function boundaries.
     // TODO: do not rely on this behavior.
     if (callerMap.find(funcOp) != callerMap.end())
       for (BlockArgument bbArg : funcOp.getArguments())
         if (bbArg.getType().isa<TensorType>())
           aliasInfo.setBufferizesToWritableMemory(bbArg);
 
+    // Set the function arguments marked with inplaceable to be known as
+    // bufferizing to a writeable memory.
+    for (BlockArgument bbArg : funcOp.getArguments()) {
+      BoolAttr inplaceAttr = funcOp.getArgAttrOfType<BoolAttr>(
+          bbArg.getArgNumber(), BufferizableOpInterface::kInplaceableAttrName);
+      if (inplaceAttr && inplaceAttr.getValue())
+        aliasInfo.setBufferizesToWritableMemory(bbArg);
+    }
+
     // Analyze and bufferize funcOp.
     if (failed(runComprehensiveBufferize(funcOp, options, state)))
       return failure();
 
     populateEquivalentFuncOpBBArgs(funcOp, state);
   }
 
   if (options.testAnalysisOnly)