Skip to content

[MLIR][buffer-deallocation] Introduce copies only for MemRef typed values. #121582

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

[MLIR][buffer-deallocation] Introduce copies only for MemRef typed values. #121582

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
5867e72
Introduce copies only for MemRef typed values.
erick-xanadu Jan 3, 2025
3da161b
Add test
erick-xanadu Jan 3, 2025
df60731
style
erick-xanadu Jan 3, 2025
4b5db41
Exclude adding non-memref values to valuesToFree
erick-xanadu Jan 13, 2025
512b6f7
Style
erick-xanadu Jan 13, 2025
cc4ca87
Fix test by moving allocation to dominate all blocks followed by while.
erick-xanadu Jan 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -308,6 +308,9 @@ class BufferDeallocation : public BufferPlacementTransformationBase {

// Add new allocs and additional clone operations.
for (Value value : valuesToFree) {
if (!isa<BaseMemRefType>(value.getType())) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you put the check earlier, so that non-memref values are not even added to valuesToFree?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

4b5db41

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @matthias-springer, I implemented the change but I noticed that my test case failed. If I do this, it appears that a deallocation will be placed in the block following the scf.while op. This means that the original test fails because the memref.alloc inside the scf.while body does not dominate the memref.dealloc operation that will be inserted. I think that no memref.dealloc should be inserted though (but I am not sure) because it is passed to a function which means that the allocated pointer may be stored on the heap or elsewhere. I'll try to take a closer look, but if you have a quick hunch about what can be wrong, please let me know. :)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not too familiar with this pass unfortunately...

continue;
}
if (failed(isa<BlockArgument>(value)
? introduceBlockArgCopy(cast<BlockArgument>(value))
: introduceValueCopyForRegionResult(value)))
Expand Down
21 changes: 21 additions & 0 deletions mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down
Loading