[Mlir-commits] [mlir] [MLIR][buffer-deallocation] Introduce copies only for MemRef typed values. (PR #121582)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jan 13 08:32:31 PST 2025
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/6] 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/6] 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/6] 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/6] 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)))
>From 512b6f759d914c7029be568993f4301654effb93 Mon Sep 17 00:00:00 2001
From: Erick Ochoa Lopez <erick.ochoalopez at xanadu.ai>
Date: Mon, 13 Jan 2025 09:35:14 -0500
Subject: [PATCH 5/6] 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 51efb6065ccfc8..eb8055ea9aa79c 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
@@ -284,7 +284,7 @@ class BufferDeallocation : public BufferPlacementTransformationBase {
(definingBlock == parentBlock && isa<BlockArgument>(value))) {
toProcess.emplace_back(value, parentBlock);
if (isa<BaseMemRefType>(value.getType())) {
- valuesToFree.insert(value);
+ valuesToFree.insert(value);
}
} else if (visitedValues.insert(std::make_tuple(value, definingBlock))
.second)
>From cc4ca873de926813152539c4bf981251597d0056 Mon Sep 17 00:00:00 2001
From: Erick Ochoa Lopez <erick.ochoalopez at xanadu.ai>
Date: Mon, 13 Jan 2025 11:31:57 -0500
Subject: [PATCH 6/6] Fix test by moving allocation to dominate all blocks
followed by while.
---
.../Dialect/Bufferization/Transforms/buffer-deallocation.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
index a6e6f724a822ac..7bcf32ef1441fa 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
@@ -1274,11 +1274,11 @@ func.func @while_two_arg(%arg0: index) {
// CHECK-LABEL: func @while_fun
func.func @while_fun() {
%c0 = arith.constant 0 : i1
+ %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref<i32>
%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
}
More information about the Mlir-commits
mailing list