[mlir][linalg][bufferize] Separate pass from ComprehensiveBufferize

[springerm]

Summary

This commit separates the bufferization from the bufferization pass in Linalg. This allows other dialects to use ComprehensiveBufferize more easily.

This commit mainly moves files to a new directory and adds a new build target.

Depends On D112991

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

+add_subdirectory(ComprehensiveBufferize)
 add_subdirectory(IR)
-add_subdirectory(Transforms)
 
 set(LLVM_TARGET_DEFINITIONS Passes.td)
 mlir_tablegen(Passes.h.inc -gen-pass-decls -name Linalg)
 mlir_tablegen(Passes.capi.h.inc -gen-pass-capi-header --prefix Linalg)
 mlir_tablegen(Passes.capi.cpp.inc -gen-pass-capi-impl --prefix Linalg)
 add_public_tablegen_target(MLIRLinalgPassIncGen)
 
 add_mlir_doc(Passes LinalgPasses ./ -gen-pass-doc)

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

 //===- BufferizableOpInterface.h - Comprehensive Bufferize ------*- 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_TRANSFORMS_BUFFERIZABLEOPINTERFACE_H_
-#define MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE_H_
+#ifndef MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZABLEOPINTERFACE_H_
+#define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZABLEOPINTERFACE_H_
 
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
 class BlockAndValueMapping;
 
 namespace linalg {
 struct AllocationCallbacks;
 class BufferizationAliasInfo;
 
 /// Specify fine-grain relationship between buffers to enable more analysis.
 enum class BufferRelation {
   None,
   // TODO: ResultContainsOperand,
   // TODO: OperandContainsResult,
   Equivalent
 };
 } // namespace linalg
 } // namespace mlir
 
-#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.h.inc"
+#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h.inc"
 
-#endif // MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE_H_
+#endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZABLEOPINTERFACE_H_

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

 //===-- BufferizableOpInterface.td - Compreh. Bufferize ----*- 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 MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE
-#define MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE
+#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 methods = [
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return true;
         }]
       >
   ];
 }
 
-#endif  // MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE
+#endif  // BUFFERIZABLE_OP_INTERFACE

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

 //===- ComprehensiveBufferize.h - Linalg bufferization pass -----*- 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_TRANSFORMS_COMPREHENSIVE_BUFFERIZE_H
-#define MLIR_DIALECT_LINALG_TRANSFORMS_COMPREHENSIVE_BUFFERIZE_H
+#ifndef MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVE_BUFFERIZE_H
+#define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVE_BUFFERIZE_H
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Value.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/SetOperations.h"
 
 namespace mlir {
 
 class DominanceInfo;
 class FuncOp;
 class GlobalCreator;
+class ModuleOp;
 
 namespace linalg {
 
+// TODO: from some HW description.
+static constexpr int64_t kBufferAlignments = 128;
+
 /// The BufferizationAliasInfo class maintains a list of buffer aliases and
 /// equivalence classes to support bufferization.
 /// ExtractSliceOps have special behavior, they act as a level of indirection
 /// for bufferization. They don't create reads or writes themselves and analysis
 /// needs to look through their uses.
 /// ExtractSliceOp + InsertSliceOp have special joint behavior: they may
 /// bufferize to the same buffer (i.e. subview), which is what introduces the
 /// need for bufferization classes.

 };
 
 /// Analyze the `ops` to determine which OpResults are inplaceable.
 LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
                               BufferizationAliasInfo &aliasInfo,
                               const DominanceInfo &domInfo,
                               unsigned analysisFuzzerSeed = 0);
 
+// TODO: Do not expose those functions in the header file.
 /// Default allocation function that is used by the comprehensive bufferization
 /// pass. The default currently creates a ranked memref using `memref.alloc`.
 Optional<Value> defaultAllocationFn(OpBuilder &b, Location loc, MemRefType type,
                                     const SmallVector<Value> &dynShape);
 
 /// Default deallocation function that is used by the comprehensive
 /// bufferization pass. It expects to recieve back the value called from the
 /// `defaultAllocationFn`.
