Revert "[mlir] Remove dialect specific bufferization passes" (#93528)

Reverts #93488

Buildbot failure:
https://lab.llvm.org/buildbot/#/builders/220/builds/39911

Author: Groverkss

mlir/include/mlir/Dialect/Arith/Transforms/Passes.h

@@ -24,6 +24,9 @@ namespace arith {
 class WideIntEmulationConverter;
 class NarrowTypeEmulationConverter;
 
+/// Create a pass to bufferize arith.constant ops.
+std::unique_ptr<Pass> createConstantBufferizePass(uint64_t alignment = 0);
+
 /// Adds patterns to emulate wide Arith and Function ops over integer
 /// types into supported ones. This is done by splitting original power-of-two
 /// i2N integer types into two iN halves.

mlir/include/mlir/Dialect/Arith/Transforms/Passes.td

@@ -11,6 +11,22 @@
 
 include "mlir/Pass/PassBase.td"
 
+def ArithBufferizePass : Pass<"arith-bufferize", "ModuleOp"> {
+  let summary = "Bufferize Arith dialect ops.";
+  let description = [{
+    This pass bufferizes arith dialect ops.
+
+    This pass needs to be a module pass because it inserts memref.global
+    ops into the module, which cannot be done safely from a function pass due to
+    multi-threading. Most other bufferization passes can run in parallel at
+    function granularity.
+  }];
+  let options = [
+    Option<"alignment", "alignment", "unsigned", /*default=*/"0",
+           "Create global memrefs with a specified alignment">,
+  ];
+}
+
 def ArithExpandOpsPass : Pass<"arith-expand"> {
   let summary = "Legalize Arith ops to be convertible to LLVM.";
   let dependentDialects = ["vector::VectorDialect"];

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

@@ -221,6 +221,9 @@ createPromoteBuffersToStackPass(std::function<bool(Value)> isSmallAlloc);
 /// insert_slice ops.
 std::unique_ptr<Pass> createEmptyTensorEliminationPass();
 
+/// Create a pass that bufferizes ops from the bufferization dialect.
+std::unique_ptr<Pass> createBufferizationBufferizePass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

@@ -350,6 +350,11 @@ def FinalizingBufferize : Pass<"finalizing-bufferize", "func::FuncOp"> {
   let constructor = "mlir::bufferization::createFinalizingBufferizePass()";
 }
 
+def BufferizationBufferize : Pass<"bufferization-bufferize", "func::FuncOp"> {
+  let summary = "Bufferize the `bufferization` dialect";
+  let constructor = "mlir::bufferization::createBufferizationBufferizePass()";
+}
+
 def DropEquivalentBufferResults : Pass<"drop-equivalent-buffer-results", "ModuleOp">  {
   let summary = "Remove MemRef return values that are equivalent to a bbArg";
   let description = [{

mlir/include/mlir/Dialect/Linalg/Passes.h

@@ -22,6 +22,10 @@ namespace func {
 class FuncOp;
 } // namespace func
 
+namespace bufferization {
+struct OneShotBufferizationOptions;
+} // namespace bufferization
+
 #define GEN_PASS_DECL
 #include "mlir/Dialect/Linalg/Passes.h.inc" // IWYU pragma: keep
 

mlir/include/mlir/Dialect/Linalg/Passes.td

@@ -89,6 +89,16 @@ def LinalgInlineScalarOperandsPass : Pass<"linalg-inline-scalar-operands"> {
   ];
 }
 
+def LinalgBufferizePass : Pass<"linalg-bufferize"> {
+  let summary = "Bufferize the linalg dialect";
+  let dependentDialects = [
+    "affine::AffineDialect",
+    "bufferization::BufferizationDialect",
+    "linalg::LinalgDialect",
+    "memref::MemRefDialect",
+  ];
+}
+
 def LinalgGeneralizeNamedOpsPass : Pass<"linalg-generalize-named-ops"> {
   let summary = "Convert named ops into generic ops";
   let dependentDialects = ["linalg::LinalgDialect"];

mlir/include/mlir/Dialect/Shape/Transforms/Passes.h

@@ -47,6 +47,13 @@ void populateShapeRewritePatterns(RewritePatternSet &patterns);
 void populateRemoveShapeConstraintsPatterns(RewritePatternSet &patterns);
 std::unique_ptr<OperationPass<func::FuncOp>> createRemoveShapeConstraintsPass();
 
+// Bufferizes shape dialect ops.
+//
+// Note that most shape dialect ops must be converted to std before
+// bufferization happens, as they are intended to be bufferized at the std
+// level.
+std::unique_ptr<OperationPass<func::FuncOp>> createShapeBufferizePass();
+
 /// Outline the shape computation part by adding shape.func and populate
 /// conrresponding mapping infomation into ShapeMappingAnalysis.
 std::unique_ptr<OperationPass<ModuleOp>> createOutlineShapeComputationPass();

mlir/include/mlir/Dialect/Shape/Transforms/Passes.td

@@ -103,4 +103,11 @@ def ShapeToShapeLowering : Pass<"shape-to-shape-lowering", "func::FuncOp"> {
   let constructor = "mlir::createShapeToShapeLowering()";
 }
 
+// TODO: Generalize this to allow any type conversions desired.
+def ShapeBufferize : Pass<"shape-bufferize", "func::FuncOp"> {
+  let summary = "Bufferize the shape dialect.";
+  let constructor = "mlir::createShapeBufferizePass()";
+  let dependentDialects = ["bufferization::BufferizationDialect",
+                           "memref::MemRefDialect"];
+}
 #endif // MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES

mlir/include/mlir/Dialect/Tensor/Transforms/Passes.h

@@ -21,6 +21,9 @@ namespace tensor {
 /// Creates an instance of the `tensor` subset folding pass.
 std::unique_ptr<Pass> createFoldTensorSubsetOpsPass();
 
+/// Creates an instance of the `tensor` dialect bufferization pass.
+std::unique_ptr<Pass> createTensorBufferizePass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td

@@ -27,4 +27,9 @@ def FoldTensorSubsetOps : Pass<"fold-tensor-subset-ops"> {
   ];
 }
 
+def TensorBufferize : Pass<"tensor-bufferize", "func::FuncOp"> {
+  let summary = "Bufferize the `tensor` dialect";
+  let constructor = "mlir::tensor::createTensorBufferizePass()";
+}
+
 #endif // MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES