[mlir][bufferization] Move one-shot bufferization to Bufferization dialect

[springerm]

Summary

This commit is the first step towards unifying core bufferization and One-Shot Bufferize.

This commit does not move over the implementations of BufferizableOpInterface yet. This will be done in separate commits. This change does also not move the unit tests yet. The tests will be moved together with op interface implementations and split into separate files.

Depends On D117478

[springerm]

The CMake files are not in order yet... See BUILD.bazel in the meantime

[pifon2a]

nice!

[springerm]

Note: The moved over bufferization code is living in three files:

• OneShotAnalysis.h (+source file): Contains the analysis of the One-Shot bufferization
• Bufferize.h (+source file): Contains code related to the existing partial bufferization and the driver for BufferizableOpInterface-based bufferizations. In the future, can be run with one-shot analysis, a different analysis or no analysis at all.
• BufferizableOpInterface.h (+source +td files): Contains the op interface + helper methods needed when implementing interface methods.

[springerm]

update CMakeFiles

[springerm]

fix CMakeFiles

[springerm]

fix CMake

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

-//===- BufferizableOpInterface.h - Comprehensive Bufferize ------*- C++ -*-===//
+//===- BufferizableOpInterface.h - Bufferizable Ops -------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZABLEOPINTERFACE_H_
-#define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZABLEOPINTERFACE_H_
+#ifndef MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_
+#define MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_
 
 #include <utility>
 
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/SetVector.h"
 
 namespace mlir {
 class BlockAndValueMapping;
 class DominanceInfo;
 class FuncOp;
 
-namespace linalg {
-namespace comprehensive_bufferize {
+namespace bufferization {
 
 // TODO: from some HW description.
 static constexpr int64_t kBufferAlignments = 128;
 
-class BufferizationAliasInfo;
 class BufferizableOpInterface;
 struct BufferizationOptions;
 class BufferizationState;
 
 /// Options for ComprehensiveBufferize.
 struct BufferizationOptions {
   using AllocationFn = std::function<FailureOr<Value>(
       OpBuilder &, Location, MemRefType, ArrayRef<Value>)>;
public:
   Optional<const StateT *> getDialectState(StringRef name) const {
     auto it = dialectState.find(name);
     if (it == dialectState.end())
       return None;
     return static_cast<const StateT *>(it->getSecond().get());
   }
 
   /// Return dialect-specific bufferization state or create one if none exists.
-  template <typename StateT> StateT &getOrCreateDialectState(StringRef name) {
+  template <typename StateT>
+  StateT &getOrCreateDialectState(StringRef name) {
     // Create state if it does not exist yet.
     if (!dialectState.count(name))
       dialectState[name] = std::make_unique<StateT>();
     return static_cast<StateT &>(*dialectState[name]);
   }
 
   /// Return a reference to the BufferizationOptions.
   const BufferizationOptions &getOptions() const { return options; }

 /// allocated using `createAlloc`.
 LogicalResult createDealloc(OpBuilder &b, Location loc, Value allocatedBuffer,
                             const BufferizationOptions &options);
 
 /// Creates a memcpy between two given buffers.
 LogicalResult createMemCpy(OpBuilder &b, Location loc, Value from, Value to,
                            const BufferizationOptions &options);
 
-} // namespace comprehensive_bufferize
-} // namespace linalg
+} // namespace bufferization
 } // namespace mlir
 
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h.inc"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h.inc"
 
 namespace mlir {
-namespace linalg {
-namespace comprehensive_bufferize {
+namespace bufferization {
 
 /// AllocationHoistingBarrierOnly is an external implementation of
 /// BufferizableOpInterface for ops that are (not yet) bufferizable, but are
 /// known to be allocation hoisting barriers. All interface methods (except for
 /// `isAllocationHoistingBarrier`) are implemented conservatively.
 template <typename OpTy>
 struct AllocationHoistingBarrierOnly
     : public BufferizableOpInterface::ExternalModel<
struct AllocationHoistingBarrierOnly
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationState &state) const {
     return failure();
   }
 
   bool isAllocationHoistingBarrier(Operation *op) const { return true; }
 };
 
-} // namespace comprehensive_bufferize
-} // namespace linalg
+} // namespace bufferization
 } // namespace mlir
 
-#endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZABLEOPINTERFACE_H_
+#endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_

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

-//===-- BufferizableOpInterface.td - Compreh. Bufferize ----*- tablegen -*-===//
+//===-- BufferizableOpInterface.td - Bufferizable Ops ------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef BUFFERIZABLE_OP_INTERFACE
 #define BUFFERIZABLE_OP_INTERFACE
 
 include "mlir/IR/OpBase.td"
 
 def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
   let description = [{
     An op interface for Comprehensive Bufferization. Ops that implement this
     interface can be bufferized using Comprehensive Bufferization.
   }];
-  let cppNamespace = "::mlir::linalg::comprehensive_bufferize";
+  let cppNamespace = "::mlir::bufferization";
   let methods = [
       InterfaceMethod<
         /*desc=*/[{
           Return `true` if the given OpOperand bufferizes to a memory read. This
           method will never be called on OpOperands that do not have a tensor
           type.
 
           Note: It is always safe to consider an OpOperand as a memory read,
bool bufferizesToAliasOnly(OpOperand &opOperand,
           && !bufferizableOp.bufferizesToMemoryWrite(opOperand, state)
           && static_cast<bool>(
               bufferizableOp.getAliasingOpResult(opOperand, state));
     }
 
     // TODO: The following two attributes should belong to the tensor dialect.
     // The corresponding verifier should also be in the tensor dialect.
     /// Attribute name used to mark region arguments that can be bufferized
-    /// in-place during linalg comprehensive bufferization.
+    /// in-place during one-shot bufferization.
     constexpr const static ::llvm::StringLiteral
       kInplaceableAttrName = "linalg.inplaceable";
 
     /// Attribute name used to mark the bufferization layout for region
-    /// arguments during linalg comprehensive bufferization.
+    /// arguments during one-shot bufferization.
     constexpr const static ::llvm::StringLiteral
       kBufferLayoutAttrName = "linalg.buffer_layout";
   }];
 }
 
 #endif  // BUFFERIZABLE_OP_INTERFACE

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

 //===- BufferizationInterfaceImpl.h - Bufferization Impl. of Op Interface -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZATIONINTERFACEIMPL_H
