[Mlir-commits] [mlir] [MLIR][buffer-deallocation] Introduce copies only for MemRef typed values. (PR #121582)

[llvmlistbot at llvm.org]

https://github.com/erick-xanadu updated https://github.com/llvm/llvm-project/pull/121582

>From 5867e7207a19e7539f9f55fc3375c510a6be9c45 Mon Sep 17 00:00:00 2001
From: Erick Ochoa Lopez <erick.ochoalopez at xanadu.ai>
Date: Fri, 3 Jan 2025 11:37:07 -0500
Subject: [PATCH 1/4] Introduce copies only for MemRef typed values.

---
 .../Dialect/Bufferization/Transforms/BufferDeallocation.cpp    | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
index a0a81d4add7121..36edbb3395eddf 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
@@ -308,6 +308,9 @@ class BufferDeallocation : public BufferPlacementTransformationBase {
 
     // Add new allocs and additional clone operations.
     for (Value value : valuesToFree) {
+      if (!isa<BaseMemRefType>(value.getType())) {
+          continue;
+      }
       if (failed(isa<BlockArgument>(value)
                      ? introduceBlockArgCopy(cast<BlockArgument>(value))
                      : introduceValueCopyForRegionResult(value)))

>From 3da161bc6fcec2aed264fbc0457f4f372c8b9a7f Mon Sep 17 00:00:00 2001
From: Erick Ochoa Lopez <erick.ochoalopez at xanadu.ai>
Date: Fri, 3 Jan 2025 11:51:50 -0500
Subject: [PATCH 2/4] Add test

---
 .../Transforms/buffer-deallocation.mlir       | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
index 3fbe3913c6549e..a6e6f724a822ac 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
@@ -1271,6 +1271,27 @@ func.func @while_two_arg(%arg0: index) {
 
 // -----
 
+// CHECK-LABEL: func @while_fun
+func.func @while_fun() {
+    %c0 = arith.constant 0 : i1
+    %4 = scf.while (%arg1 = %c0) : (i1) -> (i1) {
+      scf.condition(%c0) %arg1 : i1
+    } do {
+    ^bb0(%arg1: i1):
+      %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref<i32>
+      %7 = func.call @foo(%alloc_1, %arg1) : (memref<i32>, i1) -> (i1)
+      scf.yield %7#0: i1
+    }
+    return
+}
+
+func.func private @foo(%arg1: memref<i32>, %arg2: i1) -> (i1) {
+    return %arg2 : i1
+}
+
+// -----
+
+// CHECK-LABEL: func @while_three_arg
 func.func @while_three_arg(%arg0: index) {
 // CHECK: %[[ALLOC:.*]] = memref.alloc
   %a = memref.alloc(%arg0) : memref<?xf32>

>From df60731de2aadee0d647fd7a25c861ca907212e0 Mon Sep 17 00:00:00 2001
From: Erick Ochoa Lopez <erick.ochoalopez at xanadu.ai>
Date: Fri, 3 Jan 2025 11:56:52 -0500
Subject: [PATCH 3/4] style

---
 .../lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
index 36edbb3395eddf..c07ca72c60f4dc 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
@@ -309,7 +309,7 @@ class BufferDeallocation : public BufferPlacementTransformationBase {
     // Add new allocs and additional clone operations.
     for (Value value : valuesToFree) {
       if (!isa<BaseMemRefType>(value.getType())) {
-          continue;
+        continue;
       }
       if (failed(isa<BlockArgument>(value)
                      ? introduceBlockArgCopy(cast<BlockArgument>(value))

>From 4b5db41a84bd325c31c632bd24e29d3d197998a7 Mon Sep 17 00:00:00 2001
From: Erick Ochoa Lopez <erick.ochoalopez at xanadu.ai>
Date: Mon, 13 Jan 2025 09:26:57 -0500
Subject: [PATCH 4/4] Exclude adding non-memref values to valuesToFree

---
 .../Bufferization/Transforms/BufferDeallocation.cpp        | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
index c07ca72c60f4dc..51efb6065ccfc8 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
@@ -283,7 +283,9 @@ class BufferDeallocation : public BufferPlacementTransformationBase {
         if (!dominators.dominates(definingBlock, parentBlock) ||
             (definingBlock == parentBlock && isa<BlockArgument>(value))) {
           toProcess.emplace_back(value, parentBlock);
-          valuesToFree.insert(value);
+          if (isa<BaseMemRefType>(value.getType())) {
+              valuesToFree.insert(value);
+          }
         } else if (visitedValues.insert(std::make_tuple(value, definingBlock))
                        .second)
           toProcess.emplace_back(value, definingBlock);
@@ -308,9 +310,6 @@ class BufferDeallocation : public BufferPlacementTransformationBase {
 
     // Add new allocs and additional clone operations.
     for (Value value : valuesToFree) {
-      if (!isa<BaseMemRefType>(value.getType())) {
-        continue;
-      }
       if (failed(isa<BlockArgument>(value)
                      ? introduceBlockArgCopy(cast<BlockArgument>(value))
                      : introduceValueCopyForRegionResult(value)))