[mlir] Expose a function to populate tensor constant bufferization patterns

[bkramer]

Summary

This makes it easier to use it from other bufferization passes.

[herhut]

Thanks!

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

35

Please add a /// comment.

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

 #ifndef MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES_H_
 #define MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Bufferize.h"
 
 namespace mlir {
 
+class GlobalCreator;
 class RewritePatternSet;
 using OwningRewritePatternList = RewritePatternSet;
 
 void populateStdBufferizePatterns(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(
+    GlobalCreator &globalCreator, BufferizeTypeConverter &typeConverter,
+    RewritePatternSet &patterns);
+
 /// Creates an instance of tensor constant bufferization pass.
 std::unique_ptr<Pass> createTensorConstantBufferizePass();
 
 /// Creates an instance of the StdExpand pass that legalizes Std
 /// dialect ops to be convertible to LLVM. For example,
 /// `std.ceildivi_signed` gets transformed to a number of std operations,
 /// which can be lowered to LLVM; `memref.reshape` gets converted to
 /// `memref_reinterpret_cast`.

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

matchAndRewrite(ConstantOp op, ArrayRef<Value> operands,
     rewriter.replaceOpWithNewOp<memref::GetGlobalOp>(op, globalMemref.type(),
                                                      globalMemref.getName());
     return success();
   }
   GlobalCreator &globals;
 };
 } // namespace
 
+void mlir::populateTensorConstantBufferizePatterns(
+    GlobalCreator &globalCreator, BufferizeTypeConverter &typeConverter,
+    RewritePatternSet &patterns) {
+  patterns.add<BufferizeTensorConstantOp>(globalCreator, typeConverter,
+                                          patterns.getContext());
+}
+
 namespace {
 struct TensorConstantBufferizePass
     : public TensorConstantBufferizeBase<TensorConstantBufferizePass> {
   void runOnOperation() override {
     auto module = getOperation();
     GlobalCreator globals(module);
 
     auto *context = &getContext();
     BufferizeTypeConverter typeConverter;
     RewritePatternSet patterns(context);
     ConversionTarget target(*context);
 
     target.addLegalDialect<memref::MemRefDialect>();
-    patterns.add<BufferizeTensorConstantOp>(globals, typeConverter, context);
+    populateTensorConstantBufferizePatterns(globals, typeConverter, patterns);
     target.addDynamicallyLegalOp<ConstantOp>(
         [&](ConstantOp op) { return typeConverter.isLegal(op.getType()); });
     if (failed(applyPartialConversion(module, target, std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createTensorConstantBufferizePass() {
   return std::make_unique<TensorConstantBufferizePass>();
 }