-#define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZATIONINTERFACEIMPL_H
+#ifndef MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZATIONINTERFACEIMPL_H_
+#define MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZATIONINTERFACEIMPL_H_
 
 namespace mlir {
 
 class DialectRegistry;
 
-namespace linalg {
-namespace comprehensive_bufferize {
+namespace bufferization {
 namespace bufferization_ext {
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace bufferization_ext
-} // namespace comprehensive_bufferize
-} // namespace linalg
+} // namespace bufferization
 } // namespace mlir
 
-#endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZATIONINTERFACEIMPL_H
+#endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZATIONINTERFACEIMPL_H_

mlir/include/mlir/Dialect/Bufferization/IR/CMakeLists.txt

 add_mlir_dialect(BufferizationOps bufferization)
 add_mlir_doc(BufferizationOps BufferizationOps Dialects/ -gen-dialect-doc)
 add_mlir_interface(AllocationOpInterface)
+add_mlir_interface(BufferizableOpInterface)

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

 void populateBufferizeMaterializationLegality(ConversionTarget &target);
 
 /// Populate patterns to eliminate bufferize materializations.
 ///
 /// In particular, these are the tensor_load/buffer_cast ops.
 void populateEliminateBufferizeMaterializationsPatterns(
     BufferizeTypeConverter &typeConverter, RewritePatternSet &patterns);
 
+class BufferizationState;
+
+/// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.
+/// Whether buffer copies are needed or not is queried from `state`.
+///
+/// Note: If `allowUnknownOps` is set to false, bufferization fails when an
+/// unknown op (that does not implement `BufferizableOpInterface`) is found. No
+/// to_tensor/to_memref ops are inserted in that case.
+///
+/// Note: Tje layout map chosen to bufferize is the most dynamic canonical
+/// strided layout of the proper rank. This ensures compatibility with expected
+/// layouts after transformations. Combinations of memref.cast +
+/// canonicalization are responsible for clean ups.
+// TODO: Extract `options` from `state` and pass as separate argument.
+LogicalResult bufferizeOp(Operation *op, const BufferizationState &state);
+
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZE_H

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

-//===- ComprehensiveBufferize.h - Linalg bufferization pass -----*- C++ -*-===//
+//===- OneShotAnalysis.h - One-Shot (Single Pass) Analysis ------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVEBUFFERIZE_H
-#define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVEBUFFERIZE_H
+#ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTANALYSIS_H
+#define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTANALYSIS_H
 
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 
 namespace mlir {
-
-namespace linalg {
-namespace comprehensive_bufferize {
+namespace bufferization {
 
 class AnalysisBufferizationState;
 class BufferizationAliasInfo;
 struct AnalysisBufferizationOptions;
-class BufferizationState;
 
 /// PostAnalysisSteps can be registered with `BufferizationOptions` and are
 /// executed after the analysis, but before bufferization. They can be used to
 /// implement custom dialect-specific optimizations.
 struct PostAnalysisStep {
   virtual ~PostAnalysisStep() = default;
 
   /// Run the post analysis step. This function may modify the IR, but must keep
public:
   /// 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;
 
 private:
   /// `aliasInfo` keeps track of aliasing and equivalent values. Only internal
-  /// functions and `runComprehensiveBufferize` may access this object.
+  /// functions and `runOneShotBufferize` may access this object.
   BufferizationAliasInfo aliasInfo;
 };
 
 /// Analyze `op` and its nested ops. Bufferization decisions are stored in
 /// `state`.
 LogicalResult analyzeOp(Operation *op, AnalysisBufferizationState &state);
 
-/// Bufferize `op` and its nested ops. Bufferization decisions are stored in
-/// `state`.
-LogicalResult bufferizeOp(Operation *op, const BufferizationState &state);
-
-/// Run Comprehensive Bufferize on the given op: Analysis + Bufferization
-LogicalResult runComprehensiveBufferize(
-    Operation *op, std::unique_ptr<AnalysisBufferizationOptions> options);
+/// Run One-Shot Bufferize on the given op: Analysis + Bufferization
+LogicalResult
+runOneShotBufferize(Operation *op,
+                    std::unique_ptr<AnalysisBufferizationOptions> options);
 
-} // namespace comprehensive_bufferize
-} // namespace linalg
+} // namespace bufferization
 } // namespace mlir
 
-#endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVEBUFFERIZE_H
+#endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTANALYSIS_H

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/CMakeLists.txt

-set(LLVM_TARGET_DEFINITIONS BufferizableOpInterface.td)
-mlir_tablegen(BufferizableOpInterface.h.inc -gen-op-interface-decls)
-mlir_tablegen(BufferizableOpInterface.cpp.inc -gen-op-interface-defs)
-add_public_tablegen_target(MLIRBufferizableOpInterfaceIncGen)
-add_dependencies(mlir-headers MLIRBufferizableOpInterfaceIncGen)
+# no targets defined here
+

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

 //===- LinalgInterfaceImpl.h - Linalg Impl. of BufferizableOpInterface ----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_LINALGINTERFACEIMPL_H
 #define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_LINALGINTERFACEIMPL_H
 
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h"
+#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 
 namespace mlir {
-
 class DialectRegistry;
 
 namespace linalg {
 namespace comprehensive_bufferize {
-
-class BufferizationAliasInfo;
-
 namespace linalg_ext {
 
-struct InitTensorEliminationStep : public PostAnalysisStep {
+struct InitTensorEliminationStep : public bufferization::PostAnalysisStep {
   /// A function that matches anchor OpOperands for InitTensorOp elimination.
   using AnchorMatchFn = std::function<bool(OpOperand &)>;
 
   /// A function that rewrites matched anchors.
   using RewriteFn = std::function<Value(OpBuilder &, Location, OpOperand &)>;
 
   /// 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(Operation *op, BufferizationState &state,
-                                     BufferizationAliasInfo &aliasInfo,
-                                     AnchorMatchFn anchorMatchFunc,
-                                     RewriteFn rewriteFunc,
-                                     SmallVector<Operation *> &newOps);
+  LogicalResult
+  eliminateInitTensors(Operation *op, bufferization::BufferizationState &state,
+                       bufferization::BufferizationAliasInfo &aliasInfo,
+                       AnchorMatchFn anchorMatchFunc, RewriteFn rewriteFunc,
 };
 
 /// 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(Operation *op, BufferizationState &state,
-                    BufferizationAliasInfo &aliasInfo,
+  LogicalResult run(Operation *op, bufferization::BufferizationState &state,
+                    bufferization::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_LINALGINTERFACEIMPL_H

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

 #include <memory>
 
 namespace mlir {
 
 class DialectRegistry;
 struct LogicalResult;
 class ModuleOp;
 
+namespace bufferization {
+struct AnalysisBufferizationOptions;
+} // namespace bufferization
+
 namespace linalg {
 namespace comprehensive_bufferize {
 
-struct AnalysisBufferizationOptions;
-
 /// Run Module Bufferization on the given module. Performs a simple function
 /// call analysis to determine which function arguments are inplaceable. Then
-/// analyzes and bufferizes FuncOps one-by-one with Comprehensive Bufferization.
+/// analyzes and bufferizes FuncOps one-by-one with One-Shot Bufferize.
 LogicalResult runComprehensiveBufferize(
-    ModuleOp moduleOp, std::unique_ptr<AnalysisBufferizationOptions> options);
+    ModuleOp moduleOp,
+    std::unique_ptr<bufferization::AnalysisBufferizationOptions> options);
 
 namespace std_ext {
 
 void registerModuleBufferizationExternalModels(DialectRegistry &registry);
 
 } // namespace std_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_MODULEBUFFERIZATION_H

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

 //===- SCFInterfaceImpl.h - SCF Impl. of BufferizableOpInterface ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_SCFINTERFACEIMPL_H
 #define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_SCFINTERFACEIMPL_H
 
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h"
+#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 
 namespace mlir {
 
 class DialectRegistry;
 
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace scf_ext {
 
 /// Assert that yielded values of an scf.for op are aliasing their corresponding
 /// bbArgs. This is required because the i-th OpResult of an scf.for op is
 /// currently assumed to alias with the i-th iter_arg (in the absence of
 /// conflicts).
-struct AssertScfForAliasingProperties : public PostAnalysisStep {
-  LogicalResult run(Operation *op, BufferizationState &state,
-                    BufferizationAliasInfo &aliasInfo,
+struct AssertScfForAliasingProperties : public bufferization::PostAnalysisStep {
+  LogicalResult run(Operation *op, bufferization::BufferizationState &state,
+                    bufferization::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_SCFINTERFACEIMPL_H

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

-//===- BufferizableOpInterface.cpp - Comprehensive Bufferize --------------===//
+//===- BufferizableOpInterface.cpp - Bufferizable Ops  ---=----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
-
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Value.h"
 #include "llvm/Support/Debug.h"
 
 namespace mlir {
-namespace linalg {
-namespace comprehensive_bufferize {
+namespace bufferization {
 
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp.inc"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.cpp.inc"
 
-} // namespace comprehensive_bufferize
-} // namespace linalg
+} // namespace bufferization
 } // namespace mlir
 
 #define DEBUG_TYPE "bufferizable-op-interface"
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << (X))
 
 using namespace mlir;
-using namespace linalg::comprehensive_bufferize;
+using namespace bufferization;
 
 //===----------------------------------------------------------------------===//
 // BufferizationOptions
 //===----------------------------------------------------------------------===//
 
 // Default constructor for BufferizationOptions.
 BufferizationOptions::BufferizationOptions() {}
 
-BufferizableOpInterface mlir::linalg::comprehensive_bufferize::
-    BufferizationOptions::dynCastBufferizableOp(Operation *op) const {
+BufferizableOpInterface
   if (isOpAllowed(op))
     return dyn_cast<BufferizableOpInterface>(op);
   return nullptr;
 }
 
-BufferizableOpInterface mlir::linalg::comprehensive_bufferize::
-    BufferizationOptions::dynCastBufferizableOp(Value value) const {
+BufferizableOpInterface
   if (auto bufferizableOp = value.getDefiningOp<BufferizableOpInterface>())
     if (isOpAllowed(bufferizableOp.getOperation()))
       return bufferizableOp;
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 // Helper functions for BufferizableOpInterface
 //===----------------------------------------------------------------------===//
 
 static void setInsertionPointAfter(OpBuilder &b, Value value) {
   if (auto bbArg = value.dyn_cast<BlockArgument>()) {
     b.setInsertionPointToStart(bbArg.getOwner());
   } else {
     b.setInsertionPointAfter(value.getDefiningOp());
   }
 }
 
 /// Determine which OpOperand* will alias with `result` if the op is bufferized
 /// in place. Return an empty vector if the op is not bufferizable.
 SmallVector<OpOperand *>
-mlir::linalg::comprehensive_bufferize::BufferizationState::getAliasingOpOperand(
-    OpResult result) const {
+BufferizationState::getAliasingOpOperand(OpResult result) const {
   if (Operation *op = result.getDefiningOp())
     if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
       return bufferizableOp.getAliasingOpOperand(result, *this);
   return {};
 }
 
 /// Determine which OpResult will alias with `opOperand` if the op is bufferized
 /// in place. Return an empty OpResult if the op is not bufferizable.
-OpResult
-mlir::linalg::comprehensive_bufferize::BufferizationState::getAliasingOpResult(
-    OpOperand &opOperand) const {
+OpResult BufferizationState::getAliasingOpResult(OpOperand &opOperand) const {
   if (auto bufferizableOp =
           dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
     return bufferizableOp.getAliasingOpResult(opOperand, *this);
   return OpResult();
 }
 
 /// Return true if `opOperand` bufferizes to a memory read. Return `true` if the
 /// op is not bufferizable.
-bool mlir::linalg::comprehensive_bufferize::BufferizationState::
-    bufferizesToMemoryRead(OpOperand &opOperand) const {
+bool BufferizationState::bufferizesToMemoryRead(OpOperand &opOperand) const {
   if (auto bufferizableOp =
           dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
     return bufferizableOp.bufferizesToMemoryRead(opOperand, *this);
 
   // Unknown op that returns a tensor. The inplace analysis does not support it.
   // Conservatively return true.
   return true;
 }
 
 /// Return true if `opOperand` bufferizes to a memory write. Return
 /// `true` if the op is not bufferizable.
-bool mlir::linalg::comprehensive_bufferize::BufferizationState::
-    bufferizesToMemoryWrite(OpOperand &opOperand) const {
+bool BufferizationState::bufferizesToMemoryWrite(OpOperand &opOperand) const {
   if (auto bufferizableOp =
           dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
     return bufferizableOp.bufferizesToMemoryWrite(opOperand, *this);
 
   // Unknown op that returns a tensor. The inplace analysis does not support it.
   // Conservatively return true.
   return true;
 }
 
 /// Return true if `opOperand` does neither read nor write but bufferizes to an
 /// alias. Return false if the op is not bufferizable.
-bool mlir::linalg::comprehensive_bufferize::BufferizationState::
-    bufferizesToAliasOnly(OpOperand &opOperand) const {
+bool BufferizationState::bufferizesToAliasOnly(OpOperand &opOperand) const {
   if (auto bufferizableOp =
           dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
     return bufferizableOp.bufferizesToAliasOnly(opOperand, *this);
 
   // Unknown op that returns a tensor. The inplace analysis does not support it.
   // Conservatively return false.
   return false;
 }
 
 /// Return true if the given value is read by an op that bufferizes to a memory
 /// read. Also takes into account ops that create an alias but do not read by
 /// themselves (e.g., ExtractSliceOp).
-bool mlir::linalg::comprehensive_bufferize::BufferizationState::isValueRead(
-    Value value) const {
+bool BufferizationState::isValueRead(Value value) const {
   assert(value.getType().isa<TensorType>() && "expected TensorType");
   SmallVector<OpOperand *> workingSet;
   for (OpOperand &use : value.getUses())
     workingSet.push_back(&use);
 
   while (!workingSet.empty()) {
     OpOperand *uMaybeReading = workingSet.pop_back_val();
     // Skip over all ops that neither read nor write (but create an alias).
     if (bufferizesToAliasOnly(*uMaybeReading))
       for (OpOperand &use : getAliasingOpResult(*uMaybeReading).getUses())
         workingSet.push_back(&use);
     if (bufferizesToMemoryRead(*uMaybeReading))
       return true;
   }
 
   return false;
 }
 
 // Starting from `value`, follow the use-def chain in reverse, always selecting
 // the aliasing OpOperands. Find and return Values for which `condition`
 // evaluates to true. OpOperands of such matching Values are not traversed any
 // further.
-llvm::SetVector<Value> mlir::linalg::comprehensive_bufferize::
-    BufferizationState::findValueInReverseUseDefChain(
-        Value value, llvm::function_ref<bool(Value)> condition) const {
+llvm::SetVector<Value> BufferizationState::findValueInReverseUseDefChain(
   llvm::SetVector<Value> result, workingSet;
   workingSet.insert(value);
 
   while (!workingSet.empty()) {
     Value value = workingSet.pop_back_val();
     if (condition(value) || value.isa<BlockArgument>()) {
       result.insert(value);
       continue;
while (!workingSet.empty()) {
     for (OpOperand *o : opOperands)
       workingSet.insert(o->get());
   }
 
   return result;
 }
 
 // Find the Values of the last preceding write of a given Value.
-llvm::SetVector<Value> mlir::linalg::comprehensive_bufferize::
-    BufferizationState::findLastPrecedingWrite(Value value) const {
+llvm::SetVector<Value>
   return findValueInReverseUseDefChain(value, [&](Value value) {
     Operation *op = value.getDefiningOp();
     if (!op)
       return true;
     auto bufferizableOp = options.dynCastBufferizableOp(op);
     if (!bufferizableOp)
       return true;
     return bufferizableOp.isMemoryWrite(value.cast<OpResult>(), *this);
   });
 }
 
-mlir::linalg::comprehensive_bufferize::BufferizationState::BufferizationState(
-    const BufferizationOptions &options)
+BufferizationState::BufferizationState(const BufferizationOptions &options)
     : options(options) {}
 
 // bufferization.to_memref is not allowed to change the rank.
 static void ensureToMemrefOpIsValid(Value tensor, Type memrefType) {
 #ifndef NDEBUG
   auto rankedTensorType = tensor.getType().dyn_cast<RankedTensorType>();
   assert((!rankedTensorType || memrefType.cast<MemRefType>().getRank() ==
                                    rankedTensorType.getRank()) &&
static Value lookupBuffer(RewriterBase &rewriter, Value tensor) {
   ensureToMemrefOpIsValid(tensor, memrefType);
   return rewriter.create<bufferization::ToMemrefOp>(tensor.getLoc(), memrefType,
                                                     tensor);
 }
 
 /// 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.
-FailureOr<Value>
-mlir::linalg::comprehensive_bufferize::BufferizationState::getBuffer(
+FailureOr<Value> BufferizationState::getBuffer(
     RewriterBase &rewriter, OpOperand &opOperand, bool forceInPlace,
     Optional<Operation *> customCopyInsertionPoint) const {
   OpBuilder::InsertionGuard guard(rewriter);
   Operation *op = opOperand.getOwner();
   Location loc = op->getLoc();
   Value operand = opOperand.get();
   Value operandBuffer = lookupBuffer(rewriter, operand);
 
FailureOr<Value> BufferizationState::getBuffer(
   }
   if (failed(
           createMemCpy(rewriter, loc, operandBuffer, *resultBuffer, options)))
     return failure();
 
   return resultBuffer;
 }
 
-void mlir::linalg::comprehensive_bufferize::replaceOpWithBufferizedValues(
-    RewriterBase &rewriter, Operation *op, ValueRange values) {
+void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
+                                                  Operation *op,
+                                                  ValueRange values) {
   OpBuilder::InsertionGuard g(rewriter);
 
   // Replace all OpResults with the given values.
   for (OpResult opResult : op->getOpResults()) {
     // Skip OpResult if it has no uses.
     if (opResult.getUses().empty())
       continue;
 
if (dynShape.empty())
     moveInsertionPointToAllocationHoistingBarrier(b);
 
   return allocMemRefType;
 }
 
 /// Create an AllocOp/DeallocOp pair, where the AllocOp is after
 /// `shapedValue.getDefiningOp` (or at the top of the block in case of a
 /// bbArg) and the DeallocOp is at the end of the block.
-FailureOr<Value> mlir::linalg::comprehensive_bufferize::createAlloc(
-    OpBuilder &b, Location loc, Value shapedValue, bool deallocMemref,
+FailureOr<Value>
+bufferization::createAlloc(OpBuilder &b, Location loc, Value shapedValue,
+                           bool deallocMemref,
                            const BufferizationOptions &options) {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
 
   // 1. Create memory allocation.
   assert(shapedValue.getType().isa<ShapedType>());
   MemRefType memRefType = shapedValue.getType().dyn_cast<MemRefType>();
   SmallVector<Value> dynShape;
   // Note: getAllocationTypeAndShape also sets the insertion point.
if (deallocMemref) {
     if (failed(createDealloc(b, loc, allocated.getValue(), options)))
       return failure();
   }
 
   return casted;
 }
 
 /// Create a memref allocation.
-FailureOr<Value> mlir::linalg::comprehensive_bufferize::createAlloc(
-    OpBuilder &b, Location loc, MemRefType type, ArrayRef<Value> dynShape,
+FailureOr<Value>
+bufferization::createAlloc(OpBuilder &b, Location loc, MemRefType type,
+                           ArrayRef<Value> dynShape,
                            const BufferizationOptions &options) {
   if (options.allocationFn)
     return (*options.allocationFn)(b, loc, type, dynShape);
 
   // Default bufferallocation via AllocOp.
   Value allocated = b.create<memref::AllocOp>(
       loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
   return allocated;
 }
 
 /// Create a memref deallocation.
-LogicalResult mlir::linalg::comprehensive_bufferize::createDealloc(
-    OpBuilder &b, Location loc, Value allocatedBuffer,
+LogicalResult
+bufferization::createDealloc(OpBuilder &b, Location loc, Value allocatedBuffer,
                              const BufferizationOptions &options) {
   if (options.deallocationFn)
     return (*options.deallocationFn)(b, loc, allocatedBuffer);
 
   // Default buffer deallocation via DeallocOp.
   b.create<memref::DeallocOp>(loc, allocatedBuffer);
   return success();
 }
 
 /// Create a memory copy between two memref buffers.
-LogicalResult mlir::linalg::comprehensive_bufferize::createMemCpy(
-    OpBuilder &b, Location loc, Value from, Value to,
+LogicalResult bufferization::createMemCpy(OpBuilder &b, Location loc,
+                                          Value from, Value to,
                                           const BufferizationOptions &options) {
   if (options.memCpyFn)
     return (*options.memCpyFn)(b, loc, from, to);
 
   b.create<memref::CopyOp>(loc, from, to);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific BlockAndValueMapping support with debugging.
 //===----------------------------------------------------------------------===//
 
-bool mlir::linalg::comprehensive_bufferize::isFunctionArgument(Value value) {
+bool bufferization::isFunctionArgument(Value value) {
   auto bbArg = value.dyn_cast<BlockArgument>();
   if (!bbArg)
     return false;
   return isa<FuncOp>(bbArg.getOwner()->getParentOp());
 }
 
-MemRefType mlir::linalg::comprehensive_bufferize::getContiguousMemRefType(
-    ShapedType shapedType, MemRefLayoutAttrInterface layout,
+MemRefType
+bufferization::getContiguousMemRefType(ShapedType shapedType,
+                                       MemRefLayoutAttrInterface layout,
                                        Attribute memorySpace) {
   return MemRefType::get(shapedType.getShape(), shapedType.getElementType(),
                          layout, memorySpace);
 }
 
-UnrankedMemRefType mlir::linalg::comprehensive_bufferize::getUnrankedMemRefType(
-    Type elementType, Attribute memorySpace) {
+UnrankedMemRefType bufferization::getUnrankedMemRefType(Type elementType,
+                                                        Attribute memorySpace) {
   return UnrankedMemRefType::get(elementType, memorySpace);
 }
 
-MemRefType mlir::linalg::comprehensive_bufferize::getDynamicMemRefType(
-    RankedTensorType tensorType, unsigned addressSpace) {
+MemRefType bufferization::getDynamicMemRefType(RankedTensorType tensorType,
+                                               unsigned addressSpace) {
   // TODO: address space decisions to connect with the actual alloc.
   int64_t dynamicOffset = ShapedType::kDynamicStrideOrOffset;
   SmallVector<int64_t> dynamicStrides(tensorType.getRank(),
                                       ShapedType::kDynamicStrideOrOffset);
   AffineMap stridedLayout = makeStridedLinearLayoutMap(
       dynamicStrides, dynamicOffset, tensorType.getContext());
   return MemRefType::get(tensorType.getShape(), tensorType.getElementType(),
                          stridedLayout, addressSpace);
 }

mlir/lib/Dialect/Bufferization/IR/BufferizationInterfaceImpl.cpp

 //===- BufferizationInterfaceImpl.cpp - Bufferization Impl. of Interface --===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizationInterfaceImpl.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizationInterfaceImpl.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
-using namespace linalg;
-using namespace comprehensive_bufferize;
+using namespace mlir::bufferization;
 
 namespace mlir {
-namespace linalg {
-namespace comprehensive_bufferize {
+namespace bufferization {
 namespace bufferization_ext {
 
-// TODO: These ops should implement BufferizableOpInterface directly when moved
-// to the Bufferization dialect.
+// TODO: These ops should implement BufferizableOpInterface.
 
 /// Bufferization of bufferization.to_memref. to_memref(to_tensor(x)) is folded
 /// to x. Other to_memref ops are ignored during bufferization.
 ///
 /// ToMemrefOp casts a tensor into a memref. The resulting memref is the memory
 /// location of the incoming tensor once it will be bufferized. In the anlysis,
 /// the incoming tensor is assumed to bufferize to a memory read and to an
 /// inplace memory write, since it is unknown what will happen to the resulting
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const BufferizationState &state) const {
     // It is unknown whether the resulting MemRef will be written or not.
     return true;
   }
 
   bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand,
                             const BufferizationState &state) const {
     // ToMemrefOps always bufferize inplace.
-    // TODO: Remove ToMemrefOps from the analysis.
     return true;
   }
 
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                const BufferizationState &state) const {
     return OpResult();
   }
 
struct ToTensorOpInterface
   bool isWritable(Operation *op, Value value,
                   const BufferizationState &state) const {
     // It is unknown whether the memref operand is writable or not.
     return false;
   }
 };
 
 } // namespace bufferization_ext
-} // namespace comprehensive_bufferize
-} // namespace linalg
+} // namespace bufferization
 } // namespace mlir
 
-void mlir::linalg::comprehensive_bufferize::bufferization_ext::
-    registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
-  registry.addOpInterface<bufferization::ToMemrefOp,
-                          bufferization_ext::ToMemrefOpInterface>();
-  registry.addOpInterface<bufferization::ToTensorOp,
-                          bufferization_ext::ToTensorOpInterface>();
+void bufferization_ext::registerBufferizableOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addOpInterface<ToMemrefOp, bufferization_ext::ToMemrefOpInterface>();
+  registry.addOpInterface<ToTensorOp, bufferization_ext::ToTensorOpInterface>();
 }

mlir/lib/Dialect/Bufferization/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRBufferization
+  PARTIAL_SOURCES_INTENDED
   AllocationOpInterface.cpp
   BufferizationOps.cpp
   BufferizationDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization
 
   DEPENDS
   MLIRAllocationOpInterfaceIncGen
   MLIRBufferizationOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRDialect
   MLIRIR
   MLIRTensor
   MLIRMemRef
   )
+
+add_mlir_dialect_library(MLIRBufferizableOpInterface
+  PARTIAL_SOURCES_INTENDED
+  BufferizableOpInterface.cpp
+  BufferizationInterfaceImpl.cpp
+
+  DEPENDS
+  MLIRBufferizableOpInterfaceIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRBufferization
+  MLIRMemRef
+)

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

 //===- Bufferize.cpp - Bufferization utilities ----------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
 
 //===----------------------------------------------------------------------===//
 // BufferizeTypeConverter
 //===----------------------------------------------------------------------===//
 
struct FinalizingBufferizePass
   }
 };
 } // namespace
 
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::bufferization::createFinalizingBufferizePass() {
   return std::make_unique<FinalizingBufferizePass>();
 }
+
+static bool isaTensor(Type t) { return t.isa<TensorType>(); }
+
+/// Return true if the given op has a tensor result or a tensor operand.
+static bool hasTensorSemantics(Operation *op) {
+  bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
+  bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
+  return hasTensorResult || hasTensorOperand;
+}
+
+/// Rewrite pattern that bufferizes bufferizable ops.
+struct BufferizationPattern
+    : public OpInterfaceRewritePattern<BufferizableOpInterface> {
+  BufferizationPattern(MLIRContext *context, const BufferizationState &state,
+                       PatternBenefit benefit = 1)
+      : OpInterfaceRewritePattern<BufferizableOpInterface>(context, benefit),
+        state(state) {}
+
+  LogicalResult matchAndRewrite(BufferizableOpInterface bufferizableOp,
+                                PatternRewriter &rewriter) const override {
+    // No tensors => no buffers.
+    if (!hasTensorSemantics(bufferizableOp.getOperation()))
+      return failure();
+    if (!state.getOptions().isOpAllowed(bufferizableOp.getOperation()))
+      return failure();
+    return bufferizableOp.bufferize(rewriter, state);
+  }
+
+private:
+  const BufferizationState &state;
+};
+
+/// Check the result of bufferization. Return an error if an op was not
+/// bufferized, unless partial bufferization is allowed.
+static LogicalResult
+checkBufferizationResult(Operation *op, const BufferizationOptions &options) {
+  if (!options.allowUnknownOps) {
+    // Check if all ops were bufferized.
+    LogicalResult status = success();
+    op->walk([&](Operation *op) {
+      if (!hasTensorSemantics(op))
+        return WalkResult::advance();
+
+      // Bufferization dialect ops will canonicalize away if all other ops are
+      // bufferized.
+      if (isa<bufferization::ToMemrefOp, bufferization::ToTensorOp>(op))
+        return WalkResult::advance();
+
+      // Ops that are not in the allow list can be ignored.
+      if (!options.isOpAllowed(op))
+        return WalkResult::advance();
+
+      // Ops without any uses and no side effects will fold away.
+      if (op->getUses().empty() && MemoryEffectOpInterface::hasNoEffect(op))
+        return WalkResult::advance();
+
+      status = op->emitError("op was not bufferized");
+      return WalkResult::interrupt();
+    });
+
+    if (failed(status))
+      return status;
+  }
+
+  return success();
+}
+
+LogicalResult bufferization::bufferizeOp(Operation *op,
+                                         const BufferizationState &state) {
+  // Bufferize the op and its nested ops.
+  OwningRewritePatternList patterns(op->getContext());
+  patterns.add<BufferizationPattern>(op->getContext(), state);
+  if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns))))
+    return failure();
+
+  return checkBufferizationResult(op, state.getOptions());
+}

mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRBufferizationTransforms
   Bufferize.cpp
   BufferDeallocation.cpp
+  OneShotAnalysis.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization
 
   DEPENDS
   MLIRBufferizationPassIncGen
 
   LINK_LIBS PUBLIC
+  MLIRBufferizableOpInterface
   MLIRBufferization
+  MLIRControlFlowInterfaces
+  MLIRInferTypeOpInterface
+  MLIRIR
+  MLIRMemRef
   MLIRPass
+  MLIRStandard
   MLIRTransforms
 )

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

-//===- ComprehensiveBufferize.cpp - Single pass bufferization -------------===//
+//===- OneShotAnalysis.cpp - One-Shot (Single Pass) Analysis --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
-// Comprehensive Bufferize bufferizes function bodies. Function boundaries
-// (FuncOp bbArgs, CallOps, ReturnOps) are treated as "unknown" ops.
-// ModuleBufferization.cpp is an extension of Comprehensive Bufferize for simple
+// One-Shot Analysis analyzes function bodies. Function boundaries (FuncOp
+// bbArgs, CallOps, ReturnOps) are treated as "unknown" ops.
+// ModuleBufferization.cpp is an extension of One-Shot Analysis for simple
 // call graphs.
 //
-// Comprehensive Bufferize consists of two phases.
+// One-Shot Bufferize consists of two phases.
 //
 // 1. Analyze ops to decide which OpResults can bufferize inplace, i.e., without
 //    inserting buffer copies. The analysis queries op bufferization semantics
 //    via `BufferizableOpInterface`.
 // 2. Bufferize ops by calling `BufferizableOpInterface::bufferize`. This
 //    function does not generate buffer copies for OpResults that were decided
 //    to bufferize inplace during the analysis phase.
 //
+// This file contains only the analysis. The actual bufferization is implemented
+// via `bufferizeOp` (Bufferize.h). For convenience, this file also contains a
+// helper function `runOneShotBufferize` that analyzes an op (and its nested
+// ops) and then bufferizes it.
+//
 // Inplace bufferization decisions are passed from the analysis to the
 // bufferization phase via `BufferizationState` and `BufferizationAliasInfo`.
 // They can be printed for debugging purposes with `testAnalysisOnly`.
 //
 // Ops that do not implement `BufferizableOpInterface` can be analyzed but are
-// treated conservatively. E.g., the analysis has to assume that their
+// treated conservatively. E.g., the analysis has to assume that their tensor
 // OpOperands bufferize to memory writes. While such ops can be analyzed, they
 // are not bufferized and remain in the IR. to_tensor and to_memref ops are
 // inserted at the bufferization boundary.
 //
-// Note: If `allowUnknownOps` is set to false, bufferization fails when an
-// unknown op (that does not implement `BufferizableOpInterface`) is found. No
-// to_tensor/to_memref ops are inserted.
-//
-// 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, unless `allowReturnMemref` is enabled.
-//
-//  Lastly, note that layout map chosen to bufferize is the most dynamic
-//  canonical strided layout of the proper rank. This ensures compatibility with
-//  expected layouts after transformations. Combinations of memref.cast +
-//  canonicalization are responsible for clean ups.
+// This analysis caters to high-performance codegen where buffer reuse is deemed
+// critical: the analysis should fail if the bufferized form of the function
+// needs to return a buffer, unless `allowReturnMemref` is enabled.
 
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h"
+#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 
 #include <random>
 
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
-using namespace linalg;
-using namespace tensor;
-using namespace comprehensive_bufferize;
+using namespace mlir::bufferization;
 
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific attribute manipulation.
 // These are for testing and debugging only. Bufferization information is
 // stored in BufferizationAliasInfo. When run with `testAnalysisOnly`, the IR
 // is annotated with the results of the analysis (copied from
op->walk([&](Operation *returnOp) {
 
       return WalkResult::advance();
     });
 
     return status;
   }
 };
 
-/// Rewrite pattern that bufferizes bufferizable ops.
-struct BufferizationPattern
-    : public OpInterfaceRewritePattern<BufferizableOpInterface> {
-  BufferizationPattern(MLIRContext *context, const BufferizationState &state,
-                       PatternBenefit benefit = 1)
-      : OpInterfaceRewritePattern<BufferizableOpInterface>(context, benefit),
-        state(state) {}
-
-  LogicalResult matchAndRewrite(BufferizableOpInterface bufferizableOp,
-                                PatternRewriter &rewriter) const override {
-    // No tensors => no buffers.
-    if (!hasTensorSemantics(bufferizableOp.getOperation()))
-      return failure();
-    if (!state.getOptions().isOpAllowed(bufferizableOp.getOperation()))
-      return failure();
-    return bufferizableOp.bufferize(rewriter, state);
-  }
-
-private:
-  const BufferizationState &state;
-};
-
-/// Check the result of bufferization. Return an error if an op was not
-/// bufferized, unless partial bufferization is allowed.
-static LogicalResult
-checkBufferizationResult(Operation *op, const BufferizationOptions &options) {
-  if (!options.allowUnknownOps) {
-    // Check if all ops were bufferized.
-    LogicalResult status = success();
-    op->walk([&](Operation *op) {
-      if (!hasTensorSemantics(op))
-        return WalkResult::advance();
-
-      // Bufferization dialect ops will canonicalize away if all other ops are
-      // bufferized.
-      if (isa<bufferization::ToMemrefOp, bufferization::ToTensorOp>(op))
-        return WalkResult::advance();
-
-      // Ops that are not in the allow list can be ignored.
-      if (!options.isOpAllowed(op))
-        return WalkResult::advance();
-
-      // Ops without any uses and no side effects will fold away.
-      if (op->getUses().empty() && MemoryEffectOpInterface::hasNoEffect(op))
-        return WalkResult::advance();
-
-      status = op->emitError("op was not bufferized");
-      return WalkResult::interrupt();
-    });
-
-    if (failed(status))
-      return status;
-  }
-
-  return success();
-}
-
-LogicalResult mlir::linalg::comprehensive_bufferize::analyzeOp(
-    Operation *op, AnalysisBufferizationState &state) {
+LogicalResult bufferization::analyzeOp(Operation *op,
+                                       AnalysisBufferizationState &state) {
   DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
   const auto &options =
       static_cast<const AnalysisBufferizationOptions &>(state.getOptions());
 
   if (failed(checkAliasInfoConsistency(op, domInfo, state, aliasInfo)))
     return failure();
 
LogicalResult bufferization::analyzeOp(Operation *op,
 
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly)
     annotateOpsWithBufferizationMarkers(op, aliasInfo, state);
 
   return success();
 }
 
-LogicalResult mlir::linalg::comprehensive_bufferize::bufferizeOp(
-    Operation *op, const BufferizationState &state) {
-  // Bufferize the op and its nested ops.
-  OwningRewritePatternList patterns(op->getContext());
-  patterns.add<BufferizationPattern>(op->getContext(), state);
-  if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns))))
-    return failure();
-
-  return checkBufferizationResult(op, state.getOptions());
-}
-
-LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
+LogicalResult bufferization::runOneShotBufferize(
     Operation *op, std::unique_ptr<AnalysisBufferizationOptions> options) {
   AnalysisBufferizationState state(op, *options);
   if (failed(analyzeOp(op, state)))
     return failure();
   if (options->testAnalysisOnly)
     return success();
   return bufferizeOp(op, state);
 }

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

 //===- AffineInterfaceImpl.cpp - Affine Impl. of BufferizableOpInterface --===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/AffineInterfaceImpl.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+
+using namespace mlir::bufferization;
 
 void mlir::linalg::comprehensive_bufferize::affine_ext::
     registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
   // AffineParallelOp bufferization not implemented yet. However, never hoist
   // memref allocations across AffineParallelOp boundaries.
   registry.addOpInterface<AffineParallelOp,
                           AllocationHoistingBarrierOnly<AffineParallelOp>>();
 }

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

 //===- ArithInterfaceImpl.cpp - Arith Impl. of BufferizableOpInterface ----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ArithInterfaceImpl.h"
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Transforms/BufferUtils.h"
 
+using namespace mlir::bufferization;
+
 namespace mlir {
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace arith_ext {
 
 /// Bufferization of arith.constant. Replace with memref.get_global.
 struct ConstantOpInterface
     : public BufferizableOpInterface::ExternalModel<ConstantOpInterface,

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   AffineInterfaceImpl.cpp
   ArithInterfaceImpl.cpp
-  BufferizableOpInterface.cpp
-  BufferizationInterfaceImpl.cpp
-  ComprehensiveBufferize.cpp
   LinalgInterfaceImpl.cpp
   ModuleBufferization.cpp
   SCFInterfaceImpl.cpp
   StdInterfaceImpl.cpp
   TensorInterfaceImpl.cpp
   VectorInterfaceImpl.cpp
 )
 
-add_mlir_dialect_library(MLIRBufferizableOpInterface
-  BufferizableOpInterface.cpp
-
-  DEPENDS
-  MLIRBufferizableOpInterfaceIncGen
-
-  LINK_LIBS PUBLIC
-  MLIRIR
-  MLIRBufferization
-  MLIRMemRef
-)
-
 add_mlir_dialect_library(MLIRAffineBufferizableOpInterfaceImpl
   AffineInterfaceImpl.cpp
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRBufferizableOpInterface
 )
 
 add_mlir_dialect_library(MLIRArithBufferizableOpInterfaceImpl
   ArithInterfaceImpl.cpp
 
   LINK_LIBS PUBLIC
   MLIRArithmetic
   MLIRBufferizableOpInterface
   MLIRIR
   MLIRMemRef
   MLIRStandardOpsTransforms
 )
 
 add_mlir_dialect_library(MLIRLinalgBufferizableOpInterfaceImpl
   LinalgInterfaceImpl.cpp
 
   LINK_LIBS PUBLIC
   MLIRBufferizableOpInterface
-  MLIRComprehensiveBufferize
+  MLIRBufferizationTransforms
   MLIRIR
   MLIRLinalg
   MLIRTensor
 )
 
 add_mlir_dialect_library(MLIRSCFBufferizableOpInterfaceImpl
   SCFInterfaceImpl.cpp
 
   LINK_LIBS PUBLIC
   MLIRBufferizableOpInterface
-  MLIRComprehensiveBufferize
+  MLIRBufferizationTransforms
   MLIRIR
   MLIRSCF
 )
 
 add_mlir_dialect_library(MLIRStdBufferizableOpInterfaceImpl
   StdInterfaceImpl.cpp
 
   LINK_LIBS PUBLIC
add_mlir_dialect_library(MLIRVectorBufferizableOpInterfaceImpl
   VectorInterfaceImpl.cpp
 
   LINK_LIBS PUBLIC
   MLIRBufferizableOpInterface
   MLIRIR
   MLIRVector
 )
 
-add_mlir_dialect_library(MLIRComprehensiveBufferize
-  BufferizationInterfaceImpl.cpp
-  ComprehensiveBufferize.cpp
+add_mlir_dialect_library(MLIRModuleBufferization
   ModuleBufferization.cpp
 
   LINK_LIBS PUBLIC
   MLIRBufferizableOpInterface
-  MLIRControlFlowInterfaces
-  MLIRInferTypeOpInterface
+  MLIRBufferizationTransforms
   MLIRIR
   MLIRMemRef
   MLIRStandard
   MLIRStandardOpsTransforms
-  MLIRTransforms
 )

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

 //===- LinalgInterfaceImpl.cpp - Linalg Impl. of BufferizableOpInterface --===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
 using namespace linalg;
 using namespace comprehensive_bufferize;
+using namespace mlir::bufferization;
 
 namespace {
 
 // TODO: Ops in the linalg dialect can directly implement this interface.
 
 /// Generic conversion for any LinalgOp on tensors.
 static LogicalResult bufferizeLinalgOp(RewriterBase &rewriter, LinalgOp op,
                                        const BufferizationState &state) {

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

 // analyze the function body. In such a case, the CallOp analysis conservatively
 // assumes that each tensor OpOperand is both read and written.
 //
 // TODO: Add FuncOp attributes so that bbArgs of external FuncOps can be marked
 // as "not reading" and/or "not writing".
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h"
 
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h"
+#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
+#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
 using namespace linalg;
 using namespace tensor;
 using namespace comprehensive_bufferize;
+using namespace mlir::bufferization;
 
 namespace {
 /// The state of analysis of a FuncOp.
 enum class FuncOpAnalysisState { NotAnalyzed, InProgress, Analyzed };
 
 /// Extra bufferization state that is required for bufferization of function
 /// boundaries.
 struct ModuleBufferizationState : public DialectBufferizationState {

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

 //===- SCFInterfaceImpl.cpp - SCF Impl. of BufferizableOpInterface --------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
 
+using namespace mlir::bufferization;
+
 namespace mlir {
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace scf_ext {
 
 // bufferization.to_memref is not allowed to change the rank.
 static void ensureToMemrefOpIsValid(Value tensor, Type memrefType) {
 #ifndef NDEBUG

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

 //===- StdInterfaceImpl.cpp - Standard Impl. of BufferizableOpInterface ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/StdInterfaceImpl.h"
 
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 
+using namespace mlir::bufferization;
+
 namespace mlir {
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace std_ext {
 
 /// Bufferization of std.select. Just replace the operands.
 struct SelectOpInterface
     : public BufferizableOpInterface::ExternalModel<SelectOpInterface,

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

 //===- TensorInterfaceImpl.cpp - Tensor Impl. of BufferizableOpInterface --===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/TensorInterfaceImpl.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
+using namespace mlir::bufferization;
 
 namespace mlir {
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace tensor_ext {
 
 using tensor::ExtractSliceOp;
 using tensor::InsertSliceOp;

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

 //===- VectorInterfaceImpl.cpp - Vector Impl. of BufferizableOpInterface --===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/VectorInterfaceImpl.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 
+using namespace mlir::bufferization;
+
 namespace mlir {
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace vector_ext {
 
 /// Bufferization of vector.transfer_read. Replaced with a new
 /// vector.transfer_read that operates on a memref.
 struct TransferReadOpInterface

mlir/lib/Dialect/Linalg/IR/LinalgDialect.cpp

 //===- Dialect.cpp - Implementation of the linalg dialect and types -------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the Linalg dialect types and dialect.
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/Parser.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/InliningUtils.h"

 
 /// Attribute name used to to memoize indexing maps for named ops.
 constexpr const ::llvm::StringLiteral
     LinalgDialect::kMemoizedIndexingMapsAttrName;
 
 /// Attribute name used to mark the bufferization layout for region
 /// arguments during linalg comprehensive bufferization.
 constexpr const ::llvm::StringLiteral
-    comprehensive_bufferize::BufferizableOpInterface::kBufferLayoutAttrName;
+    bufferization::BufferizableOpInterface::kBufferLayoutAttrName;
 
 /// Attribute name used to mark region arguments that can be bufferized
 /// in-place during linalg comprehensive bufferization.
 constexpr const ::llvm::StringLiteral
-    comprehensive_bufferize::BufferizableOpInterface::kInplaceableAttrName;
+    bufferization::BufferizableOpInterface::kInplaceableAttrName;
 
 /// Trait to check if T provides a `regionBuilder` method.
 template <typename T, typename... Args>
 using has_region_builder = decltype(T::regionBuilder);
 template <typename T>
 using detect_has_region_builder = llvm::is_detected<has_region_builder, T>;
 
 /// SFINAE helper for single C++ class without a `regionBuilder` method (e.g.

 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
       >(namedStructuredOpRegionBuilders);
 
   addInterfaces<LinalgInlinerInterface>();
 }
 
 LogicalResult LinalgDialect::verifyOperationAttribute(Operation *op,
                                                       NamedAttribute attr) {
-  using comprehensive_bufferize::BufferizableOpInterface;
+  using bufferization::BufferizableOpInterface;
 
   if (attr.getName() == BufferizableOpInterface::kInplaceableAttrName) {
     if (!attr.getValue().isa<BoolAttr>()) {
       return op->emitError()
              << "'" << BufferizableOpInterface::kInplaceableAttrName
              << "' is expected to be a boolean attribute";
     }
     if (!isa<FunctionOpInterface>(op))

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRLinalgTransforms
   MLIRAffine
   MLIRAffineBufferizableOpInterfaceImpl
   MLIRAffineUtils
   MLIRAnalysis
   MLIRArithBufferizableOpInterfaceImpl
   MLIRArithmetic
   MLIRBufferizableOpInterface
   MLIRComplex
-  MLIRComprehensiveBufferize
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRMemRef
   MLIRLinalg
   MLIRLinalgAnalysis
   MLIRLinalgBufferizableOpInterfaceImpl
   MLIRLinalgUtils
+  MLIRModuleBufferization
   MLIRSCF
   MLIRSCFBufferizableOpInterfaceImpl
   MLIRSCFTransforms
   MLIRStdBufferizableOpInterfaceImpl
   MLIRPass
   MLIRStandard
   MLIRStandardOpsTransforms
   MLIRStandardToLLVM

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

 //===- ComprehensiveBufferize.cpp - Single pass bufferization -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizationInterfaceImpl.h"
+#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/AffineInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ArithInterfaceImpl.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizationInterfaceImpl.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/StdInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/TensorInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/VectorInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
+using namespace mlir::bufferization;
 using namespace mlir::linalg;
 using namespace mlir::linalg::comprehensive_bufferize;
 
 namespace {
 struct LinalgComprehensiveModuleBufferize
     : public LinalgComprehensiveModuleBufferizeBase<
           LinalgComprehensiveModuleBufferize> {
   LinalgComprehensiveModuleBufferize() = default;

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

 //===- Sparsification.cpp - Implementation of sparsification --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements converting sparse tensor types to actual sparse code.
 //
 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/SCF/Transforms.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/SparseTensor/Utils/Merger.h"

 }
 
 /// Returns true if tensor has an in-place annotation.
 static bool isInPlace(Value val) {
   if (auto arg = val.dyn_cast<BlockArgument>())
     if (auto funcOp = dyn_cast<FuncOp>(arg.getOwner()->getParentOp()))
       if (auto attr = funcOp.getArgAttrOfType<BoolAttr>(
               arg.getArgNumber(),
-              linalg::comprehensive_bufferize::BufferizableOpInterface::
-                  kInplaceableAttrName))
+              bufferization::BufferizableOpInterface::kInplaceableAttrName))
         return attr.getValue();
   return false;
 }
 
 /// Returns true if tensor materializes uninitialized into the computation.
 static bool isMaterializing(Value val) {
   return val.getDefiningOp<linalg::InitTensorOp>() ||
          val.getDefiningOp<InitOp>();

mlir/test/lib/Dialect/Linalg/CMakeLists.txt

add_mlir_library(MLIRLinalgTestPasses
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRAffineBufferizableOpInterfaceImpl
   MLIRArithBufferizableOpInterfaceImpl
   MLIRArithmetic
   MLIRBufferizableOpInterface
-  MLIRComprehensiveBufferize
+  MLIRBufferizationTransforms
   MLIRGPUTransforms
   MLIRLinalg
   MLIRLinalgBufferizableOpInterfaceImpl
   MLIRLinalgTransforms
   MLIRLLVMToLLVMIRTranslation
   MLIRMemRef
   MLIRPass
   MLIRSCF

mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp

 //===- TestComprehensiveBufferize.cpp - Test Comprehensive Bufferize ------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements logic for testing Comprehensive Bufferize.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizationInterfaceImpl.h"
+#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/AffineInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ArithInterfaceImpl.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizationInterfaceImpl.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/StdInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/TensorInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/VectorInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::linalg::comprehensive_bufferize;
+using namespace mlir::bufferization;
 
 namespace {
 /// A helper struct for FunctionBufferize and ModuleBufferize. Both passes are
 /// mostly identical.
 struct TestComprehensiveFunctionBufferize
     : public PassWrapper<TestComprehensiveFunctionBufferize,
                          OperationPass<FuncOp>> {
   StringRef getArgument() const final {
void TestComprehensiveFunctionBufferize::runOnOperation() {
   options->createDeallocs = createDeallocs;
 
   if (dialectFilter.hasValue()) {
     options->dialectFilter.emplace();
     for (const std::string &dialectNamespace : dialectFilter)
       options->dialectFilter->insert(dialectNamespace);
   }
 
-  Operation *op = getOperation().getOperation();
-  if (failed(runComprehensiveBufferize(op, std::move(options))))
+  Operation *op = getOperation();
+  if (failed(runOneShotBufferize(op, std::move(options))))
     return;
 
   if (testAnalysisOnly)
     return;
 
   OpPassManager cleanupPipeline("builtin.func");
   cleanupPipeline.addPass(createCanonicalizerPass());
   cleanupPipeline.addPass(createCSEPass());

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

gentbl_cc_library(
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td",
     deps = [":LinalgStructuredOpsTdFiles"],
 )
 
 td_library(
     name = "BufferizableOpInterfaceTdFiles",
     srcs = [
-        "include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td",
+        "include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td",
     ],
     includes = ["include"],
     deps = [
         ":OpBaseTdFiles",
     ],
 )
 
 gentbl_cc_library(
     name = "BufferizableOpInterfaceIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (
             ["-gen-op-interface-decls"],
-            "include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h.inc",
+            "include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h.inc",
         ),
         (
             ["-gen-op-interface-defs"],
-            "include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp.inc",
+            "include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.cpp.inc",
         ),
     ],
     tblgen = ":mlir-tblgen",
-    td_file = "include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td",
+    td_file = "include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td",
     deps = [
         ":BufferizableOpInterfaceTdFiles",
     ],
 )
 
 cc_library(
     name = "BufferizableOpInterface",
     srcs = [
-        "lib/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp",
+        "lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp",
+        "lib/Dialect/Bufferization/IR/BufferizationInterfaceImpl.cpp",
     ],
     hdrs = [
-        "include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h",
+        "include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h",
+        "include/mlir/Dialect/Bufferization/IR/BufferizationInterfaceImpl.h",
     ],
     includes = ["include"],
     deps = [
         ":BufferizableOpInterfaceIncGen",
         ":BufferizationDialect",
         ":IR",
         ":MemRefDialect",
         ":Support",
cc_library(
     ],
     hdrs = [
         "include/mlir/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.h",
     ],
     includes = ["include"],
     deps = [
         ":BufferizableOpInterface",
         ":BufferizationDialect",
-        ":ComprehensiveBufferize",
+        ":BufferizationTransforms",
         ":IR",
         ":LinalgOps",
         ":LinalgStructuredOpsIncGen",
         ":TensorDialect",
     ],
 )
 
 cc_library(
     name = "SCFBufferizableOpInterfaceImpl",
     srcs = [
         "lib/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.cpp",
     ],
     hdrs = [
         "include/mlir/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.h",
     ],
     includes = ["include"],
     deps = [
         ":BufferizableOpInterface",
         ":BufferizationDialect",
-        ":ComprehensiveBufferize",
+        ":BufferizationTransforms",
         ":IR",
         ":SCFDialect",
         ":Support",
         "//llvm:Support",
     ],
 )
 
 cc_library(
deps = [
         ":AffineUtils",
         ":Analysis",
         ":ArithBufferizableOpInterfaceImpl",
         ":ArithmeticDialect",
         ":BufferizableOpInterface",
         ":BufferizationDialect",
         ":BufferizationTransforms",
         ":ComplexDialect",
-        ":ComprehensiveBufferize",
         ":DialectUtils",
         ":IR",
         ":InferTypeOpInterface",
         ":LinalgBufferizableOpInterfaceImpl",
         ":LinalgOps",
         ":LinalgPassIncGen",
         ":LinalgStructuredOpsIncGen",
         ":MathDialect",
         ":MemRefDialect",
+        ":ModuleBufferization",
         ":Pass",
         ":SCFBufferizableOpInterfaceImpl",
         ":SCFDialect",
         ":SCFTransforms",
         ":StandardOps",
         ":StandardOpsTransforms",
         ":StdBufferizableOpInterfaceImpl",
         ":Support",
         ":TensorBufferizableOpInterfaceImpl",
         ":TensorDialect",
         ":TransformUtils",
         ":VectorBufferizableOpInterfaceImpl",
         ":VectorOps",
         ":VectorToSCF",
         ":X86VectorTransforms",
         "//llvm:Support",
     ],
 )
 
 cc_library(
-    name = "ComprehensiveBufferize",
+    name = "ModuleBufferization",
     srcs = [
-        "lib/Dialect/Linalg/ComprehensiveBufferize/BufferizationInterfaceImpl.cpp",
-        "lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp",
         "lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp",
     ],
     hdrs = [
-        "include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizationInterfaceImpl.h",
-        "include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h",
         "include/mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h",
     ],
     includes = ["include"],
     deps = [
         ":BufferizableOpInterface",
         ":BufferizationDialect",
-        ":ControlFlowInterfaces",
+        ":BufferizationTransforms",
         ":DialectUtils",
         ":IR",
-        ":InferTypeOpInterface",
         ":MemRefDialect",
-        ":Pass",
         ":StandardOps",
         ":Support",
-        ":Transforms",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "TilingInterface",
     srcs = ["lib/Interfaces/TilingInterface.cpp"],
     hdrs = ["include/mlir/Interfaces/TilingInterface.h"],
gentbl_cc_library(
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/Bufferization/IR/BufferizationOps.td",
     deps = [":BufferizationOpsTdFiles"],
 )
 
 cc_library(
     name = "BufferizationDialect",
-    srcs = glob(
-        [
-            "lib/Dialect/Bufferization/IR/Bufferization*.h",
-            "lib/Dialect/Bufferization/IR/Bufferization*.cpp",
-        ],
-    ),
+    srcs = [
+        "lib/Dialect/Bufferization/IR/BufferizationDialect.cpp",
+        "lib/Dialect/Bufferization/IR/BufferizationOps.cpp",
     hdrs = ["include/mlir/Dialect/Bufferization/IR/Bufferization.h"],
     includes = ["include"],
     deps = [
         ":AllocationOpInterface",
         ":BufferizationBaseIncGen",
         ":BufferizationOpsIncGen",
         ":ControlFlowInterfaces",
         ":CopyOpInterface",
srcs = glob(
             "lib/Dialect/Bufferization/Transforms/*.h",
         ],
     ),
     hdrs = glob(["include/mlir/Dialect/Bufferization/Transforms/*.h"]),
     includes = ["include"],
     deps = [
         ":AllocationOpInterface",
         ":Analysis",
+        ":BufferizableOpInterface",
         ":BufferizationDialect",
         ":BufferizationPassIncGen",
+        ":ControlFlowInterfaces",
+        ":DialectUtils",
         ":IR",
+        ":InferTypeOpInterface",
         ":MemRefDialect",
         ":Pass",
+        ":StandardOps",
+        ":Support",
         ":Transforms",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "BufferizationToMemRef",
     srcs = [

utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel

cc_library(
     deps = [
         "//llvm:Support",
         "//mlir:Affine",
         "//mlir:AffineBufferizableOpInterfaceImpl",
         "//mlir:ArithBufferizableOpInterfaceImpl",
         "//mlir:ArithmeticDialect",
         "//mlir:BufferizableOpInterface",
         "//mlir:BufferizationDialect",
-        "//mlir:ComprehensiveBufferize",
+        "//mlir:BufferizationTransforms",
         "//mlir:GPUDialect",
         "//mlir:IR",
         "//mlir:LinalgBufferizableOpInterfaceImpl",
         "//mlir:LinalgOps",
         "//mlir:LinalgTransforms",
         "//mlir:MemRefDialect",
         "//mlir:Pass",
         "//mlir:SCFBufferizableOpInterfaceImpl",