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

erick-xanadu · 2025-01-03T16:52:23Z

Hello,

I know that the buffer-deallocation pass is deprecated, but just in case anyone else is still relying on it, we found this issue with buffer-deallocation. The following test case segfaults on main:

func.func public @jit_main() {
    %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 @jitted_fun_0(%alloc_1, %arg1) : (memref<i32>, i1) -> (i1)

      scf.yield %7#0: i1
    }

    return
}

func.func private @jitted_fun_0(%arg1: memref<i32>, %arg2: i1) -> (i1) {
    return %arg2 : i1
}

It looks like it is attempting to introduce a clone for the i1. With the changes submitted, I believe the error is corrected. Happy to improve the test. It would be nice if this is still merged (if deemed a fix) despite buffer-deallocation being deprecated as other people may encounter this.

Thanks, and happy to make changes to improve this fix.

llvmbot · 2025-01-03T16:52:54Z

@llvm/pr-subscribers-mlir

Author: None (erick-xanadu)

Changes

Hello,

I know that the buffer-deallocation pass is deprecated, but just in case anyone else is still relying on it, we found this issue with buffer-deallocation. The following test case segfaults on main:

func.func public @<!-- -->jit_main() {
    %c0 = arith.constant 0 : i1

    %4 = scf.while (%arg1 = %c0) : (i1) -&gt; (i1) {
      scf.condition(%c0) %arg1 : i1
    } do {
    ^bb0(%arg1: i1):
      %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref&lt;i32&gt;

      %7 = func.call @<!-- -->jitted_fun_0(%alloc_1, %arg1) : (memref&lt;i32&gt;, i1) -&gt; (i1)

      scf.yield %7#<!-- -->0: i1
    }

    return
}

func.func private @<!-- -->jitted_fun_0(%arg1: memref&lt;i32&gt;, %arg2: i1) -&gt; (i1) {
    return %arg2 : i1
}

It looks like it is attempting to introduce a clone for the i1. With the changes submitted, I believe the error is corrected. Happy to improve the test. It would be nice if this is still merged (if deemed a fix) despite buffer-deallocation being deprecated as other people may encounter this.

Thanks, and happy to make changes to improve this fix.

Full diff: https://github.com/llvm/llvm-project/pull/121582.diff

2 Files Affected:

(modified) mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp (+3)
(modified) mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir (+21)

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)))
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>

llvmbot · 2025-01-03T16:52:55Z

@llvm/pr-subscribers-mlir-bufferization

Author: None (erick-xanadu)

Changes

Hello,

I know that the buffer-deallocation pass is deprecated, but just in case anyone else is still relying on it, we found this issue with buffer-deallocation. The following test case segfaults on main:

func.func public @<!-- -->jit_main() {
    %c0 = arith.constant 0 : i1

    %4 = scf.while (%arg1 = %c0) : (i1) -&gt; (i1) {
      scf.condition(%c0) %arg1 : i1
    } do {
    ^bb0(%arg1: i1):
      %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref&lt;i32&gt;

      %7 = func.call @<!-- -->jitted_fun_0(%alloc_1, %arg1) : (memref&lt;i32&gt;, i1) -&gt; (i1)

      scf.yield %7#<!-- -->0: i1
    }

    return
}

func.func private @<!-- -->jitted_fun_0(%arg1: memref&lt;i32&gt;, %arg2: i1) -&gt; (i1) {
    return %arg2 : i1
}

It looks like it is attempting to introduce a clone for the i1. With the changes submitted, I believe the error is corrected. Happy to improve the test. It would be nice if this is still merged (if deemed a fix) despite buffer-deallocation being deprecated as other people may encounter this.

Thanks, and happy to make changes to improve this fix.

Full diff: https://github.com/llvm/llvm-project/pull/121582.diff

2 Files Affected:

(modified) mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp (+3)
(modified) mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir (+21)

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)))
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>

github-actions · 2025-01-03T16:55:42Z

✅ With the latest revision this PR passed the C/C++ code formatter.

llvm/llvm-project#121582

matthias-springer

Are you using this pass a lot? I wanted to delete this pass from MLIR soon.

matthias-springer · 2025-01-07T08:22:01Z

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())) {


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

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. :)

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

