[mlir][bufferization] Make buffer deallocation pipeline op type independent (#67546)

The buffer deallocation pipeline now works on modules and functions.
Also add extra test cases that run the buffer deallocation pipeline on
modules and functions. (Test cases that insert a helper function.)

Author: matthias-springer

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

@@ -89,7 +89,7 @@ def BufferDeallocation : Pass<"buffer-deallocation", "func::FuncOp"> {
 }
 
 def OwnershipBasedBufferDeallocation : Pass<
-    "ownership-based-buffer-deallocation", "func::FuncOp"> {
+    "ownership-based-buffer-deallocation"> {
   let summary = "Adds all required dealloc operations for all allocations in "
                 "the input program";
   let description = [{
@@ -233,7 +233,7 @@ def OwnershipBasedBufferDeallocation : Pass<
 }
 
 def BufferDeallocationSimplification :
-    Pass<"buffer-deallocation-simplification", "func::FuncOp"> {
+    Pass<"buffer-deallocation-simplification"> {
   let summary = "Optimizes `bufferization.dealloc` operation for more "
                 "efficient codegen";
   let description = [{

mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp

@@ -20,16 +20,15 @@
 
 void mlir::bufferization::buildBufferDeallocationPipeline(
     OpPassManager &pm, const BufferDeallocationPipelineOptions &options) {
-  pm.addNestedPass<func::FuncOp>(
-      memref::createExpandReallocPass(/*emitDeallocs=*/false));
-  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
-  pm.addNestedPass<func::FuncOp>(createOwnershipBasedBufferDeallocationPass(
+  pm.addPass(memref::createExpandReallocPass(/*emitDeallocs=*/false));
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(createOwnershipBasedBufferDeallocationPass(
       options.privateFunctionDynamicOwnership.getValue()));
-  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
-  pm.addNestedPass<func::FuncOp>(createBufferDeallocationSimplificationPass());
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(createBufferDeallocationSimplificationPass());
   pm.addPass(createLowerDeallocationsPass());
-  pm.addNestedPass<func::FuncOp>(createCSEPass());
-  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+  pm.addPass(createCSEPass());
+  pm.addPass(createCanonicalizerPass());
 }
 
 //===----------------------------------------------------------------------===//

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

@@ -983,12 +983,17 @@ struct OwnershipBasedBufferDeallocationPass
     this->privateFuncDynamicOwnership.setValue(privateFuncDynamicOwnership);
   }
   void runOnOperation() override {
-    func::FuncOp func = getOperation();
-    if (func.isExternal())
-      return;
+    auto status = getOperation()->walk([&](func::FuncOp func) {
+      if (func.isExternal())
+        return WalkResult::skip();
+
+      if (failed(deallocateBuffersOwnershipBased(func,
+                                                 privateFuncDynamicOwnership)))
+        return WalkResult::interrupt();
 
-    if (failed(
-            deallocateBuffersOwnershipBased(func, privateFuncDynamicOwnership)))
+      return WalkResult::advance();
+    });
+    if (status.wasInterrupted())
       signalPassFailure();
   }
 };

mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-helper-function-invalid.mlir

@@ -0,0 +1,22 @@
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(buffer-deallocation-pipeline))" -verify-diagnostics
+
+// Test case requires a helper function but is run on a function.
+
+// CHECK-LABEL: func.func @test(
+//       CHECK:   func.call @[[helper:.*]]({{.*}}) : ({{.*}}) -> ()
+func.func @test(%lb : index, %ub: index) -> (memref<5xf32>, memref<5xf32>) {
+  %0 = memref.alloc() : memref<5xf32>
+  %1 = memref.alloc() : memref<5xf32>
+  %c1 = arith.constant 1 : index
+  %a, %b = scf.for %iv = %lb to %ub step %c1 iter_args(%c = %0, %d = %1)
+      -> (memref<5xf32>, memref<5xf32>) {
+    // 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}}
+    scf.yield %d, %c : memref<5xf32>, memref<5xf32>
+  }
+  return %a, %b : memref<5xf32>, memref<5xf32>
+}
+
+//      CHECK: func.func @[[helper]](
+// CHECK-SAME:     %{{.*}}: memref<?xindex>, %{{.*}}: memref<?xindex>, %{{.*}}: memref<?xi1>, %{{.*}}: memref<?xi1>, %{{.*}}: memref<?xi1>)
+

mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-helper-function.mlir

@@ -0,0 +1,20 @@
+// RUN: mlir-opt %s -buffer-deallocation-pipeline | FileCheck %s
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(buffer-deallocation-pipeline)" | FileCheck %s
+
+// Test case requires a helper function.
+
+// CHECK-LABEL: func.func @test(
+//       CHECK:   func.call @[[helper:.*]]({{.*}}) : ({{.*}}) -> ()
+func.func @test(%lb : index, %ub: index) -> (memref<5xf32>, memref<5xf32>) {
+  %0 = memref.alloc() : memref<5xf32>
+  %1 = memref.alloc() : memref<5xf32>
+  %c1 = arith.constant 1 : index
+  %a, %b = scf.for %iv = %lb to %ub step %c1 iter_args(%c = %0, %d = %1)
+      -> (memref<5xf32>, memref<5xf32>) {
+    scf.yield %d, %c : memref<5xf32>, memref<5xf32>
+  }
+  return %a, %b : memref<5xf32>, memref<5xf32>
+}
+
+//      CHECK: func.func private @[[helper]](
+// CHECK-SAME:     %{{.*}}: memref<?xindex>, %{{.*}}: memref<?xindex>, %{{.*}}: memref<?xi1>, %{{.*}}: memref<?xi1>, %{{.*}}: memref<?xi1>)