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

llvmlistbot at llvm.org llvmlistbot at llvm.org
Wed Sep 27 06:06:29 PDT 2023


Author: Matthias Springer
Date: 2023-09-27T15:06:25+02:00
New Revision: 5109cb28fda8076f057c70bd25af7ed82de1d9b7

URL: https://github.com/llvm/llvm-project/commit/5109cb28fda8076f057c70bd25af7ed82de1d9b7
DIFF: https://github.com/llvm/llvm-project/commit/5109cb28fda8076f057c70bd25af7ed82de1d9b7.diff

LOG: [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.)

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

Modified: 
    mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
    mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp
    mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp

Removed: 
    


################################################################################
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>)


        


More information about the Mlir-commits mailing list