From 5867e7207a19e7539f9f55fc3375c510a6be9c45 Mon Sep 17 00:00:00 2001 From: Erick Ochoa Lopez 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 a0a81d4add712..36edbb3395edd 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(value.getType())) { + continue; + } if (failed(isa(value) ? introduceBlockArgCopy(cast(value)) : introduceValueCopyForRegionResult(value))) From 3da161bc6fcec2aed264fbc0457f4f372c8b9a7f Mon Sep 17 00:00:00 2001 From: Erick Ochoa Lopez 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 3fbe3913c6549..a6e6f724a822a 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 + %7 = func.call @foo(%alloc_1, %arg1) : (memref, i1) -> (i1) + scf.yield %7#0: i1 + } + return +} + +func.func private @foo(%arg1: memref, %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 From df60731de2aadee0d647fd7a25c861ca907212e0 Mon Sep 17 00:00:00 2001 From: Erick Ochoa Lopez 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 36edbb3395edd..c07ca72c60f4d 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(value.getType())) { - continue; + continue; } if (failed(isa(value) ? introduceBlockArgCopy(cast(value)) From 4b5db41a84bd325c31c632bd24e29d3d197998a7 Mon Sep 17 00:00:00 2001 From: Erick Ochoa Lopez 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 c07ca72c60f4d..51efb6065ccfc 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(value))) { toProcess.emplace_back(value, parentBlock); - valuesToFree.insert(value); + if (isa(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(value.getType())) { - continue; - } if (failed(isa(value) ? introduceBlockArgCopy(cast(value)) : introduceValueCopyForRegionResult(value))) From 512b6f759d914c7029be568993f4301654effb93 Mon Sep 17 00:00:00 2001 From: Erick Ochoa Lopez 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 51efb6065ccfc..eb8055ea9aa79 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(value))) { toProcess.emplace_back(value, parentBlock); if (isa(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 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 a6e6f724a822a..7bcf32ef1441f 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 %4 = scf.while (%arg1 = %c0) : (i1) -> (i1) { scf.condition(%c0) %arg1 : i1 } do { ^bb0(%arg1: i1): - %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref %7 = func.call @foo(%alloc_1, %arg1) : (memref, i1) -> (i1) scf.yield %7#0: i1 }