[mlir][bufferize] Merge tensor-constant-bufferize into arith-bufferize

[springerm]

Summary

The bufferization of arith.constant ops is also switched over to BufferizableOpInterface-based bufferization. The old implementation is deleted. Both implementations utilize GlobalCreator, now renamed to just getGlobalFor.

GlobalCreator no longer maintains a set of all created allocations to avoid duplicate allocations of the same constant. Instead, getGlobalFor scans the module to see if there is already a global allocation with the same constant value.

For compatibility reasons, it is still possible to create a pass that bufferizes only arith.constant. This pass (createConstantBufferizePass) could be deleted once all users were switched over to One-Shot bufferization.

[springerm]

update

[mehdi_amini]

Can you run $ git grep tensor-constant-bufferize mlir and make sure there are no instances left?

mlir/docs/Bufferization.md

```c++
   pm.addPass(createFuncBufferizePass());
 
   // Finalizing bufferization pass.
   pm.addNestedPass<FuncOp>(createFinalizingBufferizePass());
 ```
 
 Looking first at the partial bufferization passes, we see that there are a
 sequence of `FuncOp` passes (which run in parallel on functions). These function
-passes are bracketed by `tensor-constant-bufferize` and `func-bufferize`, which
-are module passes (and thus serialize the parallel compilation process). These
-two passes must be module passes because they make changes to the top-level
-module.
+passes are bracketed by `arith-bufferize` and `func-bufferize`, which are module
+passes (and thus serialize the parallel compilation process). These two passes
+must be module passes because they make changes to the top-level module.
 
 The bulk of the bufferization work is done by the function passes. Most of these
 passes are provided as part of the upstream MLIR distribution and bufferize
 their respective dialects (e.g. `scf-bufferize` bufferizes the `scf` dialect).
 The `tcp-bufferize` pass is an exception -- it is a partial bufferization pass
 used to bufferize the downstream `tcp` dialect, and fits in perfectly with all
 the other passes provided upstream.
 
