[mlir][linalg][bufferize][NFC] Bufferize Operation* instead of FuncOp

[springerm]

Summary

This change mainly changes the API. There is no mentioning of FuncOps in ComprehensiveBufferize anymore.

Also, bufferize methods of the op interface are called for ops without tensor operands/results if they have a region.

[springerm]

rebase

[springerm]

no change

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

 /// 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` (inside `state`) consistent. Newly created operations and
   /// operations that should be re-analyzed must be stored in `newOps`.
-  virtual LogicalResult run(FuncOp funcOp, BufferizationState &state,
+  virtual LogicalResult run(Operation *op, BufferizationState &state,
                             BufferizationAliasInfo &aliasInfo,
                             SmallVector<Operation *> &newOps) = 0;
 };
 
 using PostAnalysisStepList = std::vector<std::unique_ptr<PostAnalysisStep>>;
 
 /// Options for ComprehensiveBufferize.
 struct BufferizationOptions {

 ///   that allocate and/or deallocate memref buffers.
 /// * `mapBuffer` maps a tensor value to a memref buffer during bufferization.
 /// * `lookupBuffer` returns the mapped memref buffer of a given tensor value.
 /// * `getResultBuffer` returns the memref buffer for a given tensor OpResult.
 ///   Based on inplace bufferization decisions of the analysis, it may either
 ///   directly return a mapped buffer or allocate a new brand new buffer.
 class BufferizationState {
 public:
-  BufferizationState(ModuleOp moduleOp, const BufferizationOptions &options)
-      : aliasInfo(moduleOp), options(options),
-        builder(moduleOp->getContext()) {}
+  BufferizationState(Operation *op, const BufferizationOptions &options)
+      : aliasInfo(op), options(options), builder(op->getContext()) {}
 
   // BufferizationState should be passed as a reference.
   BufferizationState(const BufferizationState &) = delete;
 
   /// Creates a memref allocation.
   Optional<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
                               ArrayRef<Value> dynShape);
 
public:
 
   /// Return a reference to the BufferizationOptions.
   const BufferizationOptions &getOptions() const { return options; }
 
   /// Return a reference to the OpBuilder.
   OpBuilder &getBuilder() { return builder; }
 
 private:
-  friend LogicalResult runComprehensiveBufferize(
-      FuncOp funcOp, const BufferizationOptions &options,
-      BufferizationState &state, const PostAnalysisStepList &extraSteps);
+  friend LogicalResult
+  runComprehensiveBufferize(Operation *op, const BufferizationOptions &options,
+                            BufferizationState &state,
+                            const PostAnalysisStepList &extraSteps);
 
   friend LogicalResult runComprehensiveBufferize(
       ModuleOp moduleOp, const BufferizationOptions &options);
 
   /// `aliasInfo` keeps track of aliasing and equivalent values. Only internal
   /// functions and `runComprehensiveBufferize` may access this object.
   BufferizationAliasInfo aliasInfo;
 

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

let methods = [
         /*desc=*/[{
           Bufferize this op, i.e., rewrite it into a memref-based equivalent.
           Tensor values should be mapped to buffer values using `state`.
 
           Implementations are required to required to bufferize nested ops
           before returning. Otherwise, nested ops will not be bufferized.
 
           This method will never be called on ops that do not have at least one
-          tensor operand or result.
+          tensor operand/result or a region.
         }],
         /*retType=*/"LogicalResult",
         /*methodName=*/"bufferize",
         /*args=*/(ins "OpBuilder &":$b,
                       "BufferizationState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           llvm_unreachable("bufferize not implemented");

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

 
 namespace linalg {
 namespace comprehensive_bufferize {
 
 struct BufferizationOptions;
 struct BufferizationState;
 struct PostAnalysisStep;
 
-/// Bufferize the given function. Does not bufferize the function boundary.
-/// Reuses an existing BufferizationState object.
-// TODO: This function is meant to be called from ModuleBufferize and not can
-// not yet be called standalone.
+/// Bufferize the given operation. Reuses an existing BufferizationState object.
 LogicalResult runComprehensiveBufferize(
-    FuncOp funcOp, const BufferizationOptions &options,
+    Operation *op, const BufferizationOptions &options,
     BufferizationState &state,
     const std::vector<std::unique_ptr<PostAnalysisStep>> &extraSteps);
 
+/// Bufferize the given operation.
+LogicalResult runComprehensiveBufferize(Operation *op,
+                                        const BufferizationOptions &options);
+
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVE_BUFFERIZE_H

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

 namespace linalg {
 namespace comprehensive_bufferize {
 
 class BufferizationAliasInfo;
 
 namespace linalg_ext {
 
 struct InitTensorEliminationStep : public PostAnalysisStep {
-  /// Try to eliminate InitTensorOps inside `funcOp`.
+  /// Try to eliminate InitTensorOps inside `op`.
   ///
   /// * `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
   ///   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 eliminateInitTensors(
-      FuncOp funcOp, BufferizationState &state,
+      Operation *op, BufferizationState &state,
       BufferizationAliasInfo &aliasInfo,
       std::function<bool(OpOperand &)> anchorMatchFunc,
       std::function<Value(OpBuilder &, Location, OpOperand &)> rewriteFunc,
       SmallVector<Operation *> &newOps);
 };
 
-/// Try to eliminate InitTensorOps inside funcOp that are anchored on an
+/// Try to eliminate InitTensorOps inside `op` that are anchored on an
 /// InsertSliceOp, i.e., if it is eventually inserted into another tensor
 /// (and some other conditions are met).
 struct InsertSliceAnchoredInitTensorEliminationStep
     : public InitTensorEliminationStep {
-  LogicalResult run(FuncOp funcOp, BufferizationState &state,
+  LogicalResult run(Operation *op, BufferizationState &state,
                     BufferizationAliasInfo &aliasInfo,
                     SmallVector<Operation *> &newOps) override;
 };
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace linalg_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_LINALG_INTERFACE_IMPL_H

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

 
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace scf_ext {
 
 /// Equivalence analysis for scf.for. Raise an error if iter_args are not
 /// equivalent to their corresponding loop yield values.
 struct AssertDestinationPassingStyle : public PostAnalysisStep {
-  LogicalResult run(FuncOp funcOp, BufferizationState &state,
+  LogicalResult run(Operation *op, BufferizationState &state,
                     BufferizationAliasInfo &aliasInfo,
                     SmallVector<Operation *> &newOps) override;
 };
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace scf_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_SCF_INTERFACE_IMPL_H

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

 
 class DialectRegistry;
 
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace tensor_ext {
 
 struct InplaceInsertSliceOpAnalysis : public PostAnalysisStep {
-  LogicalResult run(FuncOp funcOp, BufferizationState &state,
+  LogicalResult run(Operation *op, BufferizationState &state,
                     BufferizationAliasInfo &aliasInfo,
                     SmallVector<Operation *> &newOps) override;
 };
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace tensor_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_TENSOR_INTERFACE_IMPL_H

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

 mlir::linalg::comprehensive_bufferize::bufferize(Operation *op,
                                                  BufferizationState &state) {
   OpBuilder &b = state.getBuilder();
 
   // Check if op has tensor results or operands.
   auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
   bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
   bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
+  bool hasRegions = !op->getRegions().empty();
 
-  // No tensor results or operands: Simply bufferize all nested ops.
-  if (!hasTensorResult && !hasTensorOperand) {
-    for (Region &region : op->getRegions())
-      if (failed(bufferize(&region, state)))
-        return failure();
+  // No tensor results/operands or regions. We are done.
+  if (!hasTensorResult && !hasTensorOperand && !hasRegions)
     return success();
-  }
 
   // Bufferize using `BufferizableOpInterface`. Interface implementations are
   // responsible for bufferizing nested ops.
   if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op)) {
     b.setInsertionPoint(op);
     return bufferizableOp.bufferize(b, state);
   }
 
   // `op` is an unbufferizable tensor op.
   if (!state.getOptions().allowUnknownOps)
     return op->emitError() << "unsupported op with tensors";
 
   // Replace all OpOperands with "to-tensor casted" bufferized values.
   for (OpOperand &operand : op->getOpOperands()) {
     if (operand.get().getType().isa<TensorType>() &&
         state.isMapped(operand.get())) {
+      assert(state.getOptions().allowUnknownOps &&
+             "unsupported op error should have been emitted earlier");
       b.setInsertionPoint(op);
       Value toTensorOp = b.create<bufferization::ToTensorOp>(
           op->getLoc(), state.lookupBuffer(operand.get()));
       operand.set(toTensorOp);
     }
   }
 
+  // Bufferize all regions.
   for (Region &region : op->getRegions())
     if (failed(bufferize(&region, state)))
       return failure();
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//

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

op->walk<WalkOrder::PostOrder>([&](Operation *op) {
     ops.push_back(op);
   });
 
   equivalenceAnalysis(ops, aliasInfo);
 }
 
 /// Assert that the current bufferization decisions are consistent.
 static LogicalResult
-checkAliasInfoConsistency(FuncOp funcOp, const DominanceInfo &domInfo,
+checkAliasInfoConsistency(Operation *op, const DominanceInfo &domInfo,
                           const BufferizationAliasInfo &aliasInfo) {
   Operation *inconsistentOp = nullptr;
-  WalkResult walkResult = funcOp.walk([&](Operation *op) {
+  WalkResult walkResult = op->walk([&](Operation *op) {
     if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
       for (OpOperand &opOperand : op->getOpOperands())
         if (opOperand.get().getType().isa<TensorType>()) {
           OpResult opResult = bufferizableOp.getAliasingOpResult(opOperand);
           if (wouldCreateReadAfterWriteInterference(
                   opOperand, opResult, domInfo, aliasInfo,
                   /*checkConsistencyOnly=*/true)) {
             // This error can happen for two reasons. Either the input IR
annotateOpsWithBufferizationMarkers(Operation *op,
   op->walk([&](Operation *op) {
     for (OpResult opResult : op->getResults())
       if (opResult.getType().isa<TensorType>())
         setInPlaceOpResult(opResult, aliasInfo.isInPlace(opResult));
   });
 }
 
 LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
-    FuncOp funcOp, const BufferizationOptions &options,
+    Operation *op, const BufferizationOptions &options) {
+  BufferizationState state(op, options);
+  PostAnalysisStepList extraSteps;
+  return runComprehensiveBufferize(op, options, state, extraSteps);
+}
+
+LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
+    Operation *op, const BufferizationOptions &options,
     BufferizationState &state, const PostAnalysisStepList &extraSteps) {
 
-  DominanceInfo domInfo(funcOp);
+  DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.aliasInfo;
 
-  if (funcOp.body().empty())
-    return success();
-
-  if (failed(checkAliasInfoConsistency(funcOp, domInfo, aliasInfo)))
+  if (failed(checkAliasInfoConsistency(op, domInfo, aliasInfo)))
     return failure();
 
   // If the analysis fails, just return.
-  Operation *op = funcOp.getOperation();
   if (failed(inPlaceAnalysis(op, aliasInfo, state, domInfo,
                              options.analysisFuzzerSeed)))
     return failure();
   equivalenceAnalysis(op, aliasInfo);
 
   auto runPostAnalysisSteps = [&](const PostAnalysisStepList &steps) {
     for (const std::unique_ptr<PostAnalysisStep> &step : steps) {
       SmallVector<Operation *> newOps;
-      if (failed(step->run(funcOp, state, aliasInfo, newOps)))
+      if (failed(step->run(op, state, aliasInfo, newOps)))
         return failure();
       // Analyze ops that were created by the PostAnalysisStep.
       if (failed(inPlaceAnalysis(newOps, aliasInfo, state, domInfo)))
         return failure();
       equivalenceAnalysis(newOps, aliasInfo);
     }
     return success();
   };
 
   if (failed(runPostAnalysisSteps(extraSteps)))
     return failure();
   if (failed(runPostAnalysisSteps(options.postAnalysisSteps)))
     return failure();
 
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly) {
-    annotateOpsWithBufferizationMarkers(funcOp, aliasInfo);
+    annotateOpsWithBufferizationMarkers(op, aliasInfo);
     return success();
   }
 
-  // Bufferize all ops in funcOp.
-  OpBuilder b(funcOp.getContext());
-  auto bufferizableOp =
-      dyn_cast<BufferizableOpInterface>(funcOp.getOperation());
-  assert(bufferizableOp && "must use ModuleBufferization");
-  if (failed(bufferizableOp.bufferize(b, state)))
+  // Bufferize the op and its nested ops.
+  if (failed(bufferize(op, state)))
     return failure();
 
   // Erase all obsolete ops.
   state.eraseObsoleteOps();
 
   return success();
 }

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

 }
 
 /// Generic conversion for any LinalgOp on tensors.
 static LogicalResult bufferizeLinalgOp(OpBuilder &b, LinalgOp op,
                                        BufferizationState &state) {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
 
+  // Nothing to do. This op is already bufferized.
+  if (op.hasBufferSemantics())
+    return success();
+
   // Ensure op has only tensors. Allow mixed tensor-buffer mode on a per-need
   // basis.
   if (!op.hasTensorSemantics())
     return op->emitError() << "op does not have tensor semantics";
 
   Location loc = op.getLoc();
   SmallVector<Value> newInputBuffers;
   newInputBuffers.reserve(op.getNumInputs());

 
 template <>
 struct LinalgOpInterfaceHelper<> {
   static void registerOpInterface(DialectRegistry &registry) {}
 };
 
 } // namespace
 
-/// Try to eliminate InitTensorOps inside funcOp. An InitTensorOp is replaced
+/// Try to eliminate InitTensorOps inside `op`. 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::linalg_ext::
     InitTensorEliminationStep::eliminateInitTensors(
-        FuncOp funcOp, BufferizationState &state,
+        Operation *op, BufferizationState &state,
         BufferizationAliasInfo &aliasInfo,
         std::function<bool(OpOperand &)> anchorMatchFunc,
         std::function<Value(OpBuilder &, Location, OpOperand &)> rewriteFunc,
         SmallVector<Operation *> &newOps) {
-  OpBuilder b(funcOp->getContext());
+  OpBuilder b(op->getContext());
 
-  WalkResult status = funcOp->walk([&](Operation *op) {
+  WalkResult status = op->walk([&](Operation *op) {
     for (OpOperand &operand : op->getOpOperands()) {
       // Is this a matching OpOperand?
       if (!anchorMatchFunc(operand))
         continue;
 
       SetVector<Value> maybeInitTensor =
           findValueInReverseUseDefChain(operand.get(), [&](Value val) {
             // Continue traversal until this function returns true.
WalkResult status = op->walk([&](Operation *op) {
 
     // Advance to the next operation.
     return WalkResult::advance();
   });
 
   return failure(status.wasInterrupted());
 }
 
-/// Try to eliminate InitTensorOps inside funcOp. An InitTensorOp can be
+/// Try to eliminate InitTensorOps inside `op`. An InitTensorOp can be
 /// eliminated if it is eventually inserted into another tensor (and some other
 /// conditions are met).
 ///
 /// E.g.:
 /// %0 = linalg.init_tensor
 /// %1 = linalg.fill(%cst, %0) {inplace = [true]}
 /// %2 = tensor.insert_slice %1 into %t[10][20][1]
 ///

 ///   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::linalg_ext::
     InsertSliceAnchoredInitTensorEliminationStep::run(
-        FuncOp funcOp, BufferizationState &state,
+        Operation *op, BufferizationState &state,
         BufferizationAliasInfo &aliasInfo, SmallVector<Operation *> &newOps) {
   return eliminateInitTensors(
-      funcOp, state, aliasInfo,
+      op, state, aliasInfo,
       [&](OpOperand &operand) {
         auto insertSliceOp =
             dyn_cast<tensor::InsertSliceOp>(operand.getOwner());
         if (!insertSliceOp)
           return false;
         // Only inplace bufferized InsertSliceOps are eligible.
         if (!aliasInfo.isInPlace(insertSliceOp->getOpResult(0)))
           return false;

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

if (op->hasAttr(kEquivalentArgsAttr)) {
       equivBbArgs.append(op->getNumOperands(), -1);
     }
     equivBbArgs[returnVal.getOperandNumber()] = bbArg.getArgNumber();
 
     OpBuilder b(op->getContext());
     op->setAttr(kEquivalentArgsAttr, b.getI64ArrayAttr(equivBbArgs));
   }
 
-  LogicalResult run(FuncOp funcOp, BufferizationState &state,
+  LogicalResult run(Operation *op, BufferizationState &state,
                     BufferizationAliasInfo &aliasInfo,
                     SmallVector<Operation *> &newOps) override {
     ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
 
     // Support only single return-terminated block in the function.
+    auto funcOp = cast<FuncOp>(op);
     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>())
             if (aliasInfo.areEquivalentBufferizedValues(returnVal.get(),

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

for (OpOperand &operand : yieldOp->getOpOperands()) {
       // the bbArg so that it later canonicalizes away.
       operand.set(bbArg);
     }
     return success();
   }
 };
 
 LogicalResult mlir::linalg::comprehensive_bufferize::scf_ext::
-    AssertDestinationPassingStyle::run(FuncOp funcOp, BufferizationState &state,
+    AssertDestinationPassingStyle::run(Operation *op, BufferizationState &state,
                                        BufferizationAliasInfo &aliasInfo,
                                        SmallVector<Operation *> &newOps) {
   LogicalResult status = success();
-  funcOp->walk([&](scf::YieldOp yieldOp) {
+  op->walk([&](scf::YieldOp yieldOp) {
     auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
     if (!forOp)
       return WalkResult::advance();
 
     for (OpOperand &operand : yieldOp->getOpOperands()) {
       auto tensorType = operand.get().getType().dyn_cast<TensorType>();
       if (!tensorType)
         continue;

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

 };
 
 } // namespace tensor_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 LogicalResult mlir::linalg::comprehensive_bufferize::tensor_ext::
-    InplaceInsertSliceOpAnalysis::run(FuncOp funcOp, BufferizationState &state,
+    InplaceInsertSliceOpAnalysis::run(Operation *op, BufferizationState &state,
                                       BufferizationAliasInfo &aliasInfo,
                                       SmallVector<Operation *> &newOps) {
   auto &tensorState = getTensorBufferizationState(state);
-  funcOp.walk([&](InsertSliceOp insertSliceOp) {
+  op->walk([&](InsertSliceOp insertSliceOp) {
     // A copy of the source buffer is needed if either:
     //   - The producer of `source` is not inplace. This is the case where a
     //     slice is computed out of place into the inplace full tensor.
     //   - The result is not inplace. This is the case where the whole tensor is
     //     cloned and the clone needs to be updated.
     if (isSourceEquivalentToAMatchingInplaceExtractSliceOp(aliasInfo,
                                                            insertSliceOp) &&
         state.isInPlace(insertSliceOp->getResult(0)))