[mlir][bufferization] Enable BufferizationToMemRef to run on functions

[maerhart]

Summary

Modify the convert-bufferization-to-memref pass to be applicable to
builtin.module and FunctionOpInterface. In many cases the lowering using
the library function is not needed because enough static information is
available to simplify the dealloc operations enough to use the more
light-weight lowerings. In those situations, it is better to allow the pass to
run on functions for additional parallelization and for easier use with the
transform dialect (e.g., when the transform ops are in the same
builtin.module as the function to be transformed).

mlir/include/mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.h

 #include <memory>
 
 namespace mlir {
 class ModuleOp;
 
 #define GEN_PASS_DECL_CONVERTBUFFERIZATIONTOMEMREF
 #include "mlir/Conversion/Passes.h.inc"
 
-std::unique_ptr<OperationPass<ModuleOp>> createBufferizationToMemRefPass();
+std::unique_ptr<Pass> createBufferizationToMemRefPass();
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_BUFFERIZATIONTOMEMREF_BUFFERIZATIONTOMEMREF_H

mlir/include/mlir/Conversion/Passes.td

Option<"useOpaquePointers", "use-opaque-pointers", "bool",
            "instead of typed pointers">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // BufferizationToMemRef
 //===----------------------------------------------------------------------===//
 
-def ConvertBufferizationToMemRef : Pass<"convert-bufferization-to-memref",
-                                        "mlir::ModuleOp"> {
+def ConvertBufferizationToMemRef : Pass<"convert-bufferization-to-memref"> {
   let summary = "Convert operations from the Bufferization dialect to the "
                 "MemRef dialect";
   let description = [{
 
     This pass converts bufferization operations into memref operations.
 
     In the current state, this pass only transforms a `bufferization.clone`
     operation into `memref.alloc` and `memref.copy` operations. This conversion

mlir/lib/Conversion/BufferizationToMemRef/BufferizationToMemRef.cpp

matchAndRewrite(bufferization::DeallocOp op, OpAdaptor adaptor,
     }
 
     if (adaptor.getMemrefs().size() == 1 && adaptor.getRetained().empty())
       return rewriteOneMemrefNoRetainCase(op, adaptor, rewriter);
 
     if (adaptor.getMemrefs().size() == 1)
       return rewriteOneMemrefMultipleRetainCase(op, adaptor, rewriter);
 
+    if (!deallocHelperFunc)
+      return op->emitError(
+          "library function required for generic lowering, but cannot be "
+          "automatically inserted when operating on functions");
+
     return rewriteGeneralCase(op, adaptor, rewriter);
   }
 
   /// Build a helper function per compilation unit that can be called at
   /// bufferization dealloc sites to determine aliasing and ownership.
   ///
   /// The generated function takes two memrefs of indices and three memrefs of
   /// booleans as arguments:

 } // namespace
 
 namespace {
 struct BufferizationToMemRefPass
     : public impl::ConvertBufferizationToMemRefBase<BufferizationToMemRefPass> {
   BufferizationToMemRefPass() = default;
 
   void runOnOperation() override {
-    ModuleOp module = cast<ModuleOp>(getOperation());
+    if (!isa<ModuleOp, FunctionOpInterface>(getOperation())) {
+      emitError(getOperation()->getLoc(),
+                "root operation must be a builtin.module or a function");
+      signalPassFailure();
+      return;
+    }
+
+    func::FuncOp helperFuncOp;
+    if (auto module = dyn_cast<ModuleOp>(getOperation())) {
       OpBuilder builder =
           OpBuilder::atBlockBegin(&module.getBodyRegion().front());
       SymbolTable symbolTable(module);
 
       // Build dealloc helper function if there are deallocs.
-    func::FuncOp helperFuncOp;
       getOperation()->walk([&](bufferization::DeallocOp deallocOp) {
         if (deallocOp.getMemrefs().size() > 1) {
           helperFuncOp = DeallocOpConversion::buildDeallocationHelperFunction(
               builder, getOperation()->getLoc(), symbolTable);
           return WalkResult::interrupt();
         }
         return WalkResult::advance();
       });
+    }
 
     RewritePatternSet patterns(&getContext());
     patterns.add<CloneOpConversion>(patterns.getContext());
     patterns.add<DeallocOpConversion>(patterns.getContext(), helperFuncOp);
 
     ConversionTarget target(getContext());
     target.addLegalDialect<memref::MemRefDialect, arith::ArithDialect,
                            scf::SCFDialect, func::FuncDialect>();
     target.addIllegalDialect<bufferization::BufferizationDialect>();
 
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
-std::unique_ptr<OperationPass<ModuleOp>>
-mlir::createBufferizationToMemRefPass() {
+std::unique_ptr<Pass> mlir::createBufferizationToMemRefPass() {
   return std::make_unique<BufferizationToMemRefPass>();
 }

mlir/test/Conversion/BufferizationToMemRef/bufferization-to-memref-func.mlir

+// RUN: mlir-opt -verify-diagnostics --pass-pipeline="builtin.module(func.func(convert-bufferization-to-memref))" -split-input-file %s | FileCheck %s
+
+// CHECK-NOT: func @deallocHelper
+// CHECK-LABEL: func @conversion_dealloc_simple
+// CHECK-SAME: [[ARG0:%.+]]: memref<2xf32>
+// CHECK-SAME: [[ARG1:%.+]]: i1
+func.func @conversion_dealloc_simple(%arg0: memref<2xf32>, %arg1: i1) {
+  bufferization.dealloc (%arg0 : memref<2xf32>) if (%arg1)
+  return
+}
+
+//      CHECk: scf.if [[ARG1]] {
+// CHECk-NEXT:   memref.dealloc [[ARG0]] : memref<2xf32>
+// CHECk-NEXT: }
+// CHECk-NEXT: return
+
+// -----
+
+func.func @conversion_dealloc_multiple_memrefs_and_retained(%arg0: memref<2xf32>, %arg1: memref<5xf32>, %arg2: memref<1xf32>, %arg3: i1, %arg4: i1, %arg5: memref<2xf32>) -> (i1, i1) {
+  // expected-error @below {{library function required for generic lowering, but cannot be automatically inserted when operating on functions}}
+  // expected-error @below {{failed to legalize operation 'bufferization.dealloc' that was explicitly marked illegal}}
+  %0:2 = bufferization.dealloc (%arg0, %arg1 : memref<2xf32>, memref<5xf32>) if (%arg3, %arg4) retain (%arg2, %arg5 : memref<1xf32>, memref<2xf32>)
+  return %0#0, %0#1 : i1, i1
+}