-   `func-bufferize`
     [test](https://github.com/llvm/llvm-project/blob/bc8acf2ce8ad6e8c9b1d97b2e02d3f4ad26e1d9d/mlir/test/Dialect/Standard/func-bufferize.mlir#L1))
 
     -   Bufferizes `func`, `call`, and `BranchOpInterface` ops.
     -   This is an example of how to bufferize ops that have multi-block
         regions.
     -   This is an example of a pass that is not split along dialect
         subdivisions.
 
--   `tensor-constant-bufferize`
+-   `arith-bufferize`
     ([code](https://github.com/llvm/llvm-project/blob/bc8acf2ce8ad6e8c9b1d97b2e02d3f4ad26e1d9d/mlir/lib/Dialect/StandardOps/Transforms/TensorConstantBufferize.cpp#L1),
     [test](https://github.com/llvm/llvm-project/blob/bc8acf2ce8ad6e8c9b1d97b2e02d3f4ad26e1d9d/mlir/test/Dialect/Standard/tensor-constant-bufferize.mlir#L1))
 
     -   Bufferizes only `arith.constant` ops of `tensor` type.
     -   This is an example of setting up the legality so that only a subset of
         `arith.constant` ops get bufferized.
     -   This is an example of a pass that is not split along dialect
         subdivisions.

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

 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 namespace arith {
 
 /// Create a pass to bufferize Arithmetic ops.
 std::unique_ptr<Pass> createArithmeticBufferizePass();
 
+/// Create a pass to bufferize arith.constant ops.
+std::unique_ptr<Pass> createConstantBufferizePass(uint64_t alignment = 0);
+
 /// Add patterns to expand Arithmetic ops for LLVM lowering.
 void populateArithmeticExpandOpsPatterns(RewritePatternSet &patterns);
 
 /// Create a pass to legalize Arithmetic ops for LLVM lowering.
 std::unique_ptr<Pass> createArithmeticExpandOpsPass();
 
 //===----------------------------------------------------------------------===//
 // Registration

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

 //===-- Passes.td - Arithmetic pass definition file --------*- 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_ARITHMETIC_TRANSFORMS_PASSES
 #define MLIR_DIALECT_ARITHMETIC_TRANSFORMS_PASSES
 
 include "mlir/Pass/PassBase.td"
 
-def ArithmeticBufferize : Pass<"arith-bufferize", "FuncOp"> {
+def ArithmeticBufferize : Pass<"arith-bufferize", "ModuleOp"> {
   let summary = "Bufferize Arithmetic 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 constructor = "mlir::arith::createArithmeticBufferizePass()";
+  let options = [
+    Option<"alignment", "alignment", "unsigned", /*default=*/"0",
+           "Create global memrefs with a specified alignment">,
+  ];
 }
 
 def ArithmeticExpandOps : Pass<"arith-expand", "FuncOp"> {
   let summary = "Legalize Arithmetic ops to be convertible to LLVM.";
   let constructor = "mlir::arith::createArithmeticExpandOpsPass()";
   let dependentDialects = ["StandardOpsDialect"];
 }
 
 #endif // MLIR_DIALECT_ARITHMETIC_TRANSFORMS_PASSES

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

 
 namespace mlir {
 class BlockAndValueMapping;
 class DominanceInfo;
 class FuncOp;
 
 namespace bufferization {
 
-// TODO: from some HW description.
-static constexpr int64_t kBufferAlignments = 128;
-
 class BufferizableOpInterface;
 struct BufferizationOptions;
 class BufferizationState;
 
 /// Options for ComprehensiveBufferize.
 struct BufferizationOptions {
-  using AllocationFn = std::function<FailureOr<Value>(OpBuilder &, Location,
-                                                      MemRefType, ValueRange)>;
+  /// Allocator function: Generate a memref allocation with the given type,
+  /// dynamic extents and alignment.
+  using AllocationFn = std::function<FailureOr<Value>(
+      OpBuilder &, Location, MemRefType, ValueRange, unsigned int)>;
+  /// Deallocator function: Deallocate a buffer that was allocated with
+  /// AllocatorFn.
   using DeallocationFn =
       std::function<LogicalResult(OpBuilder &, Location, Value)>;
+  /// Memcpy function: Generate a memcpy between two buffers.
   using MemCpyFn =
       std::function<LogicalResult(OpBuilder &, Location, Value, Value)>;
 
   BufferizationOptions();
 
   // BufferizationOptions cannot be copied.
   BufferizationOptions(const BufferizationOptions &other) = delete;
 
   /// Return `true` if the op is allowed to be bufferized.
   bool isOpAllowed(Operation *op) const {
-    if (!dialectFilter.hasValue())
+    if (!hasFilter)
       return true;
-    return dialectFilter->contains(op->getDialect()->getNamespace());
+    return dialectFilter.contains(op->getDialect()->getNamespace()) ||
+           operationFilter.contains(op->getName().getStringRef());
   }
 
-  /// Allow-list the given dialects in the dialect filter. Only ops from
-  /// allow-listed dialects will be bufferized. If no dialect is added, ops from
-  /// any dialect will be bufferized.
+  /// Allow the given dialects and activate the filter (`hasFilter`).
   template <typename... DialectTs>
   void addToDialectFilter() {
     // The following expands a call to addToDialectFilterImpl for each dialect
     // in 'DialectTs'. This magic is necessary due to a limitation in the places
     // that a parameter pack can be expanded in c++11.
     // FIXME: In c++17 this can be simplified by using 'fold expressions'.
     (void)std::initializer_list<int>{
         0, (addToDialectFilterImpl<DialectTs>(), 0)...};
   }
 
+  /// Allow the given ops and activate the filter (`hasFilter`).
+  template <typename... OpTys> void addToOperationFilter() {
+    // FIXME: In c++17 this can be simplified by using 'fold expressions'.
+    (void)std::initializer_list<int>{0,
+                                     (addToOperationFilterImpl<OpTys>(), 0)...};
+  }
+
   /// Try to cast the given op to BufferizableOpInterface if the op is allow
   /// listed.
   BufferizableOpInterface dynCastBufferizableOp(Operation *op) const;
 
   /// Try to cast the given value to BufferizableOpInterface if the op is allow
   /// listed.
   BufferizableOpInterface dynCastBufferizableOp(Value value) const;
 
struct BufferizationOptions {
   /// If set to `true`, does not modify the IR apart from adding attributes (for
   /// checking the results of the analysis) and post analysis steps.
   bool testAnalysisOnly = false;
 
   /// If set to `true`, the IR is annotated with details about RaW conflicts.
   /// For debugging only. Should be used together with `testAnalysisOnly`.
   bool printConflicts = false;
 
-  /// Only bufferize ops from dialects that are allowed-listed by the filter.
-  /// All other ops are ignored. This option controls the scope of partial
-  /// bufferization.
+  /// Buffer alignment for new memory allocations.
+  unsigned int bufferAlignment = 128;
+
+  /// If set to `true`, only ops that belong to a filtered dialect
+  /// (`dialectFilter`) and filtered ops (`operationFilter`) are processed. All
+  /// other ops are ignored. If set to `false`, all ops are bufferized (as long
+  /// as they implement BufferizableOpInterface).
   ///
-  /// Note: If no filter is specified, all ops are bufferized (as long as they
-  /// implement BufferizableOpInterface). If a filter is specified,
-  /// `allowUnknownOps` should be enabled. Otherwise, bufferization would fail
-  /// when encountering an op that is forbidden by the filter.
-  Optional<DenseSet<StringRef>> dialectFilter;
+  /// If a filter is specified, `allowUnknownOps` should be enabled. Otherwise,
+  /// bufferization would fail when encountering a non-filtered op.
+  bool hasFilter = false;
+
+  /// A set of allowed dialects.
+  DenseSet<StringRef> dialectFilter;
+
+  /// A set of allowed ops.
+  DenseSet<StringRef> operationFilter;
 
 private:
-  /// Allow-list a dialect in the dialect filter.
+  /// Allow a dialect.
   template <typename DialectT>
   void addToDialectFilterImpl() {
-    if (!dialectFilter.hasValue())
-      dialectFilter.emplace();
-    dialectFilter->insert(DialectT::getDialectNamespace());
+    hasFilter = true;
+    dialectFilter.insert(DialectT::getDialectNamespace());
+  }
+
+  /// Allow an op.
+  template <typename OpTy> void addToOperationFilterImpl() {
+    hasFilter = true;
+    operationFilter.insert(OpTy::getOperationName());
   }
 };
 
 /// Specify fine-grain relationship between buffers to enable more analysis.
 enum class BufferRelation {
   None,
   // TODO: ResultContainsOperand,
   // TODO: OperandContainsResult,

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

protected:
   /// Stores all internally managed allocations.
   BufferPlacementAllocs allocs;
 
   /// The underlying liveness analysis to compute fine grained information
   /// about alloc and dealloc positions.
   Liveness liveness;
 };
 
-// Support class to create global ops for tensor-valued constants in the
-// program. Globals are created lazily at the top of the `moduleOp` with pretty
+// Create a global op for the given tensor-valued constant in the program.
+// Globals are created lazily at the top of the enclosing ModuleOp with pretty
 // names. Duplicates are avoided.
-class GlobalCreator {
-public:
-  GlobalCreator(ModuleOp module, unsigned alignment = 0)
-      : moduleOp(module), alignment(alignment) {}
-  memref::GlobalOp getGlobalFor(arith::ConstantOp constantOp);
+FailureOr<memref::GlobalOp> getGlobalFor(arith::ConstantOp constantOp,
+                                         uint64_t alignment);
 
-private:
-  ModuleOp moduleOp;
-  unsigned alignment;
-  // This could use memref::GlobalOp key but we avoid introducing a new
-  // dependence to the memref dialect for this.
-  DenseMap<Attribute, Operation *> globals;
-};
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERUTILS_H

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

 #ifndef MLIR_DIALECT_STANDARDOPS_TRANSFORMS_PASSES_H
 #define MLIR_DIALECT_STANDARDOPS_TRANSFORMS_PASSES_H
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 namespace bufferization {
 class BufferizeTypeConverter;
-class GlobalCreator;
 } // namespace bufferization
 
 class RewritePatternSet;
 
 void populateStdBufferizePatterns(
     bufferization::BufferizeTypeConverter &typeConverter,
     RewritePatternSet &patterns);
 
 /// Creates an instance of std bufferization pass.
 std::unique_ptr<Pass> createStdBufferizePass();
 
 /// Creates an instance of func bufferization pass.
 std::unique_ptr<Pass> createFuncBufferizePass();
 
-/// Add patterns to bufferize tensor constants into global memrefs to the given
-/// pattern list.
-void populateTensorConstantBufferizePatterns(
-    bufferization::GlobalCreator &globalCreator,
-    bufferization::BufferizeTypeConverter &typeConverter,
-    RewritePatternSet &patterns);
-
-/// Creates an instance of tensor constant bufferization pass.
-std::unique_ptr<Pass> createTensorConstantBufferizePass(unsigned alignment = 0);
-
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_STANDARDOPS_TRANSFORMS_PASSES_H

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

let description = [{
     Finally, this pass fails for unknown terminators, as we cannot decide
     whether they need rewriting.
   }];
   let constructor = "mlir::createFuncBufferizePass()";
   let dependentDialects = ["bufferization::BufferizationDialect",
                            "memref::MemRefDialect"];
 }
 
-def TensorConstantBufferize : Pass<"tensor-constant-bufferize", "ModuleOp"> {
-  let summary = "Bufferize tensor constants.";
-  let description = [{
-    This pass bufferizes tensor constants.
-
-    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 constructor = "mlir::createTensorConstantBufferizePass()";
-  let dependentDialects = ["bufferization::BufferizationDialect",
-                           "memref::MemRefDialect"];
-  let options = [
-    Option<"alignment", "alignment", "unsigned", /*default=*/"0",
-           "Create global memrefs with a specified alignment">,
-  ];
-}
-
 #endif // MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES

mlir/lib/Dialect/Arithmetic/Transforms/BufferizableOpInterfaceImpl.cpp

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     // Only constants inside a module are supported.
     auto moduleOp = constantOp->getParentOfType<ModuleOp>();
     if (!moduleOp)
       return failure();
 
     // Create global memory segment and replace tensor with memref pointing to
     // that memory segment.
-    GlobalCreator globalCreator(moduleOp);
-    auto globalMemref = globalCreator.getGlobalFor(constantOp);
+    FailureOr<memref::GlobalOp> globalOp =
+        getGlobalFor(constantOp, state.getOptions().bufferAlignment);
+    if (failed(globalOp))
+      return failure();
+    memref::GlobalOp globalMemref = globalOp.getValue();
     replaceOpWithNewBufferizedOp<memref::GetGlobalOp>(
         rewriter, op, globalMemref.type(), globalMemref.getName());
 
     return success();
   }
 
   bool isWritable(Operation *op, Value value,
                   const BufferizationState &state) const {

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

 
 using namespace mlir;
 using namespace bufferization;
 
 namespace {
 /// Pass to bufferize Arithmetic ops.
 struct ArithmeticBufferizePass
     : public ArithmeticBufferizeBase<ArithmeticBufferizePass> {
+  ArithmeticBufferizePass(uint64_t alignment = 0, bool constantOpOnly = false)
+      : ArithmeticBufferizeBase<ArithmeticBufferizePass>(),
+        constantOpOnly(constantOpOnly) {
+    this->alignment = alignment;
+  }
+
   void runOnOperation() override {
     std::unique_ptr<BufferizationOptions> options =
         getPartialBufferizationOptions();
+    if (constantOpOnly) {
+      options->addToOperationFilter<arith::ConstantOp>();
+    } else {
       options->addToDialectFilter<arith::ArithmeticDialect>();
+    }
+    options->bufferAlignment = alignment;
 
     if (failed(bufferizeOp(getOperation(), *options)))
       signalPassFailure();
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<bufferization::BufferizationDialect, memref::MemRefDialect,
                     arith::ArithmeticDialect>();
     arith::registerBufferizableOpInterfaceExternalModels(registry);
   }
+
+private:
+  bool constantOpOnly;
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::arith::createArithmeticBufferizePass() {
   return std::make_unique<ArithmeticBufferizePass>();
 }
+
+std::unique_ptr<Pass>
+mlir::arith::createConstantBufferizePass(uint64_t alignment) {
+  return std::make_unique<ArithmeticBufferizePass>(alignment,
+                                                   /*constantOpOnly=*/true);
+}

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

 }
 
 /// Create a memref allocation with the given type and dynamic extents.
 FailureOr<Value>
 bufferization::createAlloc(OpBuilder &b, Location loc, MemRefType type,
                            ValueRange dynShape,
                            const BufferizationOptions &options) {
   if (options.allocationFn)
-    return (*options.allocationFn)(b, loc, type, dynShape);
+    return (*options.allocationFn)(b, loc, type, dynShape,
+                                   options.bufferAlignment);
 
   // Default bufferallocation via AllocOp.
   Value allocated = b.create<memref::AllocOp>(
-      loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
+      loc, type, dynShape, b.getI64IntegerAttr(options.bufferAlignment));
   return allocated;
 }
 
 /// Create a memref allocation with the given type and dynamic extents. May also
 /// deallocate the memref again.
 FailureOr<Value>
 bufferization::createAlloc(OpBuilder &b, Location loc, MemRefType type,
                            ValueRange dynShape, bool deallocMemref,

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

bool BufferPlacementTransformationBase::isLoop(Operation *op) {
 
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // BufferPlacementTransformationBase
 //===----------------------------------------------------------------------===//
 
-memref::GlobalOp GlobalCreator::getGlobalFor(arith::ConstantOp constantOp) {
+FailureOr<memref::GlobalOp>
+bufferization::getGlobalFor(arith::ConstantOp constantOp, uint64_t alignment) {
   auto type = constantOp.getType().cast<RankedTensorType>();
-
-  BufferizeTypeConverter typeConverter;
+  auto moduleOp = constantOp->getParentOfType<ModuleOp>();
+  if (!moduleOp)
+    return failure();
 
   // If we already have a global for this constant value, no need to do
   // anything else.
-  auto it = globals.find(constantOp.getValue());
-  if (it != globals.end())
-    return cast<memref::GlobalOp>(it->second);
+  for (Operation &op : moduleOp.getRegion().getOps()) {
+    auto globalOp = dyn_cast<memref::GlobalOp>(&op);
+    if (!globalOp)
+      continue;
+    if (!globalOp.initial_value().hasValue())
+      continue;
+    uint64_t opAlignment =
+        globalOp.alignment().hasValue() ? globalOp.alignment().getValue() : 0;
+    Attribute initialValue = globalOp.initial_value().getValue();
+    if (opAlignment == alignment && initialValue == constantOp.getValue())
+      return globalOp;
+  }
 
   // Create a builder without an insertion point. We will insert using the
   // symbol table to guarantee unique names.
   OpBuilder globalBuilder(moduleOp.getContext());
   SymbolTable symbolTable(moduleOp);
 
   // Create a pretty name.
   SmallString<64> buf;
   llvm::raw_svector_ostream os(buf);
   interleave(type.getShape(), os, "x");
   os << "x" << type.getElementType();
 
   // Add an optional alignment to the global memref.
   IntegerAttr memrefAlignment =
       alignment > 0 ? IntegerAttr::get(globalBuilder.getI64Type(), alignment)
                     : IntegerAttr();
 
+  BufferizeTypeConverter typeConverter;
   auto global = globalBuilder.create<memref::GlobalOp>(
       constantOp.getLoc(), (Twine("__constant_") + os.str()).str(),
       /*sym_visibility=*/globalBuilder.getStringAttr("private"),
       /*type=*/typeConverter.convertType(type).cast<MemRefType>(),
       /*initial_value=*/constantOp.getValue().cast<ElementsAttr>(),
       /*constant=*/true,
       /*alignment=*/memrefAlignment);
   symbolTable.insert(global);
   // The symbol table inserts at the end of the module, but globals are a bit
   // nicer if they are at the beginning.
   global->moveBefore(&moduleOp.front());
-  globals[constantOp.getValue()] = global;
   return global;
 }

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

 static void applyEnablingTransformations(ModuleOp moduleOp) {
   RewritePatternSet patterns(moduleOp.getContext());
   patterns.add<GeneralizePadOpPattern>(moduleOp.getContext());
   (void)applyPatternsAndFoldGreedily(moduleOp, std::move(patterns));
 }
 
 static FailureOr<Value> allocationFnUsingAlloca(OpBuilder &b, Location loc,
                                                 MemRefType type,
-                                                ValueRange dynShape) {
+                                                ValueRange dynShape,
+                                                unsigned int bufferAlignment) {
   Value allocated = b.create<memref::AllocaOp>(
-      loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
+      loc, type, dynShape, b.getI64IntegerAttr(bufferAlignment));
   return allocated;
 }
 
 /// Create a linalg::GenericOp version of an n-D copy that can further tile,
 /// lower to loops or vectorize, unlike the current implementation of
 /// memref::CopyOp.
 /// Do not depend on linalg::CopyOp that is getting deprecated.
 static LogicalResult createLinalgCopyOp(OpBuilder &b, Location loc, Value from,

mlir/lib/Dialect/SparseTensor/Pipelines/CMakeLists.txt

 add_mlir_dialect_library(MLIRSparseTensorPipelines
   SparseTensorPipelines.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
 
   LINK_LIBS PUBLIC
+  MLIRArithmeticTransforms
   MLIRAffineToStandard
   MLIRBufferizationTransforms
   MLIRLinalgTransforms
   MLIRMathToLLVM
   MLIRMemRefToLLVM
   MLIRPass
   MLIRReconcileUnrealizedCasts
   MLIRSCFToStandard
   MLIRSparseTensor
   MLIRSparseTensorTransforms
   MLIRStandardOpsTransforms
   MLIRTensorTransforms
   MLIRVectorToLLVM
 )

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

 //===- SparseTensorPipelines.cpp - Pipelines for sparse tensor code -------===//
 //
 // 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/SparseTensor/Pipelines/Passes.h"
 
 #include "mlir/Conversion/Passes.h"
+#include "mlir/Dialect/Arithmetic/Transforms/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/Transforms/Passes.h"
 #include "mlir/Pass/PassManager.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 //===----------------------------------------------------------------------===//
 // Pipeline implementation.
 //===----------------------------------------------------------------------===//
 
 void mlir::sparse_tensor::buildSparseCompiler(
     OpPassManager &pm, const SparseCompilerOptions &options) {
   pm.addPass(createSparsificationPass(options.sparsificationOptions()));
   pm.addPass(createSparseTensorConversionPass());
   pm.addPass(createLinalgBufferizePass());
   pm.addPass(createConvertLinalgToLoopsPass());
   pm.addPass(createConvertVectorToSCFPass());
   pm.addPass(createLowerToCFGPass()); // --convert-scf-to-std
   pm.addPass(createFuncBufferizePass());
-  pm.addPass(createTensorConstantBufferizePass());
+  pm.addPass(arith::createConstantBufferizePass());
   pm.addPass(createTensorBufferizePass());
   pm.addPass(createStdBufferizePass());
   pm.addPass(mlir::bufferization::createFinalizingBufferizePass());
   pm.addPass(createLowerAffinePass());
   pm.addPass(createConvertVectorToLLVMPass());
   pm.addPass(createMemRefToLLVMPass());
   pm.addPass(createConvertMathToLLVMPass());
   pm.addPass(createLowerToLLVMPass()); // --convert-std-to-llvm

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

 add_mlir_dialect_library(MLIRStandardOpsTransforms
   Bufferize.cpp
   DecomposeCallGraphTypes.cpp
   FuncBufferize.cpp
   FuncConversions.cpp
-  TensorConstantBufferize.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/StandardOps/Transforms
 
   DEPENDS
   MLIRStandardTransformsIncGen
 
   LINK_LIBS PUBLIC

mlir/lib/Dialect/StandardOps/Transforms/TensorConstantBufferize.cpp

-//===- Bufferize.cpp - Bufferization for std 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
-//
-//===----------------------------------------------------------------------===//
-//
-// This file implements bufferization of tensor-valued arith.constant ops.
-//
-//===----------------------------------------------------------------------===//
-
-#include "PassDetail.h"
-#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
-#include "mlir/Dialect/Bufferization/Transforms/BufferUtils.h"
-#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
-#include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/StandardOps/IR/Ops.h"
-#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
-#include "mlir/IR/BlockAndValueMapping.h"
-#include "mlir/Transforms/DialectConversion.h"
-
-using namespace mlir;
-using namespace mlir::bufferization;
-
-namespace {
-class BufferizeTensorConstantOp
-    : public OpConversionPattern<arith::ConstantOp> {
-public:
-  BufferizeTensorConstantOp(GlobalCreator &globals,
-                            TypeConverter &typeConverter, MLIRContext *context)
-      : OpConversionPattern<arith::ConstantOp>(typeConverter, context,
-                                               /*benefit=*/1),
-        globals(globals) {}
-
-  LogicalResult
-  matchAndRewrite(arith::ConstantOp op, OpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const override {
-    auto type = op.getType().dyn_cast<RankedTensorType>();
-    if (!type)
-      return failure();
-
-    auto globalMemref = globals.getGlobalFor(op);
-    rewriter.replaceOpWithNewOp<memref::GetGlobalOp>(op, globalMemref.type(),
-                                                     globalMemref.getName());
-    return success();
-  }
-  GlobalCreator &globals;
-};
-} // namespace
-
-void mlir::populateTensorConstantBufferizePatterns(
-    GlobalCreator &globalCreator,
-    bufferization::BufferizeTypeConverter &typeConverter,
-    RewritePatternSet &patterns) {
-  patterns.add<BufferizeTensorConstantOp>(globalCreator, typeConverter,
-                                          patterns.getContext());
-}
-
-namespace {
-class TensorConstantBufferizePass
-    : public TensorConstantBufferizeBase<TensorConstantBufferizePass> {
-public:
-  explicit TensorConstantBufferizePass(unsigned alignment) {
-    if (alignment)
-      this->alignment = alignment;
-  }
-
-  void runOnOperation() override {
-    auto module = getOperation();
-    GlobalCreator globals(module, alignment);
-
-    auto *context = &getContext();
-    bufferization::BufferizeTypeConverter typeConverter;
-    RewritePatternSet patterns(context);
-    ConversionTarget target(*context);
-
-    target.addLegalDialect<memref::MemRefDialect>();
-    populateTensorConstantBufferizePatterns(globals, typeConverter, patterns);
-    target.addDynamicallyLegalOp<arith::ConstantOp>([&](arith::ConstantOp op) {
-      return typeConverter.isLegal(op.getType());
-    });
-    if (failed(applyPartialConversion(module, target, std::move(patterns))))
-      signalPassFailure();
-  }
-};
-} // namespace
-
-std::unique_ptr<Pass>
-mlir::createTensorConstantBufferizePass(unsigned alignment) {
-  return std::make_unique<TensorConstantBufferizePass>(alignment);
-}

mlir/test/Dialect/Arithmetic/bufferize.mlir

-// RUN: mlir-opt %s -arith-bufferize | FileCheck %s
+// RUN: mlir-opt %s -arith-bufferize -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -arith-bufferize=alignment=64 -split-input-file | FileCheck --check-prefix=ALIGNED %s
 
 // CHECK-LABEL:   func @index_cast(
 // CHECK-SAME:  %[[TENSOR:.*]]: tensor<i32>, %[[SCALAR:.*]]: i32
 func @index_cast(%tensor: tensor<i32>, %scalar: i32) -> (tensor<index>, index) {
   %index_tensor = arith.index_cast %tensor : tensor<i32> to tensor<index>
   %index_scalar = arith.index_cast %scalar : i32 to index
   return %index_tensor, %index_scalar : tensor<index>, index
 }
 // CHECK:  %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<i32>
 // CHECK-NEXT: %[[INDEX_MEMREF:.*]] = arith.index_cast %[[MEMREF]]
 // CHECK-SAME:   memref<i32> to memref<index>
 // CHECK-NEXT: %[[INDEX_TENSOR:.*]] = bufferization.to_tensor %[[INDEX_MEMREF]]
 // CHECK: return %[[INDEX_TENSOR]]
+
+// -----
+
+// CHECK-LABEL: module {
+
+// We check the debug name too since we put some effort into making that readable.
+// The name isn't load-bearing though.
+
+// CHECK: memref.global "private" constant @__constant_3x4xf32 : memref<3x4xf32> = dense<7.000000e+00>
+// CHECK-NOT: alignment
+
+// ALIGNED: memref.global "private" constant @__constant_3x4xf32 : memref<3x4xf32> = dense<7.000000e+00>
+// ALIGNED-SAME: {alignment = 64 : i64}
+
+// CHECK: @basic
+func @basic() -> tensor<3x4xf32> {
+  // CHECK: %[[MEMREF:.*]] = memref.get_global @__constant_3x4xf32 : memref<3x4xf32>
+  // CHECK: %[[TENSOR:.*]] = bufferization.to_tensor %[[MEMREF]]
+  %0 = arith.constant dense<7.0> : tensor<3x4xf32>
+  // CHECK: return %[[TENSOR]]
+  return %0 : tensor<3x4xf32>
+}
+
+// CHECK: }
+
+// -----
+
+// CHECK-LABEL: module {
+
+// Only one global is created.
+// CHECK: memref.global
+// CHECK-NOT: memref.global
+func @duplicate_constants() -> (tensor<3x4xf32>, tensor<3x4xf32>) {
+  %0 = arith.constant dense<7.0> : tensor<3x4xf32>
+  %1 = arith.constant dense<7.0> : tensor<3x4xf32>
+  return %0, %1 : tensor<3x4xf32>, tensor<3x4xf32>
+}
+
+// CHECK: }
+
+// -----
+
+// CHECK-LABEL: module {
+
+// Two globals are created.
+// CHECK: memref.global
+// CHECK: memref.global
+// CHECK-NOT: memref.global
+func @multiple_constants() -> (tensor<3x4xf32>, tensor<3x4xf32>) {
+  %0 = arith.constant dense<7.0> : tensor<3x4xf32>
+  %1 = arith.constant dense<8.0> : tensor<3x4xf32>
+  return %0, %1 : tensor<3x4xf32>, tensor<3x4xf32>
+}
+
+// CHECK: }
+
+// -----
+
+// CHECK-LABEL: module {
+// We don't convert non-tensor globals.
+// CHECK-NOT: memref.global
+func @non_tensor() {
+    %0 = arith.constant 7 : i32
+    return
+}
+
+// CHECK: }

mlir/test/Dialect/SparseTensor/sparse_lower.mlir

 // RUN: mlir-opt %s -sparsification | FileCheck %s --check-prefix=CHECK-HIR
 //
 // RUN: mlir-opt %s -sparsification --sparse-tensor-conversion | \
 // RUN: FileCheck %s --check-prefix=CHECK-MIR
 //
 // RUN: mlir-opt %s -sparsification --sparse-tensor-conversion \
-// RUN: --func-bufferize --tensor-constant-bufferize           \
+// RUN: --func-bufferize --arith-bufferize           \
 // RUN: --tensor-bufferize --finalizing-bufferize |            \
 // RUN: FileCheck %s --check-prefix=CHECK-LIR
 
 #CSR = #sparse_tensor.encoding<{dimLevelType = [ "dense", "compressed" ]}>
 
 #trait_matvec = {
   indexing_maps = [
     affine_map<(i,j) -> (i,j)>,  // A

mlir/test/Dialect/SparseTensor/sparse_lower_col.mlir

 // RUN: mlir-opt %s -sparsification | FileCheck %s --check-prefix=CHECK-HIR
 //
 // RUN: mlir-opt %s -sparsification --sparse-tensor-conversion | \
 // RUN: FileCheck %s --check-prefix=CHECK-MIR
 //
 // RUN: mlir-opt %s -sparsification --sparse-tensor-conversion \
-// RUN: --func-bufferize --tensor-constant-bufferize           \
+// RUN: --func-bufferize --arith-bufferize           \
 // RUN: --tensor-bufferize --finalizing-bufferize |            \
 // RUN: FileCheck %s --check-prefix=CHECK-LIR
 
 #CSC = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ],
   dimOrdering = affine_map<(i,j) -> (j,i)>
 }>
 

mlir/test/Dialect/SparseTensor/sparse_lower_inplace.mlir

 // RUN: mlir-opt %s -sparsification | FileCheck %s --check-prefix=CHECK-HIR
 //
 // RUN: mlir-opt %s -sparsification --sparse-tensor-conversion | \
 // RUN: FileCheck %s --check-prefix=CHECK-MIR
 //
 // RUN: mlir-opt %s -sparsification --sparse-tensor-conversion \
-// RUN: --func-bufferize --tensor-constant-bufferize           \
+// RUN: --func-bufferize --arith-bufferize           \
 // RUN: --tensor-bufferize --finalizing-bufferize |            \
 // RUN: FileCheck %s --check-prefix=CHECK-LIR
 
 #CSR = #sparse_tensor.encoding<{dimLevelType = [ "dense", "compressed" ]}>
 
 #trait_matvec = {
   indexing_maps = [
     affine_map<(i,j) -> (i,j)>,  // A

mlir/test/Dialect/Standard/tensor-constant-bufferize.mlir

-// RUN: mlir-opt %s -tensor-constant-bufferize -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -tensor-constant-bufferize=alignment=64 -split-input-file | FileCheck --check-prefix=ALIGNED %s
-
-// CHECK-LABEL: module {
-
-// We check the debug name too since we put some effort into making that readable.
-// The name isn't load-bearing though.
-
-// CHECK: memref.global "private" constant @__constant_3x4xf32 : memref<3x4xf32> = dense<7.000000e+00>
-// CHECK-NOT: alignment
-
-// ALIGNED: memref.global "private" constant @__constant_3x4xf32 : memref<3x4xf32> = dense<7.000000e+00>
-// ALIGNED-SAME: {alignment = 64 : i64}
-
-// CHECK: @basic
-func @basic() -> tensor<3x4xf32> {
-  // CHECK: %[[MEMREF:.*]] = memref.get_global @__constant_3x4xf32 : memref<3x4xf32>
-  // CHECK: %[[TENSOR:.*]] = bufferization.to_tensor %[[MEMREF]]
-  %0 = arith.constant dense<7.0> : tensor<3x4xf32>
-  // CHECK: return %[[TENSOR]]
-  return %0 : tensor<3x4xf32>
-}
-
-// CHECK: }
-
-// -----
-
-// CHECK-LABEL: module {
-
-// Only one global is created.
-// CHECK: memref.global
-// CHECK-NOT: memref.global
-func @duplicate_constants() -> (tensor<3x4xf32>, tensor<3x4xf32>) {
-  %0 = arith.constant dense<7.0> : tensor<3x4xf32>
-  %1 = arith.constant dense<7.0> : tensor<3x4xf32>
-  return %0, %1 : tensor<3x4xf32>, tensor<3x4xf32>
-}
-
-// CHECK: }
-
-// -----
-
-// CHECK-LABEL: module {
-
-// Two globals are created.
-// CHECK: memref.global
-// CHECK: memref.global
-// CHECK-NOT: memref.global
-func @multiple_constants() -> (tensor<3x4xf32>, tensor<3x4xf32>) {
-  %0 = arith.constant dense<7.0> : tensor<3x4xf32>
-  %1 = arith.constant dense<8.0> : tensor<3x4xf32>
-  return %0, %1 : tensor<3x4xf32>, tensor<3x4xf32>
-}
-
-// CHECK: }
-
-// -----
-
-// CHECK-LABEL: module {
-// We don't convert non-tensor globals.
-// CHECK-NOT: memref.global
-func @non_tensor() {
-    %0 = arith.constant 7 : i32
-    return
-}
-
-// CHECK: }

mlir/test/Integration/Dialect/Linalg/CPU/test-collapse-tensor.mlir

 // RUN: mlir-opt %s -linalg-bufferize -std-bufferize \
-// RUN: -tensor-constant-bufferize -tensor-bufferize -func-bufferize \
+// RUN: -arith-bufferize -tensor-bufferize -func-bufferize \
 // RUN: -finalizing-bufferize -buffer-deallocation -convert-linalg-to-llvm \
 // RUN: -convert-memref-to-llvm -convert-std-to-llvm -reconcile-unrealized-casts | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 
 func @main() {

mlir/test/Integration/Dialect/Linalg/CPU/test-elementwise.mlir

 // RUN: mlir-opt %s -convert-elementwise-to-linalg -std-bufferize \
-// RUN: -tensor-constant-bufferize -linalg-bufferize -tensor-bufferize \
+// RUN: -arith-bufferize -linalg-bufferize -tensor-bufferize \
 // RUN: -func-bufferize -buffer-deallocation -convert-linalg-to-loops \
 // RUN: -convert-linalg-to-llvm --convert-memref-to-llvm -convert-std-to-llvm \
 // RUN: -reconcile-unrealized-casts | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @main() {

mlir/test/Integration/Dialect/Linalg/CPU/test-expand-tensor.mlir

 // RUN: mlir-opt %s -linalg-bufferize -std-bufferize \
-// RUN: -tensor-constant-bufferize -tensor-bufferize -func-bufferize \
+// RUN: -arith-bufferize -tensor-bufferize -func-bufferize \
 // RUN: -finalizing-bufferize -buffer-deallocation -convert-linalg-to-llvm \
 // RUN: -convert-memref-to-llvm -convert-std-to-llvm -reconcile-unrealized-casts | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 
 func @main() {

mlir/test/Integration/Dialect/Linalg/CPU/test-padtensor.mlir

 // RUN: mlir-opt %s -linalg-bufferize -std-bufferize \
-// RUN: -tensor-constant-bufferize -tensor-bufferize -func-bufferize \
+// RUN: -arith-bufferize -tensor-bufferize -func-bufferize \
 // RUN: -finalizing-bufferize -buffer-deallocation \
 // RUN: -convert-linalg-to-loops -convert-scf-to-std -convert-linalg-to-llvm -convert-memref-to-llvm -convert-std-to-llvm -reconcile-unrealized-casts | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 
 func @main() {

mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert-multiple-uses.mlir

 // RUN: mlir-opt %s -linalg-bufferize -std-bufferize \
-// RUN: -tensor-constant-bufferize -tensor-bufferize -func-bufferize \
+// RUN: -arith-bufferize -tensor-bufferize -func-bufferize \
 // RUN: -finalizing-bufferize -buffer-deallocation \
 // RUN: -convert-linalg-to-loops -convert-scf-to-std -convert-linalg-to-llvm --convert-memref-to-llvm -convert-std-to-llvm -reconcile-unrealized-casts | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @main() {
   %const = arith.constant dense<10.0> : tensor<2xf32>

mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert.mlir

 // RUN: mlir-opt %s -linalg-bufferize -std-bufferize \
-// RUN: -tensor-constant-bufferize -tensor-bufferize -func-bufferize \
+// RUN: -arith-bufferize -tensor-bufferize -func-bufferize \
 // RUN: -finalizing-bufferize -buffer-deallocation \
 // RUN: -convert-linalg-to-loops -convert-scf-to-std -convert-linalg-to-llvm --convert-memref-to-llvm -convert-std-to-llvm -reconcile-unrealized-casts | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @main() {
   %const = arith.constant dense<10.0> : tensor<2xf32>

mlir/test/Integration/Dialect/Linalg/CPU/test-tensor-e2e.mlir

-// RUN: mlir-opt %s -tensor-constant-bufferize -std-bufferize -linalg-bufferize \
+// RUN: mlir-opt %s -arith-bufferize -std-bufferize -linalg-bufferize \
 // RUN: -tensor-bufferize -func-bufferize -finalizing-bufferize -buffer-deallocation -convert-linalg-to-loops \
 // RUN: -convert-linalg-to-llvm --convert-memref-to-llvm -convert-std-to-llvm -reconcile-unrealized-casts | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @foo() -> tensor<4xf32> {
   %0 = arith.constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>

mlir/test/Integration/Dialect/Linalg/CPU/test-tensor-matmul.mlir

 // UNSUPPORTED: asan
-// RUN: mlir-opt %s -linalg-bufferize -std-bufferize -tensor-constant-bufferize \
+// RUN: mlir-opt %s -linalg-bufferize -std-bufferize -arith-bufferize \
 // RUN: -tensor-bufferize -func-bufferize -finalizing-bufferize -buffer-deallocation -convert-linalg-to-loops -convert-scf-to-std \
 // RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-std --convert-memref-to-llvm -convert-std-to-llvm -reconcile-unrealized-casts | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s  -linalg-tile="tile-sizes=1,2,3" -linalg-bufferize \
-// RUN: -scf-bufferize -std-bufferize -tensor-constant-bufferize -tensor-bufferize \
+// RUN: -scf-bufferize -std-bufferize -arith-bufferize -tensor-bufferize \
 // RUN: -func-bufferize \
 // RUN: -finalizing-bufferize -convert-linalg-to-loops -convert-scf-to-std -convert-scf-to-std \
 // RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-std --convert-memref-to-llvm -convert-std-to-llvm -reconcile-unrealized-casts | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @main() {

mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py

class SparseCompiler:
 
   def __init__(self, options: str):
     pipeline = (
         f'sparsification{{{options}}},'
         f'sparse-tensor-conversion,'
         f'builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),'
         f'convert-scf-to-std,'
         f'func-bufferize,'
-        f'tensor-constant-bufferize,'
+        f'arith-bufferize,'
         f'builtin.func(tensor-bufferize,std-bufferize,finalizing-bufferize),'
         f'convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},'
         f'lower-affine,'
         f'convert-memref-to-llvm,'
         f'convert-std-to-llvm,'
         f'reconcile-unrealized-casts')
     self.pipeline = pipeline
 

mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py

class SparseCompiler:
   def __init__(self, options: str):
     pipeline = (
         f'builtin.func(linalg-generalize-named-ops,linalg-fuse-elementwise-ops),'
         f'sparsification{{{options}}},'
         f'sparse-tensor-conversion,'
         f'builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),'
         f'convert-scf-to-std,'
         f'func-bufferize,'
-        f'tensor-constant-bufferize,'
+        f'arith-bufferize,'
         f'builtin.func(tensor-bufferize,std-bufferize,finalizing-bufferize),'
         f'convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},'
         f'lower-affine,'
         f'convert-memref-to-llvm,'
         f'convert-std-to-llvm,'
         f'reconcile-unrealized-casts')
     self.pipeline = pipeline
 

mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py

class _SparseCompiler:
 
   def __init__(self):
     self.pipeline = (
         f'sparsification,'
         f'sparse-tensor-conversion,'
         f'builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),'
         f'convert-scf-to-std,'
         f'func-bufferize,'
-        f'tensor-constant-bufferize,'
+        f'arith-bufferize,'
         f'builtin.func(tensor-bufferize,std-bufferize,finalizing-bufferize),'
         f'convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},'
         f'lower-affine,'
         f'convert-memref-to-llvm,'
         f'convert-std-to-llvm,'
         f'reconcile-unrealized-casts')
 
   def __call__(self, module: ir.Module):

mlir/test/Integration/Dialect/SparseTensor/python/test_output.py

class SparseCompiler:
   def __init__(self):
     pipeline = (
         f'builtin.func(linalg-generalize-named-ops,linalg-fuse-elementwise-ops),'
         f'sparsification,'
         f'sparse-tensor-conversion,'
         f'builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),'
         f'convert-scf-to-std,'
         f'func-bufferize,'
-        f'tensor-constant-bufferize,'
+        f'arith-bufferize,'
         f'builtin.func(tensor-bufferize,std-bufferize,finalizing-bufferize),'
         f'convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},'
         f'lower-affine,'
         f'convert-memref-to-llvm,'
         f'convert-std-to-llvm,'
         f'reconcile-unrealized-casts')
     self.pipeline = pipeline
 

mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py

def __init__(self, sparsification_options: str, support_lib: str):
     self._support_lib = support_lib
     self._pipeline = (
         f'builtin.func(linalg-generalize-named-ops,linalg-fuse-elementwise-ops),'
         f'sparsification{{{sparsification_options}}},'
         f'sparse-tensor-conversion,'
         f'builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),'
         f'convert-scf-to-std,'
         f'func-bufferize,'
-        f'tensor-constant-bufferize,'
+        f'arith-bufferize,'
         f'builtin.func(tensor-bufferize,std-bufferize,finalizing-bufferize),'
         f'convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},'
         f'lower-affine,'
         f'convert-memref-to-llvm,'
         f'convert-std-to-llvm,'
         f'reconcile-unrealized-casts')
     # Must be in the scope of a `with ir.Context():`
     self._passmanager = PassManager.parse(self._pipeline)

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

dtype_to_irtype = {
       Type.FLOAT32: ir.F32Type.get(),
       Type.FLOAT64: ir.F64Type.get()
   }
   return dtype_to_irtype[dtype.kind]
 
 
 def _compile_mlir(module: ir.Module) -> ir.Module:
   """Compiles an MLIR module and returns the compiled module."""
-  # TODO: Replace this with a pipeline implemented for 
+  # TODO: Replace this with a pipeline implemented for
   #   https://github.com/llvm/llvm-project/issues/51751.
   pipeline = (
       f"sparsification,"
       f"sparse-tensor-conversion,"
       f"builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),"
       f"convert-scf-to-std,"
       f"func-bufferize,"
-      f"tensor-constant-bufferize,"
+      f"arith-bufferize,"
       f"builtin.func(tensor-bufferize,std-bufferize,finalizing-bufferize),"
       f"convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},"
       f"lower-affine,"
       f"convert-memref-to-llvm,"
       f"convert-std-to-llvm,"
       f"reconcile-unrealized-casts")
   PassManager.parse(pipeline).run(module)
   return module

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

void TestComprehensiveFunctionBufferize::runOnOperation() {
   options->allowReturnMemref = allowReturnMemref;
   options->allowUnknownOps = allowUnknownOps;
   options->testAnalysisOnly = testAnalysisOnly;
   options->analysisFuzzerSeed = analysisFuzzerSeed;
   options->fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
   options->createDeallocs = createDeallocs;
 
   if (dialectFilter.hasValue()) {
-    options->dialectFilter.emplace();
+    options->hasFilter = true;
     for (const std::string &dialectNamespace : dialectFilter)
-      options->dialectFilter->insert(dialectNamespace);
+      options->dialectFilter.insert(dialectNamespace);
   }
 
   Operation *op = getOperation();
   if (failed(runOneShotBufferize(op, std::move(options))))
     return;
 
   if (testAnalysisOnly)
     return;