[Mlir-commits] [mlir] [mlir][bufferization] Add `buffer_deallocation` transform op (PR #74734)

[Jiefeng Wang]

https://github.com/jiefwo created https://github.com/llvm/llvm-project/pull/74734

None

>From 33e8c62065ba4c91b09e8dbc04625fcbf3b4bc93 Mon Sep 17 00:00:00 2001
From: Jiefeng Wang <wangjiefeng at huawei.com>
Date: Fri, 8 Dec 2023 00:22:33 +0800
Subject: [PATCH] [mlir][bufferization] Add `buffer_deallocation` transform op

---
 .../TransformOps/BufferizationTransformOps.td | 23 +++++++++++++++++++
 .../BufferizationTransformOps.cpp             | 20 ++++++++++++++++
 .../Transforms/transform-ops.mlir             | 19 +++++++++++++++
 3 files changed, 62 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td b/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
index 9b588eb610e51b..e9136f78a41bf2 100644
--- a/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
@@ -19,6 +19,29 @@ include "mlir/IR/OpBase.td"
 def Transform_EmptyOp : Transform_ConcreteOpType<"tensor.empty">;
 def Transform_AllocTensorOp : Transform_ConcreteOpType<"bufferization.alloc_tensor">;
 
+//===----------------------------------------------------------------------===//
+// BufferDeallocationOp
+//===----------------------------------------------------------------------===//
+
+def BufferDeallocationOp
+    : Op<Transform_Dialect, "bufferization.buffer_deallocation",
+        [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+         DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let description = [{
+    Introduce all required deallocation operations for all buffers within the
+    targeted op.
+
+    #### Return modes
+
+    This operation consumes the `target` handle and produces the `transformed`
+    handle.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target);
+  let results = (outs TransformHandleTypeInterface:$transformed);
+  let assemblyFormat = "$target attr-dict `:` functional-type(operands, results)";
+}
+
 //===----------------------------------------------------------------------===//
 // BufferLoopHoistingOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp b/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
index cbb36639a3383d..5ed58d3b550611 100644
--- a/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
@@ -23,6 +23,26 @@ using namespace mlir;
 using namespace mlir::bufferization;
 using namespace mlir::transform;
 
+//===----------------------------------------------------------------------===//
+// BufferDeallocationOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::BufferDeallocationOp::apply(transform::TransformRewriter &rewriter,
+                                       TransformResults &transformResults,
+                                       TransformState &state) {
+  auto payloadOps = state.getPayloadOps(getTarget());
+  for (Operation *target : payloadOps) {
+    if (!isa<FunctionOpInterface>(target))
+      return emitSilenceableError() << "expected function target";
+    if (failed(bufferization::deallocateBuffers(target)))
+      return emitSilenceableError() << "buffer deallocation failed";
+  }
+
+  transformResults.set(cast<OpResult>(getTransformed()), payloadOps);
+  return DiagnosedSilenceableFailure::success();
+}
+
 //===----------------------------------------------------------------------===//
 // BufferLoopHoistingOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
index 3c50a9e72d9d9b..ccc915b54abcd5 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
@@ -238,3 +238,22 @@ func.func @empty_to_tensor_alloc() -> tensor<2x2xf32> {
   %0 = bufferization.alloc_tensor() : tensor<2x2xf32>
   return %0 : tensor<2x2xf32>
 }
+
+// -----
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    transform.bufferization.buffer_deallocation %0 : (!transform.any_op) -> !transform.any_op
+    transform.yield
+  }
+}
+
+// CHECK-LABEL: func @buffer_deallocation(
+//       CHECK:   memref.alloc
+//       CHECK:   memref.dealloc
+//       CHECK:   return
+func.func @buffer_deallocation(%arg0: memref<4xf32>) {
+  %0 = memref.alloc() : memref<4xf32>
+  return
+}