[Mlir-commits] [mlir] [mlir] run buffer deallocation in transform tutorial (PR #67978)

[Oleksandr Alex Zinenko]

https://github.com/ftynse created https://github.com/llvm/llvm-project/pull/67978

Buffer deallocation pipeline previously was incorrect when applied to functions. It has since been fixed. Make sure it is exercised in the tutorial to avoid leaking allocations.

>From b13ae2565cfd5690ad036f15aaa1026d5551b085 Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko at google.com>
Date: Mon, 2 Oct 2023 12:19:50 +0000
Subject: [PATCH] [mlir] run buffer deallocation in transform tutorial

Buffer deallocation pipeline previously was incorrect when applied to
functions. It has since been fixed. Make sure it is exercised in the
tutorial to avoid leaking allocations.
---
 mlir/docs/Tutorials/transform/ChH.md       | 14 ++++++++++++++
 mlir/test/Examples/transform/ChH/full.mlir | 13 +++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/mlir/docs/Tutorials/transform/ChH.md b/mlir/docs/Tutorials/transform/ChH.md
index 0696853dda1cf75..7c12728ac32466d 100644
--- a/mlir/docs/Tutorials/transform/ChH.md
+++ b/mlir/docs/Tutorials/transform/ChH.md
@@ -497,6 +497,20 @@ bufferization is directly available as a transform operation.
   function_boundary_type_conversion = 1 : i32 }
 ```
 
+One-shot bufferization itself does not produce buffer deallocations, which may
+lead to leaks. So we have to run the buffer deallocation pass pipeline to avoid
+them. Note that the transform dialect seamlessly runs named passes and pass
+pipelines: if desired, one could replace complex `--pass-pipeline expressions`
+with operations. Note that we apply the pipeline to functions rather than entire
+module to avoid running it on the transform IR that is contained in the module.
+
+```mlir
+%f = transform.structured.match ops{["func.func"]} in %arg1
+  : (!transform.any_op) -> !transform.any_op
+transform.apply_registered_pass "buffer-deallocation-pipeline" to %f
+  : (!transform.any_op) -> !transform.any_op
+```
+
 In this particular case, the transformed IR could be directly bufferized. This
 is not always the case in general as some operations, in particular
 `tensor.empty` may not be bufferizable. Such operations need to be removed
diff --git a/mlir/test/Examples/transform/ChH/full.mlir b/mlir/test/Examples/transform/ChH/full.mlir
index ed0509fcef6f12b..d90d740b4453126 100644
--- a/mlir/test/Examples/transform/ChH/full.mlir
+++ b/mlir/test/Examples/transform/ChH/full.mlir
@@ -1,6 +1,7 @@
 // RUN: mlir-opt %s --test-transform-dialect-interpreter \
 // RUN:             --test-transform-dialect-erase-schedule \
 // RUN:             --math-uplift-to-fma \
+// RUN:             --convert-bufferization-to-memref \
 // RUN:             --test-lower-to-llvm |\
 // RUN: FileCheck %s
 
@@ -307,10 +308,22 @@ module attributes { transform.with_named_sequence } {
     // transformation process, so invalidation is not an issue. However, if
     // other transformations, such as loop unrolling, are required after
     // bufferization, new handles should be produced using the match operations.
+    //
+    // One-shot bufferization itself does not produce buffer deallocations,
+    // which may lead to leaks. So we have to run the buffer deallocation pass
+    // pipeline to avoid them. Note that the transform dialect seamlessly runs
+    // named passes and pass pipelines: if desired, one could replace complex
+    // --pass-pipeline expressions with operations. Note that we apply the
+    // pipeline to functions rather than entire module to avoid running it
+    // on the transform IR that is contained in the module.
     %arg1 = transform.bufferization.one_shot_bufferize %arg0 {
       bufferize_function_boundaries = true,
       function_boundary_type_conversion = 1 : i32 }
       : (!transform.any_op) -> !transform.any_op
+    %f = transform.structured.match ops{["func.func"]} in %arg1
+      : (!transform.any_op) -> !transform.any_op
+    transform.apply_registered_pass "buffer-deallocation-pipeline" to %f
+      : (!transform.any_op) -> !transform.any_op
 
     // Apply general canonicalization and CSE to each function after
     // bufferization as new simplification opportunities may have appeared.