using AllocationFn = std::function<Optional<Value>(
       OpBuilder &, Location, MemRefType, const SmallVector<Value> &)>;
   using DeallocationFn = std::function<void(OpBuilder &, Location, Value)>;
   using MemCpyFn = std::function<void(OpBuilder &, Location, Value, Value)>;
 
   AllocationCallbacks(AllocationFn allocFn, DeallocationFn deallocFn,
                       MemCpyFn copyFn)
       : allocationFn(allocFn), deallocationFn(deallocFn), memCpyFn(copyFn) {}
 
+  AllocationCallbacks()
+      : allocationFn(defaultAllocationFn),
+        deallocationFn(defaultDeallocationFn), memCpyFn(defaultMemCpyFn) {}
+
   AllocationFn allocationFn;
   DeallocationFn deallocationFn;
   MemCpyFn memCpyFn;
 };
 
 /// Bufferize one particular op.
 /// `bufferizedFunctionTypes` (resp. `globalCreator`) are expected to be
 /// non-null if `op` is a CallOpInterface (resp. GlobalCreator).
LogicalResult initTensorElimination(
     bool skipAnalysis = false);
 
 /// Try to eliminate InitTensorOps inside funcOp that are anchored on an
 /// InsertSliceOp, i.e., if it is eventually inserted into another tensor
 /// (and some other conditions are met).
 LogicalResult eliminateInsertSliceAnchoredInitTensorOps(
     FuncOp funcOp, BufferizationAliasInfo &aliasInfo, DominanceInfo &domInfo);
 
+struct BufferizationOptions {
+  BufferizationOptions()
+      : allocationFns(std::make_unique<AllocationCallbacks>()) {}
+
+  std::unique_ptr<AllocationCallbacks> allocationFns;
+  bool allowReturnMemref = false;
+  unsigned analysisFuzzerSeed = 0;
+  bool testAnalysisOnly = false;
+};
+
+LogicalResult runComprehensiveBufferize(ModuleOp moduleOp,
+                                        const BufferizationOptions &options);
+
 } // namespace linalg
 } // namespace mlir
 
-#endif // define MLIR_DIALECT_LINALG_TRANSFORMS_COMPREHENSIVE_BUFFERIZE_H
+#endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVE_BUFFERIZE_H

mlir/lib/Dialect/Linalg/CMakeLists.txt

 add_subdirectory(Analysis)
+add_subdirectory(ComprehensiveBufferize)
 add_subdirectory(IR)
 add_subdirectory(Transforms)
 add_subdirectory(Utils)

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

 //===- BufferizableOpInterface.cpp - 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
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.h"
+#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
 
 namespace mlir {
 namespace linalg {
 
-#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.cpp.inc"
+#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp.inc"
 
 } // namespace linalg
 } // namespace mlir

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

+set(LLVM_OPTIONAL_SOURCES
+  BufferizableOpInterface.cpp
+  ComprehensiveBufferize.cpp
+)
+
+add_mlir_dialect_library(MLIRBufferizableOpInterface
+  BufferizableOpInterface.cpp
+
+  DEPENDS
+  MLIRBufferizableOpInterfaceIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+)
+
+add_mlir_dialect_library(MLIRComprehensiveBufferize
+  ComprehensiveBufferize.cpp
+
+  LINK_LIBS PUBLIC
+  MLIRBufferizableOpInterface
+  MLIRInferTypeOpInterface
+  MLIRIR
+  MLIRMemRef
+  MLIRLinalg
+  MLIRSCF
+  MLIRStandard
+  MLIRStandardOpsTransforms
+  MLIRTensor
+  MLIRVector
+)

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

 //  ops to encode the desired semantics at function boundaries for i., ii. and
 //  iii.
 //
 //  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.
 
