[Mlir-commits] [mlir] [mlir][bufferization] Make buffer deallocation pipeline op type independent (PR #67546)

[Matthias Springer]

https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/67546

>From 11ac8efde0d7141f603576e42cea8be2651eb2e9 Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Wed, 27 Sep 2023 13:54:50 +0200
Subject: [PATCH] [mlir][bufferization] Make buffer deallocation pipeline op
 type independent

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.)
---
 .../Bufferization/Transforms/Passes.td        |  4 ++--
 .../Pipelines/BufferizationPipelines.cpp      | 15 ++++++-------
 .../OwnershipBasedBufferDeallocation.cpp      | 15 ++++++++-----
 .../dealloc-helper-function-invalid.mlir      | 22 +++++++++++++++++++
 .../dealloc-helper-function.mlir              | 20 +++++++++++++++++
 5 files changed, 61 insertions(+), 15 deletions(-)
 create mode 100644 mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-helper-function-invalid.mlir
 create mode 100644 mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-helper-function.mlir

diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 62383e376f6f7a3..e01f36b8daa18d6 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/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 = [{
diff --git a/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp b/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp
index b2a60feb9a7f011..a2878f0b80fa1cd 100644
--- a/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp
+++ b/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());
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index f698ef5f9ba08c4..fd36716163d0ad4 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/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();
   }
 };
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-helper-function-invalid.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-helper-function-invalid.mlir
new file mode 100644
index 000000000000000..d22fc7d9b74e9e2
--- /dev/null
+++ b/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>)
+
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-helper-function.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-helper-function.mlir
new file mode 100644
index 000000000000000..eaf6625c838c0fc
--- /dev/null
+++ b/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>)