dime10 · 2025-01-07T18:40:07Z

Are you using this pass a lot? I wanted to delete this pass from MLIR soon.

Thanks for the review @matthias-springer ! We've been meaning to update to the new bufferization infrastructure although we haven't completed the migration yet (currently we are indeed relying on this pass heavily). Is the new deallocation pass useable without one-shot bufferization?

matthias-springer · 2025-01-08T08:34:43Z

Yes, the new deallocation pass (-buffer-deallocation-pipeline) can be used without One-Shot Bufferize. See documentation here.

…ation (#1408) **Context:** This patch fixes a bug in upstream MLIR that prevents buffer deallocation from working in certain cases. The patch was also submitted upstream: llvm/llvm-project#121582 **Description of the Change:** This PR adds a patch to the llvm-project source, and updates the relevant build recipes. The llvm step during wheel builds is now also handled through Make instead of replicating the CMake configuration in each wheel script. **Benefits:** Unblocks our friends over at Qrisp. **Possible Drawbacks:** Patching dependencies is not great, but hopefully temporary. **Related GitHub Issues:** [sc-81444]

erick-xanadu · 2025-01-13T14:29:27Z

Are you using this pass a lot? I wanted to delete this pass from MLIR soon.

Yes, but we are planning on updating soon. Thanks!

erick-xanadu added 2 commits January 3, 2025 11:37

Introduce copies only for MemRef typed values.

5867e72

Add test

3da161b

erick-xanadu requested a review from matthias-springer as a code owner January 3, 2025 16:52

llvmbot added mlir mlir:bufferization Bufferization infrastructure labels Jan 3, 2025

style

df60731

dime10 added a commit to PennyLaneAI/catalyst that referenced this pull request Jan 3, 2025

Patch mlir buffer-deallocation bug

1677a25

llvm/llvm-project#121582

dime10 mentioned this pull request Jan 3, 2025

Allow quantum registers returned from functions during buffer deallocation PennyLaneAI/catalyst#1408

Merged

matthias-springer approved these changes Jan 7, 2025

View reviewed changes

Exclude adding non-memref values to valuesToFree

4b5db41

erick-xanadu added 2 commits January 13, 2025 09:35

Style

512b6f7

Fix test by moving allocation to dominate all blocks followed by while.

cc4ca87

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

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

Uh oh!

erick-xanadu commented Jan 3, 2025

Uh oh!

llvmbot commented Jan 3, 2025

Uh oh!

llvmbot commented Jan 3, 2025

Uh oh!

github-actions bot commented Jan 3, 2025 •

edited

Loading

Uh oh!

matthias-springer left a comment

Uh oh!

matthias-springer Jan 7, 2025

Uh oh!

erick-xanadu Jan 13, 2025

Uh oh!

erick-xanadu Jan 14, 2025

Uh oh!

matthias-springer Jan 15, 2025

Uh oh!

dime10 commented Jan 7, 2025 •

edited

Loading

Uh oh!

matthias-springer commented Jan 8, 2025

Uh oh!

erick-xanadu commented Jan 13, 2025

Uh oh!

Uh oh!

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

Are you sure you want to change the base?

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

Uh oh!

Conversation

erick-xanadu commented Jan 3, 2025

Uh oh!

llvmbot commented Jan 3, 2025

Uh oh!

llvmbot commented Jan 3, 2025

Uh oh!

github-actions bot commented Jan 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

matthias-springer left a comment

Choose a reason for hiding this comment

Uh oh!

matthias-springer Jan 7, 2025

Choose a reason for hiding this comment

Uh oh!

erick-xanadu Jan 13, 2025

Choose a reason for hiding this comment

Uh oh!

erick-xanadu Jan 14, 2025

Choose a reason for hiding this comment

Uh oh!

matthias-springer Jan 15, 2025

Choose a reason for hiding this comment

Uh oh!

dime10 commented Jan 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

matthias-springer commented Jan 8, 2025

Uh oh!

erick-xanadu commented Jan 13, 2025

Uh oh!

Uh oh!

github-actions bot commented Jan 3, 2025 •

edited

Loading

dime10 commented Jan 7, 2025 •

edited

Loading