-#include "mlir/Dialect/Linalg/Transforms/ComprehensiveBufferize.h"
+#include "mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h"
 
 #include <random>
 
-#include "PassDetail.h"
+#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
-#include "mlir/Dialect/Linalg/Passes.h"
-#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.h"
-#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
-#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/BufferUtils.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/FormatVariadic.h"
 
 #define DEBUG_TYPE "comprehensive-module-bufferize"
 
 using namespace mlir;
 using namespace linalg;
 using namespace tensor;
 
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X)
 
-// TODO: from some HW description.
-static constexpr int64_t kBufferAlignments = 128;
-
 // Forward declarations.
 static std::string printOperationInfo(Operation *, bool prefix = true);
 static std::string printValueInfo(Value, bool prefix = true);
 
 //===----------------------------------------------------------------------===//
 // Generic helpers.
 //===----------------------------------------------------------------------===//
 
Operation *getFirstParentOfType(Value v) {
   while (parent) {
     if (isa<Args...>(parent))
       return parent;
     parent = parent->getParentOp();
   }
   return nullptr;
 }
 
+/// Helper function that creates a memref::DimOp or tensor::DimOp depending on
+/// the type of `source`.
+static Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source,
+                               int64_t dim) {
+  if (source.getType().isa<UnrankedMemRefType, MemRefType>())
+    return b.createOrFold<memref::DimOp>(loc, source, dim);
+  if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
+    return b.createOrFold<tensor::DimOp>(loc, source, dim);
+  llvm_unreachable("Expected MemRefType or TensorType");
+}
+
 /// Compute the type of the `memref` to use for allocating the buffer for
 /// `shapedValue`. Also returns (by reference in `dynShape`), the value for the
 /// dynamic dimensions in the returned `memref` type. The function also sets the
 /// insertion point of the builder `b` to the position where the allocation is
 /// to be inserted.
 static MemRefType getAllocationTypeAndShape(OpBuilder &b, Location loc,
                                             Value shapedValue,
                                             SmallVectorImpl<Value> &dynShape) {

 Optional<Value>
 mlir::linalg::defaultAllocationFn(OpBuilder &b, Location loc, MemRefType type,
                                   const SmallVector<Value> &dynShape) {
   Value allocated = b.create<memref::AllocOp>(
       loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
   return allocated;
 }
 
-static Optional<Value>
-allocationFnUsingAlloca(OpBuilder &b, Location loc, MemRefType type,
-                        const SmallVector<Value> &dynShape) {
-  Value allocated = b.create<memref::AllocaOp>(
-      loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
-  return allocated;
-}
-
 void mlir::linalg::defaultDeallocationFn(OpBuilder &b, Location loc,
                                          Value allocatedBuffer) {
   b.create<memref::DeallocOp>(loc, allocatedBuffer);
 }
 
 void mlir::linalg::defaultMemCpyFn(OpBuilder &b, Location loc, Value from,
                                    Value to) {
   b.create<CopyOp>(loc, from, to);
while (!numberCallOpsContainedInFuncOp.empty()) {
     orderedFuncOps.push_back(it->getFirst());
     for (auto callee : calledBy[it->getFirst()])
       numberCallOpsContainedInFuncOp[callee]--;
     numberCallOpsContainedInFuncOp.erase(it);
   }
   return success();
 }
 
-namespace {
-struct LinalgComprehensiveModuleBufferize
-    : public LinalgComprehensiveModuleBufferizeBase<
-          LinalgComprehensiveModuleBufferize> {
-  LinalgComprehensiveModuleBufferize() {}
-
-  LinalgComprehensiveModuleBufferize(
-      const LinalgComprehensiveModuleBufferize &p) {}
-
-  void runOnOperation() override;
-
-  void getDependentDialects(DialectRegistry &registry) const override {
-    registry
-        .insert<linalg::LinalgDialect, memref::MemRefDialect,
-                tensor::TensorDialect, vector::VectorDialect, scf::SCFDialect,
-                arith::ArithmeticDialect, StandardOpsDialect>();
-    registerBufferizableOpInterfaceExternalModels(registry);
-  }
-
-private:
-  std::unique_ptr<AllocationCallbacks> allocationFns;
-};
-} // end namespace
-
-static void applyEnablingTransformations(ModuleOp moduleOp) {
-  RewritePatternSet patterns(moduleOp.getContext());
-  patterns.add<GeneralizePadTensorOpPattern>(moduleOp.getContext());
-  (void)applyPatternsAndFoldGreedily(moduleOp, std::move(patterns));
-}
-
 static void
 foreachCaller(const DenseMap<FuncOp, DenseSet<Operation *>> &callerMap,
               FuncOp callee, llvm::function_ref<void(Operation *)> doit) {
   auto itCallers = callerMap.find(callee);
   if (itCallers == callerMap.end())
     return;
   for (Operation *caller : itCallers->second)
     doit(caller);
return initTensorElimination(
         auto insertSliceOp = cast<InsertSliceOp>(operand.getOwner());
         auto extractOp = b.create<tensor::ExtractSliceOp>(
             loc, insertSliceOp.dest(), insertSliceOp.getMixedOffsets(),
             insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
         return extractOp.result();
       });
 }
 
-void LinalgComprehensiveModuleBufferize::runOnOperation() {
-  if (!allocationFns) {
-    // The allocation functions to use needs to be set here. The flag for the
-    // pass and flag for the use of alloca map to LLVM command line
-    // options. These being static global objects have no set order in which
-    // they are defined. So ideally this should be in the constructor, but the
-    // constructor might be called before the flag is initialized using the
-    // command line option. So this is set up at the start of the pass.
-    if (useAlloca) {
-      AllocationCallbacks allocaAllocationFns = {
-          allocationFnUsingAlloca, [](OpBuilder &b, Location loc, Value v) {},
-          defaultMemCpyFn};
-      allocationFns =
-          std::make_unique<AllocationCallbacks>(std::move(allocaAllocationFns));
-    } else {
-      allocationFns = std::make_unique<AllocationCallbacks>(
-          defaultAllocationFn, defaultDeallocationFn, defaultMemCpyFn);
-    }
-  }
-  ModuleOp moduleOp = getOperation();
-  applyEnablingTransformations(moduleOp);
-
+LogicalResult
+mlir::linalg::runComprehensiveBufferize(ModuleOp moduleOp,
+                                        const BufferizationOptions &options) {
   SmallVector<FuncOp> orderedFuncOps;
   DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
   DenseMap<FuncOp, FunctionType> bufferizedFunctionTypes;
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
-    return signalPassFailure();
+    return failure();
 
   DominanceInfo domInfo(moduleOp);
   BufferizationAliasInfo aliasInfo(moduleOp);
   // Interestingly, all function args that are not visible outside of a module
   // can be fully bufferized inplace by guaranteeing the CallOp is bufferized
   // inplace. Therefore, we just bufferize funcOp as if none of its results were
   // inplaceable, detect which operands are cloned internally and decide what to
   // do at call sites.
for (FuncOp funcOp : orderedFuncOps) {
     // TODO: do not rely on this behavior.
     if (callerMap.find(funcOp) != callerMap.end())
       for (BlockArgument bbArg : funcOp.getArguments())
         if (bbArg.getType().isa<TensorType>())
           setInPlaceFuncArgument(bbArg);
 
     // If the analysis fails, just return.
     if (failed(inPlaceAnalysisFuncOpBody(funcOp, aliasInfo, domInfo,
-                                         analysisFuzzerSeed))) {
-      signalPassFailure();
-      return;
-    }
+                                         options.analysisFuzzerSeed)))
+      return failure();
 
     // Try to eliminate InitTensorOps to avoid new allocations during the
     // bufferization phase.
     if (failed(eliminateInsertSliceAnchoredInitTensorOps(funcOp, aliasInfo,
-                                                         domInfo))) {
-      signalPassFailure();
-      return;
-    }
+                                                         domInfo)))
+      return failure();
 
     // Bufferization phase.
-    if (!testAnalysisOnly) {
+    if (!options.testAnalysisOnly) {
       BlockAndValueMapping tensorToBufferMap;
       if (failed(bufferizeFuncOpInternals(funcOp, tensorToBufferMap, aliasInfo,
-                                          *allocationFns,
-                                          bufferizedFunctionTypes))) {
-        signalPassFailure();
-        return;
-      }
+                                          *options.allocationFns,
+                                          bufferizedFunctionTypes)))
+        return failure();
     }
   }
   // Don't drop the attributes if we only want to report the analysis.
-  if (testAnalysisOnly)
-    return;
+  if (options.testAnalysisOnly)
+    return success();
 
   for (FuncOp funcOp : orderedFuncOps) {
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.
     if (failed(bufferizeFuncOpBoundary(funcOp, aliasInfo,
-                                       bufferizedFunctionTypes))) {
-      signalPassFailure();
-      return;
-    }
-    if (!allowReturnMemref &&
+                                       bufferizedFunctionTypes)))
+      return failure();
+
+    if (!options.allowReturnMemref &&
         llvm::any_of(funcOp.getType().getResults(), [](Type t) {
           return t.isa<MemRefType, UnrankedMemRefType>();
         })) {
       funcOp->emitError("memref return type is unsupported");
-      signalPassFailure();
-      return;
+      return failure();
     }
   }
 
   // Perform a post-processing pass of layout modification at function boundary
   // according to the kBufferLayoutAttrName.
   layoutPostProcessing(moduleOp);
 
   // Post-pass cleanup of inplaceable and buffer_layout attributes.
   moduleOp.walk(
       [&](Operation *op) { op->removeAttr(kInPlaceResultsAttrName); });
   moduleOp.walk([&](FuncOp op) {
     for (BlockArgument bbArg : op.getArguments())
       removeBufferizationFuncArguments(bbArg);
   });
 
-  OpPassManager cleanupPipeline("builtin.module");
-  cleanupPipeline.addPass(createCanonicalizerPass());
-  cleanupPipeline.addPass(createCSEPass());
-  cleanupPipeline.addPass(createLoopInvariantCodeMotionPass());
-  (void)runPipeline(cleanupPipeline, moduleOp);
-}
-
-std::unique_ptr<Pass> mlir::createLinalgComprehensiveModuleBufferizePass() {
-  return std::make_unique<LinalgComprehensiveModuleBufferize>();
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
 // BufferizableOpInterface Implementations
 //===----------------------------------------------------------------------===//
 
 // TODO: Move these to a different file and BUILD target, so that they are
 // decoupled from ComprehensiveBufferize.

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

-set(LLVM_OPTIONAL_SOURCES
-  Bufferize.cpp
-  BufferizableOpInterface.cpp
-  CodegenStrategy.cpp
-  ComprehensiveBufferize.cpp
-  Detensorize.cpp
-  Distribution.cpp
-  DropUnitDims.cpp
-  ElementwiseOpFusion.cpp
-  ElementwiseToLinalg.cpp
-  Fusion.cpp
-  FusionOnTensors.cpp
-  Generalization.cpp
-  Hoisting.cpp
-  HoistPadding.cpp
-  InlineScalarOperands.cpp
-  Interchange.cpp
-  Loops.cpp
-  LinalgStrategyPasses.cpp
-  Promotion.cpp
-  Tiling.cpp
-  Transforms.cpp
-  Vectorization.cpp
-)
-
-add_mlir_dialect_library(MLIRBufferizableOpInterface
-  BufferizableOpInterface.cpp
-
-  DEPENDS
-  MLIRBufferizableOpInterfaceIncGen
-
-  LINK_LIBS PUBLIC
-  MLIRIR
-)
-
 add_mlir_dialect_library(MLIRLinalgTransforms
   Bufferize.cpp
   CodegenStrategy.cpp
-  ComprehensiveBufferize.cpp
+  ComprehensiveBufferizePass.cpp
   Detensorize.cpp
   Distribution.cpp
   DropUnitDims.cpp
   ElementwiseOpFusion.cpp
   ElementwiseToLinalg.cpp
   Fusion.cpp
   FusionOnTensors.cpp
   Generalization.cpp
add_mlir_dialect_library(MLIRLinalgTransforms
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRAffineUtils
   MLIRAnalysis
   MLIRArithmetic
   MLIRBufferizableOpInterface
   MLIRComplex
+  MLIRComprehensiveBufferize
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRMemRef
   MLIRLinalgAnalysis
   MLIRLinalg
   MLIRLinalgUtils
   MLIRSCF
   MLIRSCFTransforms

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/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.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::linalg;
+
+namespace {
+struct LinalgComprehensiveModuleBufferize
+    : public LinalgComprehensiveModuleBufferizeBase<
+          LinalgComprehensiveModuleBufferize> {
+  LinalgComprehensiveModuleBufferize() {}
+
+  LinalgComprehensiveModuleBufferize(
+      const LinalgComprehensiveModuleBufferize &p) {}
+
+  void runOnOperation() override;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry
+        .insert<linalg::LinalgDialect, memref::MemRefDialect,
+                tensor::TensorDialect, vector::VectorDialect, scf::SCFDialect,
+                arith::ArithmeticDialect, StandardOpsDialect>();
+    registerBufferizableOpInterfaceExternalModels(registry);
+  }
+};
+} // end namespace
+
+static void applyEnablingTransformations(ModuleOp moduleOp) {
+  RewritePatternSet patterns(moduleOp.getContext());
+  patterns.add<GeneralizePadTensorOpPattern>(moduleOp.getContext());
+  (void)applyPatternsAndFoldGreedily(moduleOp, std::move(patterns));
+}
+
+static Optional<Value>
+allocationFnUsingAlloca(OpBuilder &b, Location loc, MemRefType type,
+                        const SmallVector<Value> &dynShape) {
+  Value allocated = b.create<memref::AllocaOp>(
+      loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
+  return allocated;
+}
+
+void LinalgComprehensiveModuleBufferize::runOnOperation() {
+  BufferizationOptions options;
+  if (useAlloca) {
+    options.allocationFns->allocationFn = allocationFnUsingAlloca;
+    options.allocationFns->deallocationFn = [](OpBuilder &b, Location loc,
+                                               Value v) {};
+  }
+  options.allowReturnMemref = allowReturnMemref;
+  options.analysisFuzzerSeed = analysisFuzzerSeed;
+  options.testAnalysisOnly = testAnalysisOnly;
+
+  ModuleOp moduleOp = getOperation();
+  applyEnablingTransformations(moduleOp);
+
+  if (failed(runComprehensiveBufferize(moduleOp, options))) {
+    signalPassFailure();
+    return;
+  }
+
+  if (options.testAnalysisOnly)
+    return;
+
+  OpPassManager cleanupPipeline("builtin.module");
+  cleanupPipeline.addPass(createCanonicalizerPass());
+  cleanupPipeline.addPass(createCSEPass());
+  cleanupPipeline.addPass(createLoopInvariantCodeMotionPass());
+  (void)runPipeline(cleanupPipeline, moduleOp);
+}
+
+std::unique_ptr<Pass> mlir::createLinalgComprehensiveModuleBufferizePass() {
+  return std::make_unique<LinalgComprehensiveModuleBufferize>();
+}

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/Transforms/BufferizableOpInterface.td",
+        "include/mlir/Dialect/Linalg/ComprehensiveBufferize/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/Transforms/BufferizableOpInterface.h.inc",
+            "include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h.inc",
         ),
         (
             ["-gen-op-interface-defs"],
-            "include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.cpp.inc",
+            "include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp.inc",
         ),
     ],
     tblgen = ":mlir-tblgen",
-    td_file = "include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.td",
+    td_file = "include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td",
     deps = [
         ":BufferizableOpInterfaceTdFiles",
     ],
 )
 
 cc_library(
     name = "BufferizableOpInterface",
     srcs = [
-        "lib/Dialect/Linalg/Transforms/BufferizableOpInterface.cpp",
+        "lib/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp",
     ],
     hdrs = [
-        "include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.h",
+        "include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h",
     ],
     includes = ["include"],
     deps = [
         ":BufferizableOpInterfaceIncGen",
         ":IR",
         ":Support",
         "//llvm:Support",
     ],
gentbl_cc_library(
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/Linalg/Passes.td",
     deps = [":PassBaseTdFiles"],
 )
 
 cc_library(
     name = "LinalgTransforms",
-    srcs = glob(
-        [
-            "lib/Dialect/Linalg/Transforms/*.cpp",
-            "lib/Dialect/Linalg/Transforms/*.h",
-        ],
-        exclude = ["lib/Dialect/Linalg/Transforms/BufferizableOpInterface.cpp"],
-    ) + [
+    srcs = glob([
+    ]) + [
         "lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp",
         "lib/Dialect/Linalg/Utils/Utils.cpp",
     ],
     hdrs = [
         "include/mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h",
         "include/mlir/Dialect/Linalg/Passes.h",
         "include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h",
-        "include/mlir/Dialect/Linalg/Transforms/ComprehensiveBufferize.h",
         "include/mlir/Dialect/Linalg/Transforms/HoistPadding.h",
         "include/mlir/Dialect/Linalg/Transforms/Hoisting.h",
         "include/mlir/Dialect/Linalg/Transforms/Transforms.h",
         "include/mlir/Dialect/Linalg/Utils/Utils.h",
     ],
     includes = ["include"],
     deps = [
         ":Affine",
         ":AffineUtils",
         ":Analysis",
         ":ArithmeticDialect",
-        ":BufferizableOpInterface",
         ":ComplexDialect",
+        ":ComprehensiveBufferize",
         ":DialectUtils",
         ":IR",
         ":InferTypeOpInterface",
         ":LinalgOps",
         ":LinalgPassIncGen",
         ":LinalgStructuredOpsIncGen",
         ":MathDialect",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
         ":SCFTransforms",
         ":StandardOps",
         ":StandardOpsTransforms",
         ":Support",
         ":TensorDialect",
         ":TransformUtils",
         ":VectorOps",
         ":VectorToSCF",
         "//llvm:Support",
     ],
 )
 
 cc_library(
+    name = "ComprehensiveBufferize",
+    srcs = [
+        "lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp",
+    ],
+    hdrs = [
+        "include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h",
+    ],
+    includes = ["include"],
+    deps = [
+        ":ArithmeticDialect",
+        ":BufferizableOpInterface",
+        ":DialectUtils",
+        ":IR",
+        ":InferTypeOpInterface",
+        ":LinalgOps",
+        ":LinalgStructuredOpsIncGen",
+        ":MemRefDialect",
+        ":Pass",
+        ":SCFDialect",
+        ":StandardOps",
+        ":Support",
+        ":TensorDialect",
+        ":TransformUtils",
+        ":VectorOps",
+        "//llvm:Support",
+    ],
+)
+
+cc_library(
     name = "TilingInterface",
     srcs = ["lib/Interfaces/TilingInterface.cpp"],
     hdrs = ["include/mlir/Interfaces/TilingInterface.h"],
     includes = ["include"],
     deps = [
         ":IR",
         ":Support",
         ":TensorDialect",